ktransformers/.github/model-e2e/test_contracts.py

import copy
import fcntl
import json
import subprocess
import tempfile
import unittest
import zipfile
from pathlib import Path
from unittest.mock import patch

from contracts import (
    CASES,
    PACKAGES,
    REPOSITORIES,
    digest,
    suite_passed,
    validate_answer,
    validate_lora,
    validate_request,
    verify_candidate,
    write_json,
)
from host import docker_command, export_evidence
from recipes import training_config
from resource_queue import ResourceUnavailable, gpu_busy, reservation


def request():
    return {
        "schema": 1,
        "mode": "pr",
        "harness_sha": "a" * 40,
        "pr": {"repository": REPOSITORIES["sglang"], "number": 88, "sha": "b" * 40},
        "sources": {
            key: {"repository": repo, "sha": "b" * 40}
            for key, repo in REPOSITORIES.items()
        },
        "build_run_id": 123,
        "build_workflow_sha": "c" * 40,
    }


def rank(index=0):
    return {
        "rank": index,
        "global_step": 1,
        "optimizer_steps": 1,
        "train_end": True,
        "raw_losses": [2.5],
    }


def answer(text="Paris", reason="stop", tokens=4):
    return {
        "choices": [{"message": {"content": text}, "finish_reason": reason}],
        "usage": {"completion_tokens": tokens},
    }


class Contracts(unittest.TestCase):
    def test_pypi_default(self):
        validate_request(
            {
                "schema": 1,
                "mode": "pypi",
                "version": "latest",
                "expected_version": "0.7.0.post3",
                "harness_sha": "a" * 40,
            }
        )

    def test_pypi_without_expected_version_rejected(self):
        with self.assertRaises(ValueError):
            validate_request(
                {
                    "schema": 1,
                    "mode": "pypi",
                    "version": "latest",
                    "harness_sha": "a" * 40,
                }
            )

    def test_pr_snapshot(self):
        validate_request(request())

    def test_changed_pr_head(self):
        data = request()
        data["pr"]["sha"] = "c" * 40
        with self.assertRaises(ValueError):
            validate_request(data)

    def test_foreign_repository(self):
        data = request()
        data["pr"]["repository"] = "attacker/fork"
        with self.assertRaises(ValueError):
            validate_request(data)

    def test_short_sha(self):
        data = request()
        data["harness_sha"] = "abc123"
        with self.assertRaises(ValueError):
            validate_request(data)

    def test_loss_only_requires_one_finite_step(self):
        self.assertEqual(validate_lora([rank()], 1)["status"], "passed")

    def test_all_ranks_required(self):
        with self.assertRaises(ValueError):
            validate_lora([rank()], 8)

    def test_duplicate_rank(self):
        with self.assertRaises(ValueError):
            validate_lora([rank(), rank()], 2)

    def test_bad_lora_evidence(self):
        changes = [
            {"global_step": 0},
            {"global_step": 2},
            {"optimizer_steps": 0},
            {"train_end": False},
            {"raw_losses": []},
            {"raw_losses": [float("nan")]},
            {"raw_losses": [float("inf")]},
            {"raw_losses": [True]},
        ]
        for change in changes:
            with self.subTest(change=change), self.assertRaises(ValueError):
                validate_lora([rank() | change], 1)

    def test_answer_semantic_smoke(self):
        self.assertEqual(validate_answer(answer())["status"], "passed")
        validate_answer(answer("巴黎"))

    def test_bad_answers(self):
        for value in [
            answer(""),
            answer(None),
            answer("London"),
            answer("Paris\ufffd"),
            answer("Paris\x00"),
            answer(reason="length"),
            answer(tokens=0),
        ]:
            with self.subTest(value=value), self.assertRaises(ValueError):
                validate_answer(value)

    def test_incomplete_suite_never_green(self):
        complete = {case: {"status": "passed"} for case in CASES}
        self.assertTrue(suite_passed(complete))
        del complete[CASES[0]]
        self.assertFalse(suite_passed(complete))

    def test_recipes_have_approved_boundary(self):
        for case in CASES[:2]:
            config = training_config(case, "/tmp/test")
            self.assertEqual(config["max_steps"], 1)
            self.assertEqual(config["gradient_accumulation_steps"], 1)
            self.assertFalse(config["logging_nan_inf_filter"])
            self.assertEqual(config["finetuning_type"], "lora")
            self.assertTrue(config["use_kt"])
            self.assertEqual(config["kt_config"]["kt_num_gpu_experts"], 0)


