From 0ab5b161db4a20f0945eccf2470c48e2e7c6813d Mon Sep 17 00:00:00 2001 From: Lion Summerbell Date: Thu, 15 Oct 2026 09:51:25 +0200 Subject: [PATCH] feature(translation-worker): cache hunyuan models --- workers/translation-worker/pyproject.toml | 4 ++ .../tests/test_model_cache.py | 49 ++++++++++++++++++ .../translation_worker/cli.py | 48 ++++++++++++++++++ .../translation_worker/model_cache.py | 50 +++++++++++++++++++ workers/translation-worker/uv.dist.lock | 2 + workers/translation-worker/uv.lock | 2 + 6 files changed, 155 insertions(+) create mode 100644 workers/translation-worker/tests/test_model_cache.py create mode 100644 workers/translation-worker/translation_worker/cli.py create mode 100644 workers/translation-worker/translation_worker/model_cache.py diff --git a/workers/translation-worker/pyproject.toml b/workers/translation-worker/pyproject.toml index d482957e..eb551d05 100644 --- a/workers/translation-worker/pyproject.toml +++ b/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" diff --git a/workers/translation-worker/tests/test_model_cache.py b/workers/translation-worker/tests/test_model_cache.py new file mode 100644 index 00000000..5ea2401e --- /dev/null +++ b/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" + ) diff --git a/workers/translation-worker/translation_worker/cli.py b/workers/translation-worker/translation_worker/cli.py new file mode 100644 index 00000000..76287506 --- /dev/null +++ b/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) diff --git a/workers/translation-worker/translation_worker/model_cache.py b/workers/translation-worker/translation_worker/model_cache.py new file mode 100644 index 00000000..568a4289 --- /dev/null +++ b/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) diff --git a/workers/translation-worker/uv.dist.lock b/workers/translation-worker/uv.dist.lock index e099d770..a9720a1d 100644 --- a/workers/translation-worker/uv.dist.lock +++ b/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" }, diff --git a/workers/translation-worker/uv.lock b/workers/translation-worker/uv.lock index 62a59e13..affbe0f1 100644 --- a/workers/translation-worker/uv.lock +++ b/workers/translation-worker/uv.lock @@ -531,6 +531,7 @@ name = "datashare-translation-worker" source = { editable = "." } dependencies = [ { name = "datashare-python" }, + { name = "huggingface-hub" }, { name = "language-data" }, { name = "pydantic-extra-types", extra = ["pycountry"] }, ] @@ -573,6 +574,7 @@ requires-dist = [ { name = "argostranslate", marker = "extra == 'c2translate'", specifier = "==1.11.0" }, { name = "argostranslate", marker = "extra == 'inference'", specifier = "==1.11.0" }, { name = "datashare-python", editable = "../../datashare-python" }, + { 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" },