diff --git a/src/pipeline/contracts.py b/src/pipeline/contracts.py index f441a7fb..4a152409 100644 --- a/src/pipeline/contracts.py +++ b/src/pipeline/contracts.py @@ -122,6 +122,14 @@ class RecipeSpec(_Strict): stages: list[StageSpec] = Field(min_length=1) gates: GateSpec negative_results: list[str] = Field(default_factory=list) + deps: list[str] = Field( + default_factory=lambda: [ + "torch==2.5.1", + "transformers==4.46.3", + "accelerate>=1.1.0", + "sentencepiece", + ] + ) channel: Literal["gh-releases"] = "gh-releases" @field_validator("corpora") diff --git a/src/pipeline/planner.py b/src/pipeline/planner.py index 79efc8c3..1ff166cb 100644 --- a/src/pipeline/planner.py +++ b/src/pipeline/planner.py @@ -36,6 +36,11 @@ def plan_stage( prefix = stage.env_prefix env = {f"{prefix}_{k.upper()}": str(v) for k, v in stage.params.items()} env[f"{prefix}_RUN"] = run + # legacy scripts need their wheels; the recipe declares them and the + # planner installs before exec (the bare-python launch, 6ad494b9, + # died on `No module named torch` — this is that scar, encoded) + bootstrap = "pip install -q " + " ".join(recipe.deps) + command = ["bash", "-c", f"{bootstrap} && python {CKPT_MOUNT}/{run}/{script}"] volumes = [Volume(CKPT_BUCKET, CKPT_MOUNT, "rw")] seen: set[str] = set() idx = 0 @@ -58,7 +63,7 @@ def plan_stage( idx += 1 return JobSpec( name=run, - command=["python", f"{CKPT_MOUNT}/{run}/{script}"], + command=command, flavor=flavor, timeout_h=timeout_h, volumes=volumes, diff --git a/tests/test_pipeline_plan_gate.py b/tests/test_pipeline_plan_gate.py index aee7cd33..95464e5a 100644 --- a/tests/test_pipeline_plan_gate.py +++ b/tests/test_pipeline_plan_gate.py @@ -26,8 +26,8 @@ def test_plan_sft_stage_builds_lawful_job() -> None: sft = recipe.stages[0] spec = plan_stage(recipe, sft, run="run-066-urd-diac-base") assert spec.name == "run-066-urd-diac-base" - assert spec.command[0] == "python" - assert spec.command[1].endswith("train_urd_diac_base_hf.py") + assert spec.command[0] == "bash" and spec.command[1] == "-c" + assert spec.command[2].endswith("train_urd_diac_base_hf.py") mounts = [v.mount for v in spec.volumes] assert "/ckpt" in mounts and "/data-dia" in mounts ro = {v.mount: v.mode for v in spec.volumes} @@ -120,3 +120,15 @@ def test_gate_results_files_end_to_end(tmp_path: Path) -> None: ) by_name = {g.name: g.passed for g in results} assert by_name == {"quality": True, "parity": True, "license": False} + + +def test_plan_wraps_legacy_deps() -> None: + """The first dogfooded launch (stage-2, 6ad494b9) died on a bare + python command — the spine must install the recipe's deps.""" + recipe = load_recipe(RECIPES / "ara-diac.yaml") + spec = plan_stage(recipe, recipe.stages[1], run="run-x") + assert spec.command[0] == "bash" and spec.command[1] == "-c" + cmd = spec.command[2] + assert cmd.startswith("pip install -q ") + assert "torch==" in cmd and "transformers==" in cmd + assert cmd.endswith("python /ckpt/run-x/continue_oracle_train_hf.py")