class Candidates(unittest.TestCase):
    def setUp(self):
        self.temp = tempfile.TemporaryDirectory()
        self.addCleanup(self.temp.cleanup)
        self.root = Path(self.temp.name)
        self.request = request()
        self.manifest = {
            "schema": 1,
            "sources": self.request["sources"],
            "build_run_id": 123,
            "build_workflow_sha": "c" * 40,
            "wheels": {},
        }
        for package in PACKAGES:
            name = package.replace("-", "_")
            path = self.root / f"{name}-1.0-py3-none-any.whl"
            with zipfile.ZipFile(path, "w") as wheel:
                wheel.writestr(
                    f"{name}-1.0.dist-info/METADATA", f"Name: {package}\nVersion: 1.0\n"
                )
            self.manifest["wheels"][package] = {
                "file": path.name,
                "version": "1.0",
                "sha256": digest(path),
            }
        self.save()

    def save(self):
        write_json(self.root / "candidate.json", self.manifest)

    def test_final_five_wheels(self):
        verify_candidate(self.root, self.request)

    def test_stale_main_snapshot(self):
        changed = copy.deepcopy(self.request)
        changed["sources"]["accelerate"]["sha"] = "d" * 40
        with self.assertRaises(ValueError):
            verify_candidate(self.root, changed)

    def test_wrong_build(self):
        with self.assertRaises(ValueError):
            verify_candidate(self.root, self.request | {"build_run_id": 999})

    def test_tampered_wheel(self):
        path = self.root / self.manifest["wheels"]["ktransformers"]["file"]
        path.write_bytes(b"changed")
        with self.assertRaises(ValueError):
            verify_candidate(self.root, self.request)

    def test_path_traversal(self):
        self.manifest["wheels"]["ktransformers"]["file"] = "../escape.whl"
        self.save()
        with self.assertRaises(ValueError):
            verify_candidate(self.root, self.request)

    def test_symlink_wheel(self):
        path = self.root / self.manifest["wheels"]["ktransformers"]["file"]
        renamed = path.with_suffix(".real")
        path.rename(renamed)
        path.symlink_to(renamed)
        with self.assertRaises(ValueError):
            verify_candidate(self.root, self.request)

    def test_wrong_metadata_version(self):
        self.manifest["wheels"]["ktransformers"]["version"] = "999"
        self.save()
        with self.assertRaises(ValueError):
            verify_candidate(self.root, self.request)

    def test_raw_native_extra_rejected(self):
        (self.root / "sgl_kernel_kt-1.0.whl").write_bytes(b"raw")
        with self.assertRaises(ValueError):
            verify_candidate(self.root, self.request)


