Spaces:
Running
Running
| from __future__ import annotations | |
| import json | |
| import sys | |
| from pathlib import Path | |
| from tempfile import TemporaryDirectory | |
| ROOT = Path(__file__).resolve().parents[1] | |
| if str(ROOT) not in sys.path: | |
| sys.path.insert(0, str(ROOT)) | |
| from training.trace_logging import TraceArtifactLogger | |
| def main() -> None: | |
| with TemporaryDirectory() as tmpdir: | |
| output_dir = Path(tmpdir) | |
| logger = TraceArtifactLogger( | |
| run_id="run-123", | |
| output_dir=output_dir, | |
| training_config={"max_steps": 6, "model_name": "demo-model"}, | |
| model_identifiers={"model_name": "demo-model", "generator_mode": "reward_aware"}, | |
| system_prompt="You are the Solver Agent.", | |
| checkpoint_interval_steps=2, | |
| ) | |
| manifest_path = output_dir / "logs" / "run_manifest.json" | |
| assert manifest_path.exists() | |
| manifest = json.loads(manifest_path.read_text(encoding="utf-8")) | |
| assert manifest["run_id"] == "run-123" | |
| assert manifest["training_config"]["max_steps"] == 6 | |
| logger.log_event( | |
| { | |
| "phase": "train", | |
| "step": 0, | |
| "train_episode_index": 1, | |
| "problem_id": "sum_even_numbers_1", | |
| "problem_family": "sum_even_numbers", | |
| "difficulty": "easy", | |
| "teacher_prompt": "Problem: sum the even numbers", | |
| "solver_completion": "print(sum(x for x in nums if x % 2 == 0))", | |
| "extracted_code": "print(sum(x for x in nums if x % 2 == 0))", | |
| "reward": 0.94, | |
| "pass_rate": 1.0, | |
| "visible_pass_rate": 1.0, | |
| "execution_status": "completed", | |
| "efficiency_score": 0.94, | |
| "optimization_hints": ["Avoid materializing temporary containers."], | |
| "feedback": "All hidden tests passed, but the solution can still be optimized further.", | |
| } | |
| ) | |
| logger.record_progress( | |
| { | |
| "phase": "train", | |
| "completed_steps": 2, | |
| "total_steps": 6, | |
| "remaining_steps": 4, | |
| "progress_ratio": 0.3333, | |
| "current_epoch": 2.0, | |
| "current_difficulty": "easy", | |
| "curriculum_level": 1, | |
| "train_episode_index": 1, | |
| "last_problem_id": "sum_even_numbers_1", | |
| "last_problem_family": "sum_even_numbers", | |
| "last_execution_status": "completed", | |
| } | |
| ) | |
| artifact_paths = logger.artifact_paths() | |
| events_path = Path(artifact_paths["events_path"]) | |
| latest_checkpoint_path = Path(artifact_paths["latest_checkpoint_path"]) | |
| assert events_path.exists() | |
| assert latest_checkpoint_path.exists() | |
| event_line = events_path.read_text(encoding="utf-8").strip().splitlines()[0] | |
| event = json.loads(event_line) | |
| assert event["problem_id"] == "sum_even_numbers_1" | |
| assert event["teacher_prompt"] == "Problem: sum the even numbers" | |
| assert "training_config" not in event | |
| checkpoint = json.loads(latest_checkpoint_path.read_text(encoding="utf-8")) | |
| assert checkpoint["step"] == 2 | |
| assert checkpoint["rolling_metrics"]["avg_reward"] == 0.94 | |
| assert "training_config" not in checkpoint | |
| reward_curve = output_dir / "reward_curve.csv" | |
| reward_curve.write_text("step,episode_reward\n0,0.94\n", encoding="utf-8") | |
| summary_paths = logger.finalize( | |
| reward_curve_csv=reward_curve, | |
| final_metrics={"completed_steps": 6}, | |
| ) | |
| summary_path = Path(summary_paths) | |
| assert summary_path.exists() | |
| summary = json.loads(summary_path.read_text(encoding="utf-8")) | |
| assert summary["final_metrics"]["completed_steps"] == 6 | |
| print("Trace logging smoke tests passed") | |
| if __name__ == "__main__": | |
| main() | |