feature(translation-worker): cache hunyuan models

workers/translation-worker/pyproject.toml

@@ -12,10 +12,14 @@ requires-python = ">=3.11.0, <3.13"
 
 dependencies = [
   "datashare-python~=0.12.6",
+  "huggingface-hub>=1.0.0,<2",
   "language-data>=1.4.0",
   "pydantic-extra-types[pycountry]==2.11.1",
 ]
 
+[project.scripts]
+translation-worker = "translation_worker.cli:cli_app"
+
 [project.entry-points."datashare.workflows"]
 workflows = "translation_worker.workflows:WORKFLOWS"
 

workers/translation-worker/tests/test_model_cache.py

@@ -0,0 +1,49 @@
+import sys
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import pytest
+from translation_worker import model_cache
+from translation_worker.cli import cli_app
+from translation_worker.model_cache import cache_models
+from translation_worker.objects import DEFAULT_HUNYUAN_MODEL_REF, TranslationModel
+from typer.testing import CliRunner
+
+
+@pytest.mark.parametrize(
+    ("model_ref", "expected_repo"),
+    [(None, DEFAULT_HUNYUAN_MODEL_REF), ("tencent/HY-MT1.5-7B", "tencent/HY-MT1.5-7B")],
+)
+def test_cache_models_hunyuan(
+    tmp_path: Path, model_ref: str | None, expected_repo: str
+) -> None:
+    # Given
+    hf_hub = MagicMock()
+    hf_hub.get_token.return_value = None
+    # When
+    with patch.dict(sys.modules, {"huggingface_hub": hf_hub}):
+        cache_models(TranslationModel.HUNYUAN, tmp_path, model_ref=model_ref)
+    # Then
+    hf_hub.snapshot_download.assert_called_once_with(
+        repo_id=expected_repo, cache_dir=tmp_path, token=None
+    )
+
+
+def test_cache_models_should_raise_for_uncachable_model() -> None:
+    with pytest.raises(ValueError, match="can't be cached"):
+        cache_models(TranslationModel.ARGOS)
+
+
+def test_cli_cache_models(tmp_path: Path) -> None:
+    # Given
+    runner = CliRunner()
+    cache = MagicMock()
+    cmd = ["model", "cache", "hunyuan", str(tmp_path), "--model-ref", "some/ref"]
+    # When
+    with patch.object(model_cache, "cache_models", cache):
+        result = runner.invoke(cli_app, cmd, catch_exceptions=False)
+    # Then
+    assert result.exit_code == 0, result.output
+    cache.assert_called_once_with(
+        TranslationModel.HUNYUAN, tmp_path, model_ref="some/ref"
+    )

workers/translation-worker/translation_worker/cli.py

@@ -0,0 +1,48 @@
+import logging
+from pathlib import Path
+from typing import Annotated
+
+import typer
+from icij_common.logging_utils import setup_loggers
+
+import translation_worker
+
+from . import model_cache
+from .objects import TranslationModel
+
+_CACHE_MODELS_HELP = "cache translation models in the provided location"
+_CACHE_MODELS_MODEL_HELP = "translation model to cache"
+_CACHE_MODELS_CACHE_DIR_HELP = (
+    "cache directory, if not provided defaults to HF cache directory"
+)
+_CACHE_MODELS_MODEL_REF_HELP = "model HF repo, defaults to the model default"
+
+logger = logging.getLogger(__name__)
+
+cli_app = typer.Typer(
+    context_settings={"help_option_names": ["-h", "--help"]},
+    pretty_exceptions_enable=False,
+)
+models_app = typer.Typer(name="model")
+cli_app.add_typer(models_app)
+
+
+@cli_app.callback()
+def main() -> None:
+    setup_loggers(["__main__", translation_worker.__name__])
+
+
+@models_app.command(help=_CACHE_MODELS_HELP)
+def cache(
+    model: Annotated[
+        TranslationModel,
+        typer.Argument(help=_CACHE_MODELS_MODEL_HELP, case_sensitive=False),
+    ],
+    cache_dir: Annotated[
+        Path | None, typer.Argument(help=_CACHE_MODELS_CACHE_DIR_HELP)
+    ] = None,
+    model_ref: Annotated[
+        str | None, typer.Option(help=_CACHE_MODELS_MODEL_REF_HELP)
+    ] = None,
+) -> None:
+    model_cache.cache_models(model, cache_dir, model_ref=model_ref)

workers/translation-worker/translation_worker/model_cache.py

@@ -0,0 +1,50 @@
+import logging
+from collections.abc import Callable
+from pathlib import Path
+
+from .objects import DEFAULT_HUNYUAN_MODEL_REF, TranslationModel
+
+logger = logging.getLogger(__name__)
+
+
+def cache_hunyuan_model(
+    model_ref: str | None = None, cache_dir: Path | None = None
+) -> None:
+    """Download a Hunyuan model to the HF hub cache.
+
+    :param model_ref: HF repo of the model, defaults to
+        ``DEFAULT_HUNYUAN_MODEL_REF``
+    :param cache_dir: HF hub cache directory, defaults to ``HF_HUB_CACHE``
+    """
+    from huggingface_hub import get_token, snapshot_download  # noqa: PLC0415
+
+    if model_ref is None:
+        model_ref = DEFAULT_HUNYUAN_MODEL_REF
+    logger.info("caching %s to %s", model_ref, cache_dir)
+    snapshot_download(repo_id=model_ref, cache_dir=cache_dir, token=get_token())
+
+
+_MODEL_CACHE_FNS: dict[TranslationModel, Callable[[str | None, Path | None], None]] = {
+    TranslationModel.HUNYUAN: cache_hunyuan_model,
+}
+
+
+def cache_models(
+    model: TranslationModel,
+    cache_dir: Path | None = None,
+    *,
+    model_ref: str | None = None,
+) -> None:
+    """Download a translation model to the HF hub cache so workers can share it.
+
+    :param model: the translation model
+    :param cache_dir: HF hub cache directory, defaults to ``HF_HUB_CACHE``
+    :param model_ref: the model's HF repo, defaults to the model's default
+    :raises ValueError: when the model can't be cached
+    """
+    if model not in _MODEL_CACHE_FNS:
+        msg = (
+            f"{model} models can't be cached, cachable models: {list(_MODEL_CACHE_FNS)}"
+        )
+        raise ValueError(msg)
+    _MODEL_CACHE_FNS[model](model_ref, cache_dir)

workers/translation-worker/uv.dist.lock

@@ -503,6 +503,7 @@ name = "datashare-translation-worker"
 source = { editable = "." }
 dependencies = [
     { name = "datashare-python" },
+    { name = "huggingface-hub" },
     { name = "language-data" },
     { name = "pydantic-extra-types", extra = ["pycountry"] },
 ]
@@ -545,6 +546,7 @@ requires-dist = [
     { name = "argostranslate", marker = "extra == 'c2translate'", specifier = "==1.11.0" },
     { name = "argostranslate", marker = "extra == 'inference'", specifier = "==1.11.0" },
     { name = "datashare-python", specifier = "~=0.12.6" },
+    { name = "huggingface-hub", specifier = ">=1.0.0,<2" },
     { name = "language-data", specifier = ">=1.4.0" },
     { name = "pydantic-extra-types", extras = ["pycountry"], specifier = "==2.11.1" },
     { name = "torch", marker = "extra == 'torch'", specifier = "==2.11.0" },