class QueueAndSandbox(unittest.TestCase):
    def setUp(self):
        self.temp = tempfile.TemporaryDirectory()
        self.addCleanup(self.temp.cleanup)
        self.root = Path(self.temp.name)

    def test_busy_then_idle_waits(self):
        states = iter([(True, {}), (False, {}), (False, {})])
        with patch("resource_queue.time.sleep") as sleep:
            with reservation(
                self.root / "lock",
                self.root / "queue",
                interval=0,
                idle_samples=2,
                probe=lambda: next(states),
            ):
                self.assertEqual(sleep.call_count, 2)
        self.assertTrue((self.root / "lock").exists())

    def test_queue_timeout_is_not_a_model_failure(self):
        with self.assertRaises(ResourceUnavailable):
            with reservation(self.root / "lock", self.root / "queue", timeout=0):
                self.fail("should not start")

    def test_lock_held_through_work_and_released_on_exception(self):
        path = self.root / "lock"
        with self.assertRaisesRegex(RuntimeError, "owned failure"):
            with reservation(
                path, self.root / "queue", idle_samples=1, probe=lambda: (False, {})
            ):
                with path.open("a") as other:
                    with self.assertRaises(BlockingIOError):
                        fcntl.flock(other, fcntl.LOCK_EX | fcntl.LOCK_NB)
                raise RuntimeError("owned failure")
        with path.open("a") as other:
            fcntl.flock(other, fcntl.LOCK_EX | fcntl.LOCK_NB)

    def test_compute_process_always_queues(self):
        with patch(
            "resource_queue.subprocess.check_output", side_effect=["123\n", "2, 0\n"]
        ):
            self.assertTrue(gpu_busy()[0])

    def test_non_compute_vram_queues(self):
        with patch(
            "resource_queue.subprocess.check_output", side_effect=["", "4096, 0\n"]
        ):
            self.assertTrue(gpu_busy()[0])

    def test_idle(self):
        with patch(
            "resource_queue.subprocess.check_output", side_effect=["", "2, 0\n2, 0\n"]
        ):
            self.assertFalse(gpu_busy()[0])

    def test_driver_error_never_means_idle(self):
        with patch(
            "resource_queue.subprocess.check_output",
            side_effect=subprocess.CalledProcessError(1, "nvidia-smi"),
        ):
            with self.assertRaises(subprocess.CalledProcessError):
                gpu_busy()

    def test_unknown_stats_never_mean_idle(self):
        with patch(
            "resource_queue.subprocess.check_output", side_effect=["", "N/A, N/A\n"]
        ):
            with self.assertRaises(ResourceUnavailable):
                gpu_busy()

    def test_container_has_no_host_secrets_or_privileges(self):
        config = {
            "image": "test@sha256:" + "a" * 64,
            "models": {"glm53": {"path": "/mnt/models/glm53"}},
        }
        command = docker_command(
            config, "/harness", "/input", Path("/output"), "owned-name"
        )
        self.assertIn("--cap-drop=ALL", command)
        self.assertIn("--read-only", command)
        self.assertIn("--no-healthcheck", command)
        self.assertNotIn("--privileged", command)
        self.assertNotIn("--network=host", command)
        self.assertNotIn("--ipc=host", command)
        self.assertIn(
            "type=bind,src=/mnt/models/glm53,dst=/models/glm53,readonly", command
        )
        self.assertNotIn("docker.sock", " ".join(command))
        self.assertIn("--device=nvidia.com/gpu=all", command)
        self.assertIn("NVIDIA_VISIBLE_DEVICES=void", command)
        self.assertNotIn("--gpus=all", command)

    def test_rootless_user_mapping(self):
        config = {"rootless": True, "image": "test@sha256:" + "a" * 64, "models": {}}
        command = docker_command(
            config, "/harness", "/input", Path("/output"), "owned-name"
        )
        self.assertEqual(command[command.index("--user") + 1], "0:0")

    def test_evidence_export_refuses_symlinks(self):
        source, destination = self.root / "source", self.root / "export"
        source.mkdir()
        destination.mkdir()
        (source / "leak").symlink_to("/etc/passwd")
        with self.assertRaises(ValueError):
            export_evidence(source, destination)

    def test_regular_evidence_export(self):
        source, destination = self.root / "source", self.root / "export"
        source.mkdir()
        destination.mkdir()
        write_json(source / "suite.json", {"status": "failed"})
        export_evidence(source, destination)
        self.assertEqual(
            json.loads((destination / "suite.json").read_text())["status"], "failed"
        )


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