DGXAI/driftcall

1"""Cell 01 — Install pinned dependencies.2 3Runs once at notebook boot. On Colab the notebook kernel is a bare Python 34install, so we ``pip install`` the flat pin set from ``requirements.txt``.5Locally we skip reinstall if every pin is already importable.6 7Also authenticates with the Hugging Face Hub when an ``HF_TOKEN`` environment8variable is set; on interactive sessions the user can run ``hf auth login``9separately. No network calls are attempted when ``HF_TOKEN`` is absent — the10cell remains a no-op so offline unit tests pass.11"""12 13from __future__ import annotations14 15import importlib.util16import os17import subprocess18import sys19from pathlib import Path20 21REQUIREMENTS_FILENAME = "requirements.txt"22 23# Packages whose import name differs from their distribution name. Only list24# the handful we actually probe with ``is_installed``; everything else uses25# the distribution name verbatim.26_IMPORT_ALIASES: dict[str, str] = {27    "faster-whisper": "faster_whisper",28    "huggingface_hub": "huggingface_hub",29    "uvicorn[standard]": "uvicorn",30    "pytest-cov": "pytest_cov",31}32 33 34def is_installed(distribution: str) -> bool:35    """Return True iff the import name behind *distribution* is available."""36 37    base = distribution.split("[", 1)[0].split(">", 1)[0].split("<", 1)[0]38    base = base.split("==", 1)[0].split("~=", 1)[0].strip()39    module = _IMPORT_ALIASES.get(distribution, _IMPORT_ALIASES.get(base, base))40    module = module.replace("-", "_")41    return importlib.util.find_spec(module) is not None42 43 44def _find_requirements() -> Path | None:45    """Locate ``requirements.txt`` alongside the project root (worktree-safe)."""46 47    candidates = [48        Path.cwd() / REQUIREMENTS_FILENAME,49        Path(__file__).resolve().parent.parent / REQUIREMENTS_FILENAME,50    ]51    for candidate in candidates:52        if candidate.is_file():53            return candidate54    return None55 56 57def is_colab() -> bool:58    """Detect Google Colab runtime (``google.colab`` is always importable there)."""59 60    return importlib.util.find_spec("google.colab") is not None61 62 63def pip_install(requirements_path: Path) -> int:64    """Invoke ``pip install -r <requirements_path>`` via the current interpreter."""65 66    cmd = [sys.executable, "-m", "pip", "install", "--quiet", "-r", str(requirements_path)]67    completed = subprocess.run(cmd, check=False)68    return completed.returncode69 70 71def hf_login_if_token_present() -> bool:72    """Log into HF Hub using ``HF_TOKEN`` env var. Returns True on success."""73 74    token = os.environ.get("HF_TOKEN")75    if not token:76        return False77    try:78        from huggingface_hub import login79    except ImportError:80        return False81    login(token=token, add_to_git_credential=False)82    return True83 84 85def install(force: bool = False) -> int:86    """Top-level cell body. Idempotent: skips reinstall when pins already import.87 88    :param force: Reinstall even if every dependency is importable.89    :returns: 0 when deps already satisfied or pip succeeded; non-zero on pip failure.90    """91 92    requirements_path = _find_requirements()93    if requirements_path is None:94        return 095 96    if not force and not is_colab():97        declared = [98            line.strip()99            for line in requirements_path.read_text(encoding="utf-8").splitlines()100            if line.strip() and not line.strip().startswith("#")101        ]102        if declared and all(is_installed(pkg) for pkg in declared):103            hf_login_if_token_present()104            return 0105 106    rc = pip_install(requirements_path)107    if rc == 0:108        hf_login_if_token_present()109    return rc110 111 112# Cell body: execute on import so the Colab notebook runs end-to-end.113# Skip the side effect when the cell is being imported under the pytest114# runner or when a caller opts out via ``DRIFTCALL_SKIP_INSTALL=1``.115_skip_marker = "pytest" in sys.modules or os.environ.get("DRIFTCALL_SKIP_INSTALL") == "1"116_rc = 0 if _skip_marker else install()117