|
| 1 | +"""On-disk artifact management for optimization runs. |
| 2 | +
|
| 3 | +Centralizes the directory layout and write logic for all artifacts |
| 4 | +produced during an optimization run under .runs/<run_id>/. |
| 5 | +""" |
| 6 | + |
| 7 | +import json |
| 8 | +import pathlib |
| 9 | +from datetime import datetime |
| 10 | + |
| 11 | + |
| 12 | +def _sanitize_artifact_path(path_value: str) -> pathlib.Path: |
| 13 | + """Convert a source path into a safe relative artifact path. |
| 14 | +
|
| 15 | + Strips traversal components (..), absolute prefixes, and Windows |
| 16 | + drive letters so that artifacts are always written under the |
| 17 | + intended directory. |
| 18 | + """ |
| 19 | + normalized = path_value.replace("\\", "/") |
| 20 | + parts = pathlib.PurePosixPath(normalized).parts |
| 21 | + safe_parts: list[str] = [] |
| 22 | + for part in parts: |
| 23 | + if part in ("", ".", "/"): |
| 24 | + continue |
| 25 | + if part == "..": |
| 26 | + continue |
| 27 | + if not safe_parts and ":" in part: |
| 28 | + part = part.replace(":", "_") |
| 29 | + safe_parts.append(part) |
| 30 | + |
| 31 | + if not safe_parts: |
| 32 | + return pathlib.Path("unnamed_file") |
| 33 | + return pathlib.Path(*safe_parts) |
| 34 | + |
| 35 | + |
| 36 | +class RunArtifacts: |
| 37 | + """Manages the on-disk artifact layout for a single optimization run. |
| 38 | +
|
| 39 | + Layout:: |
| 40 | +
|
| 41 | + <root>/ |
| 42 | + steps/<step>/ |
| 43 | + files/<relative_path> # actual code files |
| 44 | + manifest.json # machine-readable index |
| 45 | + best/ |
| 46 | + files/<relative_path> |
| 47 | + manifest.json |
| 48 | + outputs/ |
| 49 | + step_<n>.out.txt # execution stdout/stderr |
| 50 | + exec_output.jsonl # centralized output index |
| 51 | + """ |
| 52 | + |
| 53 | + def __init__(self, log_dir: str, run_id: str) -> None: |
| 54 | + self.root = pathlib.Path(log_dir) / run_id |
| 55 | + self.root.mkdir(parents=True, exist_ok=True) |
| 56 | + |
| 57 | + # ------------------------------------------------------------------ |
| 58 | + # Code snapshots |
| 59 | + # ------------------------------------------------------------------ |
| 60 | + |
| 61 | + def save_step_code(self, step: int, file_map: dict[str, str]) -> pathlib.Path: |
| 62 | + """Write code snapshot + manifest for a given step. |
| 63 | +
|
| 64 | + Returns the bundle directory path. |
| 65 | + """ |
| 66 | + return self._write_code_bundle(file_map, label=("steps", str(step))) |
| 67 | + |
| 68 | + def save_best_code(self, file_map: dict[str, str]) -> pathlib.Path: |
| 69 | + """Write code snapshot + manifest for the best result. |
| 70 | +
|
| 71 | + Returns the bundle directory path. |
| 72 | + """ |
| 73 | + return self._write_code_bundle(file_map, label=("best",)) |
| 74 | + |
| 75 | + # ------------------------------------------------------------------ |
| 76 | + # Execution output |
| 77 | + # ------------------------------------------------------------------ |
| 78 | + |
| 79 | + def save_execution_output(self, step: int, output: str) -> None: |
| 80 | + """Save execution output as a per-step file and append to the JSONL index.""" |
| 81 | + timestamp = datetime.now().isoformat() |
| 82 | + |
| 83 | + outputs_dir = self.root / "outputs" |
| 84 | + # Keep raw execution output per step for easy local inspection. |
| 85 | + outputs_dir.mkdir(parents=True, exist_ok=True) |
| 86 | + |
| 87 | + step_file = outputs_dir / f"step_{step}.out.txt" |
| 88 | + # Store full stdout/stderr for this exact step. |
| 89 | + step_file.write_text(output, encoding="utf-8") |
| 90 | + |
| 91 | + jsonl_file = self.root / "exec_output.jsonl" |
| 92 | + entry = { |
| 93 | + "step": step, |
| 94 | + "timestamp": timestamp, |
| 95 | + "output_file": step_file.relative_to(self.root).as_posix(), |
| 96 | + "output_length": len(output), |
| 97 | + } |
| 98 | + # Append compact metadata so tooling can stream/index outputs. |
| 99 | + with open(jsonl_file, "a", encoding="utf-8") as f: |
| 100 | + f.write(json.dumps(entry) + "\n") |
| 101 | + |
| 102 | + # ------------------------------------------------------------------ |
| 103 | + # Internal helpers |
| 104 | + # ------------------------------------------------------------------ |
| 105 | + |
| 106 | + def _write_code_bundle(self, file_map: dict[str, str], label: tuple[str, ...]) -> pathlib.Path: |
| 107 | + bundle_dir = self.root.joinpath(*label) |
| 108 | + files_dir = bundle_dir / "files" |
| 109 | + files_dir.mkdir(parents=True, exist_ok=True) |
| 110 | + |
| 111 | + files_manifest: list[dict[str, str | int]] = [] |
| 112 | + for source_path, content in sorted(file_map.items()): |
| 113 | + artifact_rel = _sanitize_artifact_path(source_path) |
| 114 | + artifact_path = files_dir / artifact_rel |
| 115 | + artifact_path.parent.mkdir(parents=True, exist_ok=True) |
| 116 | + artifact_path.write_text(content, encoding="utf-8") |
| 117 | + files_manifest.append( |
| 118 | + {"path": source_path, "artifact_path": artifact_rel.as_posix(), "bytes": len(content.encode("utf-8"))} |
| 119 | + ) |
| 120 | + |
| 121 | + is_step = label[0] == "steps" |
| 122 | + manifest: dict = { |
| 123 | + "type": "step_code_snapshot" if is_step else "best_code_snapshot", |
| 124 | + "created_at": datetime.now().isoformat(), |
| 125 | + "file_count": len(files_manifest), |
| 126 | + "files": files_manifest, |
| 127 | + } |
| 128 | + if is_step: |
| 129 | + manifest["step"] = int(label[1]) |
| 130 | + |
| 131 | + manifest_path = bundle_dir / "manifest.json" |
| 132 | + manifest_path.write_text(json.dumps(manifest, indent=2) + "\n", encoding="utf-8") |
| 133 | + return bundle_dir |
0 commit comments