| 1 | """Tiny-model fixture for integration tests. |
| 2 | |
| 3 | Mirrors ``dlm.tests.fixtures.tiny_model``: session-scoped snapshot of |
| 4 | SmolLM2-135M-Instruct, reused across the whole test run. The model is |
| 5 | small enough (~280 MB on disk, ~600 MB in fp32 VRAM) to make integration |
| 6 | tests feasible in CI. |
| 7 | |
| 8 | Tests using this fixture must carry ``@pytest.mark.slow`` and |
| 9 | ``@pytest.mark.online`` — the default test selection excludes both. |
| 10 | """ |
| 11 | |
| 12 | from __future__ import annotations |
| 13 | |
| 14 | import os |
| 15 | from collections.abc import Iterator |
| 16 | from pathlib import Path |
| 17 | |
| 18 | import pytest |
| 19 | |
| 20 | TINY_MODEL_HF_ID = "HuggingFaceTB/SmolLM2-135M-Instruct" |
| 21 | TINY_MODEL_REVISION = os.environ.get("DLM_SWAY_TINY_MODEL_REVISION", "main") |
| 22 | |
| 23 | |
| 24 | def _offline_mode() -> bool: |
| 25 | return os.environ.get("SWAY_OFFLINE", "0") == "1" |
| 26 | |
| 27 | |
| 28 | @pytest.fixture(scope="session") |
| 29 | def tiny_model_dir(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Path]: |
| 30 | """Download (or reuse) the tiny model; yield the cached directory. |
| 31 | |
| 32 | Test opts in via ``@pytest.mark.online`` — the session-wide offline |
| 33 | env vars are cleared inside this fixture so ``snapshot_download`` |
| 34 | actually fetches. |
| 35 | """ |
| 36 | from huggingface_hub import snapshot_download |
| 37 | |
| 38 | # Clear offline env guards (set by the unit-test autouse fixture). |
| 39 | prior = { |
| 40 | k: os.environ.pop(k, None) |
| 41 | for k in ("HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE", "HF_DATASETS_OFFLINE") |
| 42 | } |
| 43 | try: |
| 44 | path = snapshot_download( |
| 45 | repo_id=TINY_MODEL_HF_ID, |
| 46 | revision=TINY_MODEL_REVISION, |
| 47 | local_files_only=_offline_mode(), |
| 48 | ) |
| 49 | yield Path(path) |
| 50 | finally: |
| 51 | for k, v in prior.items(): |
| 52 | if v is not None: |
| 53 | os.environ[k] = v |