code/self_healing_finetuning/tests/test_hf_runtime_deps.py

from __future__ import annotations

import sys
import unittest
import inspect
from pathlib import Path
from types import SimpleNamespace

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from training.hf_runtime_deps import BASE_RUNTIME_PACKAGES, TRL_RUNTIME_PACKAGES, runtime_packages
from training import hf_fingpt_train
from training.hf_fingpt_train import gradient_parameter_summary, trainable_parameter_summary


class HfRuntimeDependencyPinsTest(unittest.TestCase):
    def test_transformers_is_bounded_below_float8_e8m0fnu_requirement(self) -> None:
        # transformers>=4.54 imports torch.float8_e8m0fnu, absent from the
        # torch 2.6 HF Jobs image; the pin must stay below that boundary.
        transformers = [item for item in BASE_RUNTIME_PACKAGES if item.startswith("transformers")]
        self.assertEqual(transformers, ["transformers>=4.51.0,<4.54.0"])

    def test_runtime_packages_are_version_constrained(self) -> None:
        for package in runtime_packages(include_trl=True):
            self.assertIn(">=", package, package)
            self.assertIn("<", package, package)

    def test_eval_stack_excludes_trl_but_training_stack_includes_it(self) -> None:
        self.assertNotIn(TRL_RUNTIME_PACKAGES[0], runtime_packages(include_trl=False))
        self.assertIn(TRL_RUNTIME_PACKAGES[0], runtime_packages(include_trl=True))

    def test_trainable_parameter_summary_detects_trainable_lora_params(self) -> None:
        class FakeParam:
            def __init__(self, count: int, requires_grad: bool) -> None:
                self._count = count
                self.requires_grad = requires_grad

            def numel(self) -> int:
                return self._count

        model = SimpleNamespace(
            named_parameters=lambda: iter(
                [
                    ("base_model.model.layers.0.self_attn.q_proj.base_layer.weight", FakeParam(1000, False)),
                    ("base_model.model.layers.0.self_attn.q_proj.lora_A.default.weight", FakeParam(160, True)),
                    ("base_model.model.layers.0.self_attn.q_proj.lora_B.default.weight", FakeParam(160, True)),
                ]
            )
        )

        summary = trainable_parameter_summary(model)

        self.assertTrue(summary["ok"])
        self.assertEqual(summary["total_params"], 1320)
        self.assertEqual(summary["trainable_params"], 320)
        self.assertEqual(summary["trainable_tensors"], 2)

    def test_gradient_parameter_summary_requires_nonzero_lora_gradient(self) -> None:
        class FakeGrad:
            def __init__(self, value: float) -> None:
                self.value = value

            def detach(self) -> "FakeGrad":
                return self

            def float(self) -> "FakeGrad":
                return self

            def norm(self) -> "FakeGrad":
                return self

            def item(self) -> float:
                return self.value

        class FakeParam:
            def __init__(self, requires_grad: bool, grad_value: float | None) -> None:
                self.requires_grad = requires_grad
                self.grad = None if grad_value is None else FakeGrad(grad_value)

        model = SimpleNamespace(
            named_parameters=lambda: iter(
                [
                    ("base_model.model.layers.0.self_attn.q_proj.base_layer.weight", FakeParam(False, None)),
                    ("base_model.model.layers.0.self_attn.q_proj.lora_A.default.weight", FakeParam(True, 0.0)),
                    ("base_model.model.layers.0.self_attn.q_proj.lora_B.default.weight", FakeParam(True, 0.25)),
                ]
            )
        )

        summary = gradient_parameter_summary(model)

        self.assertTrue(summary["ok"])
        self.assertEqual(summary["trainable_tensors"], 2)
        self.assertEqual(summary["grad_tensors"], 2)
        self.assertEqual(summary["nonzero_grad_tensors"], 1)
        self.assertEqual(summary["lora_trainable_tensors"], 2)
        self.assertEqual(summary["lora_nonzero_grad_tensors"], 1)

    def test_gradient_parameter_summary_blocks_all_zero_lora_gradients(self) -> None:
        class FakeGrad:
            def detach(self) -> "FakeGrad":
                return self

            def float(self) -> "FakeGrad":
                return self

            def norm(self) -> "FakeGrad":
                return self

            def item(self) -> float:
                return 0.0

        class FakeParam:
            requires_grad = True
            grad = FakeGrad()

        model = SimpleNamespace(
            named_parameters=lambda: iter(
                [
                    ("base_model.model.layers.0.self_attn.q_proj.lora_A.default.weight", FakeParam()),
                    ("base_model.model.layers.0.self_attn.q_proj.lora_B.default.weight", FakeParam()),
                ]
            )
        )

        summary = gradient_parameter_summary(model)

        self.assertFalse(summary["ok"])
        self.assertEqual(summary["nonzero_grad_tensors"], 0)
        self.assertEqual(summary["lora_nonzero_grad_tensors"], 0)

    def test_training_uses_non_reentrant_gradient_checkpointing(self) -> None:
        source = inspect.getsource(hf_fingpt_train.run_live_training)

        self.assertIn('gradient_checkpointing_kwargs={"use_reentrant": False}', source)

    def test_zero_gradient_stall_guard_is_hard_halt_enabled(self) -> None:
        source = inspect.getsource(hf_fingpt_train.run_live_training)

        self.assertIn("ZeroGradientStallCallback(local_metrics_path, hard_halt_enabled=True)", source)
        self.assertNotIn("ZeroGradientStallCallback(local_metrics_path, hard_halt_enabled=False)", source)


if __name__ == "__main__":
    unittest.main()