| 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 "<redacted>" 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() |
|
|
