from __future__ import annotations import json import os import subprocess import sys from pathlib import Path import pytest from scripts.run_manifest import PlannedJob, execute_local_jobs, load_manifest, plan_jobs from scripts.train_dovla import _parse_loss_weights def test_manifest_dry_run_writes_resolved_manifest(tmp_path: Path) -> None: out_dir = tmp_path / "run" result = subprocess.run( [ sys.executable, "scripts/run_manifest.py", "manifests/scaling_k_sweep.yaml", "--dry-run", "--out", str(out_dir), ], check=True, capture_output=True, text=True, ) assert "manifest: scaling_k_sweep" in result.stdout assert "planned jobs:" in result.stdout assert (out_dir / "resolved_manifest.yaml").exists() jobs_path = out_dir / "planned_jobs.json" assert jobs_path.exists() jobs = json.loads(jobs_path.read_text(encoding="utf-8")) assert any(job["stage"] == "dataset_generation" for job in jobs) assert any(job["stage"] == "training" for job in jobs) assert any(job["stage"] == "evaluation" for job in jobs) assert any(job["stage"] == "scaling_sweeps" for job in jobs) generation_job = next(job for job in jobs if job["stage"] == "dataset_generation") assert "--num-tasks" in generation_job["command"] def test_manifest_redacts_secrets_from_outputs(tmp_path: Path) -> None: secret = "manifest_secret_value" manifest = tmp_path / "manifest.yaml" manifest.write_text( """ name: secret_test run_dir: ${TEST_RUN_DIR:-runs/secret_test} dataset_generation: backend: toy simulator_params: {} task_source: builtins num_tasks: 1 num_states_per_task: 1 k: 2 shard_size: 8 output_path: outputs/secret_test/cil seed: 0 vlm_annotation: enabled: false cache_path: outputs/secret_test/cache.json model_env_var: OPENCLAUDE_MODEL api_key: ${OPENCLAUDE_API_KEY} training: model_size: tiny hidden_dim: 32 batch_groups: 1 records_per_group: 2 learning_rate: 0.001 loss_weights: {bc: 1.0} epochs: 1 steps: null checkpoint_path: outputs/secret_test/train/best.pt evaluation: causalstress: enabled: true backend: toy num_tasks: 1 k: 2 output_path: outputs/secret_test/eval/causalstress.json libero: {enabled: false, placeholder: true} maniskill: {enabled: false, placeholder: true} simpler: {enabled: false, placeholder: true} baselines: enabled: false output_root: outputs/secret_test/baselines names: [] scaling_sweeps: enabled: false output_path: outputs/secret_test/scaling total_records: 4 k_values: [1, 2] epochs: 1 """, encoding="utf-8", ) out_dir = tmp_path / "run" env = {**os.environ, "OPENCLAUDE_API_KEY": secret, "TEST_RUN_DIR": str(out_dir)} result = subprocess.run( [ sys.executable, "scripts/run_manifest.py", str(manifest), "--dry-run", "--emit-slurm", "--out", str(out_dir), ], check=True, capture_output=True, text=True, env=env, ) all_text = result.stdout + result.stderr all_text += (out_dir / "resolved_manifest.yaml").read_text(encoding="utf-8") all_text += (out_dir / "planned_jobs.json").read_text(encoding="utf-8") for path in (out_dir / "slurm").glob("*.sbatch"): all_text += path.read_text(encoding="utf-8") assert secret not in all_text assert "" in (out_dir / "resolved_manifest.yaml").read_text(encoding="utf-8") def test_manifest_files_exist() -> None: for path in ( "manifests/cil_160m.yaml", "manifests/cil_1b_template.yaml", "manifests/scaling_k_sweep.yaml", "manifests/baselines_full.yaml", ): assert Path(path).exists() def test_large_manifests_plan_measured_maniskill_generation() -> None: expected_budgets = { "manifests/cil_160m.yaml": 160_000_000, "manifests/cil_1b_template.yaml": 1_000_000_000, } for path, expected_budget in expected_budgets.items(): manifest = load_manifest(path) generation = manifest["dataset_generation"] assert ( generation["num_tasks"] * generation["num_states_per_task"] * generation["k"] == expected_budget ) job = next(job for job in plan_jobs(manifest) if job.stage == "dataset_generation") assert job.command[1] == "scripts/generate_maniskill_lattice.py" assert "--demo" in job.command assert "--state-batch-size" in job.command assert "--parallel-branches" in job.command groups_index = job.command.index("--num-groups") + 1 assert int(job.command[groups_index]) * generation["k"] == expected_budget evaluation_job = next( item for item in plan_jobs(manifest) if item.name == "eval_causalstress" ) assert not evaluation_job.local_executable def test_training_plan_carries_manifest_loss_weights() -> None: manifest = load_manifest("manifests/scaling_k_sweep.yaml") job = next(job for job in plan_jobs(manifest) if job.stage == "training") encoded = [ job.command[index + 1] for index, value in enumerate(job.command[:-1]) if value == "--loss-weight" ] assert "bc=1.0" in encoded assert "rank=1.0" in encoded parsed = _parse_loss_weights(encoded) assert parsed.bc == 1.0 assert parsed.rank == 1.0 def test_manifest_validation_rejects_invalid_budget(tmp_path: Path) -> None: payload = Path("manifests/scaling_k_sweep.yaml").read_text(encoding="utf-8") invalid = tmp_path / "invalid.yaml" invalid.write_text(payload.replace("num_tasks: 10", "num_tasks: 0"), encoding="utf-8") with pytest.raises(ValueError, match="num_tasks"): load_manifest(invalid) def test_emitted_slurm_directives_are_concrete(tmp_path: Path) -> None: out_dir = tmp_path / "run" subprocess.run( [ sys.executable, "scripts/run_manifest.py", "manifests/scaling_k_sweep.yaml", "--dry-run", "--emit-slurm", "--out", str(out_dir), ], check=True, capture_output=True, text=True, ) scripts = list((out_dir / "slurm").glob("*.sbatch")) assert scripts for script in scripts: text = script.read_text(encoding="utf-8") directives = "\n".join(line for line in text.splitlines() if line.startswith("#SBATCH")) assert "${" not in directives def test_local_execution_uses_current_interpreter(tmp_path: Path) -> None: result_path = tmp_path / "interpreter.txt" job = PlannedJob( name="interpreter_probe", stage="test", command=[ "python", "-c", f"import pathlib,sys; pathlib.Path({str(result_path)!r}).write_text(sys.executable)", ], local_executable=True, ) execute_local_jobs([job]) assert Path(result_path.read_text(encoding="utf-8")).resolve() == Path(sys.executable).resolve()