Source code for coseda.ici.run_contract

"""Canonical file layout helpers for COSEDA ICI runs."""

from __future__ import annotations

import os
import json
import shutil
import datetime
from pathlib import Path

from coseda.logging_utils import log_print

CELL_FILENAME = "cellfile.cell"
GEOMETRY_FILENAME = "geometry.geom"
SETTINGS_FILENAME = "ici_settings.json"
ORCHESTRATOR_LOG_FILENAME = "orchestrator.log"
IMAGE_RUN_LOG_FILENAME = "image_run_log.csv"
IMAGE_RUN_STATE_FILENAME = "image_run_state.json"
EARLY_BREAK_STREAM_FILENAME = "early_break.stream"
DONE_STREAM_FILENAME = "done.stream"
MANIFEST_FILENAME = "ici_run_manifest.json"


[docs] def run_root_path(run_root: str) -> str: """Return the normalized top-level ICI run directory.""" return os.path.abspath(os.path.expanduser(run_root))
[docs] def run_name(run_root: str) -> str: """Return the ICI run folder name.""" return os.path.basename(run_root_path(run_root).rstrip(os.sep))
[docs] def canonical_stream_filename(run_root: str) -> str: """Return the COSEDA-compatible final stream filename for an ICI run.""" return f"{run_name(run_root)}.stream"
[docs] def canonical_stream_path(run_root: str) -> str: """Return the top-level stream path expected by COSEDA run consumers.""" return os.path.join(run_root_path(run_root), canonical_stream_filename(run_root))
[docs] def early_break_stream_path(run_root: str) -> str: return os.path.join(run_root_path(run_root), EARLY_BREAK_STREAM_FILENAME)
[docs] def done_stream_path(run_root: str) -> str: return os.path.join(run_root_path(run_root), DONE_STREAM_FILENAME)
[docs] def manifest_path(run_root: str) -> str: return os.path.join(run_root_path(run_root), MANIFEST_FILENAME)
[docs] def resolve_final_stream_path(run_root: str) -> str | None: """ Resolve the best available final stream for an ICI run. Preference is converged output, then current best early-break output, then the canonical alias for backwards compatibility with existing COSEDA code. """ for candidate in ( done_stream_path(run_root), early_break_stream_path(run_root), canonical_stream_path(run_root), ): if os.path.isfile(candidate): return candidate return None
def _replace_with_relative_symlink(target: str, link_path: str) -> bool: """Create a relative symlink, replacing an existing alias if possible.""" link = Path(link_path) target_path = Path(target) try: if link.exists() or link.is_symlink(): link.unlink() rel_target = os.path.relpath(target_path, start=link.parent) link.symlink_to(rel_target) return True except OSError: return False
[docs] def update_final_stream_alias(run_root: str) -> str | None: """ Maintain the canonical ``<ici_run_folder>.stream`` alias. ICI keeps its native streams (``early_break.stream`` and ``done.stream``), while this alias makes the run look like a normal COSEDA indexing run to existing downstream tools. """ source = resolve_final_stream_path(run_root) if source is None: return None alias = canonical_stream_path(run_root) if os.path.abspath(source) == os.path.abspath(alias): return alias if _replace_with_relative_symlink(source, alias): log_print(f"[ici-contract] Updated final stream alias: {alias} -> {os.path.basename(source)}") return alias try: shutil.copy2(source, alias) log_print(f"[ici-contract] Updated final stream copy: {alias} from {os.path.basename(source)}") return alias except OSError as exc: log_print(f"[ici-contract] WARNING: failed to update final stream alias {alias}: {exc}") return None
def _existing_path(path: str) -> str | None: return path if os.path.exists(path) else None def _latest_run_number(run_root: str) -> int | None: root = run_root_path(run_root) latest = None try: for name in os.listdir(root): if not name.startswith("run_"): continue suffix = name[4:] if len(suffix) != 3 or not suffix.isdigit(): continue full = os.path.join(root, name) if not os.path.isdir(full): continue value = int(suffix) latest = value if latest is None else max(latest, value) except OSError: return None return latest def _load_image_state_summary(run_root: str) -> dict: state_path = os.path.join(run_root_path(run_root), IMAGE_RUN_STATE_FILENAME) summary = { "state_file": _existing_path(state_path), "total_events": 0, "done_events": 0, "events_with_trials": 0, } if not os.path.isfile(state_path): return summary try: with open(state_path, "r", encoding="utf-8") as handle: state = json.load(handle) except (OSError, json.JSONDecodeError): return summary events = state.get("events", {}) if not isinstance(events, dict): return summary summary["total_events"] = len(events) done_events = 0 events_with_trials = 0 for event_state in events.values(): if not isinstance(event_state, dict): continue if event_state.get("trials"): events_with_trials += 1 latest_status = event_state.get("latest_status", ["", ""]) if ( isinstance(latest_status, (list, tuple)) and len(latest_status) >= 2 and str(latest_status[0]).lower() == "done" and str(latest_status[1]).lower() == "done" ): done_events += 1 summary["done_events"] = done_events summary["events_with_trials"] = events_with_trials return summary def _json_safe(value): if value is None or isinstance(value, (str, int, float, bool)): return value if isinstance(value, dict): return {str(k): _json_safe(v) for k, v in value.items()} if isinstance(value, (list, tuple)): return [_json_safe(v) for v in value] return str(value)
[docs] def write_run_manifest( run_root: str, *, geom: str | None = None, cell: str | None = None, h5_sources: list[str] | tuple[str, ...] | None = None, flags: list[str] | tuple[str, ...] | None = None, params: dict | None = None, argv: list[str] | tuple[str, ...] | None = None, phase: str | None = None, status: str | None = None, exit_code: int | None = None, ) -> str: """Write or refresh the top-level ICI run manifest.""" root = run_root_path(run_root) os.makedirs(root, exist_ok=True) path = manifest_path(root) existing = {} if os.path.isfile(path): try: with open(path, "r", encoding="utf-8") as handle: existing = json.load(handle) except (OSError, json.JSONDecodeError): existing = {} existing_inputs = existing.get("inputs", {}) if isinstance(existing.get("inputs"), dict) else {} existing_settings = existing.get("settings", {}) if isinstance(existing.get("settings"), dict) else {} existing_params = existing_settings.get("params", {}) existing_flags = existing_settings.get("flags", []) existing_argv = existing_settings.get("argv", []) geom = geom if geom is not None else existing_inputs.get("geometry") cell = cell if cell is not None else existing_inputs.get("cell") h5_sources = h5_sources if h5_sources is not None else existing_inputs.get("h5_sources", []) flags = flags if flags is not None else existing_flags params = params if params is not None else existing_params argv = argv if argv is not None else existing_argv alias = canonical_stream_path(root) final_source = resolve_final_stream_path(root) latest_run = _latest_run_number(root) converged = os.path.isfile(done_stream_path(root)) stream_alias_exists = os.path.exists(alias) files = { "cell": _existing_path(cell or os.path.join(root, CELL_FILENAME)), "geometry": _existing_path(geom or os.path.join(root, GEOMETRY_FILENAME)), "settings": _existing_path(os.path.join(root, SETTINGS_FILENAME)), "orchestrator_log": _existing_path(os.path.join(root, ORCHESTRATOR_LOG_FILENAME)), "image_run_log": _existing_path(os.path.join(root, IMAGE_RUN_LOG_FILENAME)), "image_run_state": _existing_path(os.path.join(root, IMAGE_RUN_STATE_FILENAME)), "early_break_stream": _existing_path(early_break_stream_path(root)), "done_stream": _existing_path(done_stream_path(root)), "canonical_stream": _existing_path(alias), "final_stream": alias if stream_alias_exists else final_source, "final_stream_source": final_source, } manifest = { "schema_version": 1, "run_type": "ici", "run_name": run_name(root), "run_root": root, "updated_at": datetime.datetime.now().isoformat(timespec="seconds"), "phase": phase, "status": status or ("converged" if converged else "running"), "exit_code": exit_code, "latest_run": latest_run, "converged": converged, "files": files, "inputs": { "h5_sources": [os.path.abspath(os.path.expanduser(str(p))) for p in (h5_sources or [])], "geometry": files["geometry"] or geom, "cell": files["cell"] or cell, }, "settings": { "params": _json_safe(params or {}), "flags": _json_safe(list(flags or [])), "argv": _json_safe(list(argv or [])), }, "summary": _load_image_state_summary(root), } tmp_path = path + ".tmp" with open(tmp_path, "w", encoding="utf-8") as handle: json.dump(manifest, handle, indent=2, sort_keys=True) handle.write("\n") os.replace(tmp_path, path) log_print(f"[ici-contract] Wrote run manifest: {path}") return path