code/self_healing_finetuning/training/hf_runtime_deps.py

from __future__ import annotations

import os
import subprocess
import sys


BASE_RUNTIME_PACKAGES = [
    # transformers>=4.54 references torch.float8_e8m0fnu (added in torch 2.7),
    # which is absent from the torch 2.6 HF Jobs image and crashes the trainer
    # at import. Keep transformers below that boundary.
    "transformers>=4.51.0,<4.54.0",
    "peft>=0.15.0,<0.19.0",
    "datasets>=3.6.0,<5.0.0",
    "accelerate>=1.5.0,<2.0.0",
    "bitsandbytes>=0.45.0,<0.50.0",
    "safetensors>=0.5.0,<0.8.0",
    "huggingface-hub>=0.30.0,<1.0.0",
]

TRL_RUNTIME_PACKAGES = [
    "trl>=0.15.0,<1.0.0",
]


def runtime_packages(*, include_trl: bool) -> list[str]:
    packages = list(BASE_RUNTIME_PACKAGES)
    if include_trl:
        packages.extend(TRL_RUNTIME_PACKAGES)
    return packages


def install_runtime_packages(*, include_trl: bool) -> None:
    """Install the HF training/eval stack with bounded versions.

    The HF Jobs image supplies PyTorch. Do not allow unbounded dependency
    upgrades here: newer Transformers builds can assume torch symbols that are
    absent from the selected PyTorch image.
    """

    if os.environ.get("SHFT_SKIP_RUNTIME_PIP_INSTALL", "").lower() in {"1", "true", "yes"}:
        print("[SHFT HF deps] Skipping runtime pip install because SHFT_SKIP_RUNTIME_PIP_INSTALL is set.")
        return
    command = [
        sys.executable,
        "-m",
        "pip",
        "install",
        "--disable-pip-version-check",
        "--upgrade",
        *runtime_packages(include_trl=include_trl),
    ]
    print("[SHFT HF deps] Installing bounded runtime stack:")
    for package in runtime_packages(include_trl=include_trl):
        print(f"[SHFT HF deps]   {package}")
    subprocess.check_call(command)