Skills › Scientific Agent Skills › fluidsim › restart_compatibility.py
restart_compatibility.py
482 lines · skills/fluidsim/scripts/restart_compatibility.py at 5f2bc03
1#!/usr/bin/env python3
2"""Check FluidSim restart metadata against a validated target configuration."""
4from __future__ import annotations
6import argparse
7import math
8import re
9from collections.abc import Mapping
10from pathlib import Path
11from typing import Any
13try:
14 from ._common import (
15 GIB,
16 ToolError,
17 bounded_int,
18 checked_input,
19 checked_root,
20 emit_json,
21 fail_json,
22 finite_float,
23 iter_local_files,
24 load_json,
25 relative_display,
26 sha256_file,
27 validate_keys,
28 )
29 from ._schema import FLUIDSIM_VERSION, SOLVER_IMPORTS, normalized_copy
30 from ._profiles import PROFILES
31except ImportError: # Direct script execution.
32 from _common import (
33 GIB,
34 ToolError,
35 bounded_int,
36 checked_input,
37 checked_root,
38 emit_json,
39 fail_json,
40 finite_float,
41 iter_local_files,
42 load_json,
43 relative_display,
44 sha256_file,
45 validate_keys,
46 )
47 from _schema import FLUIDSIM_VERSION, SOLVER_IMPORTS, normalized_copy
48 from _profiles import PROFILES
51TOOL = "fluidsim-restart-compatibility"
52_MODULE_TO_SOLVER = {module: key for key, module in SOLVER_IMPORTS.items()}
53_STATE_SUFFIXES = (".nc", ".h5", ".hdf5")
54_COMPATIBLE_OPER_KEYS = ("nx", "ny", "nz", "Lx", "Ly", "Lz")
55_PHYSICS_KEYS = ("N", "beta", "c2", "f", "nu_2", "nu_4", "nu_8", "nu_m4")
58def _safe_scalar(value: Any) -> Any:
59 if hasattr(value, "item"):
60 try:
61 value = value.item()
62 except (TypeError, ValueError):
63 return None
64 if isinstance(value, bytes):
65 return value.decode("utf-8", errors="replace")[:500]
66 if value is None or isinstance(value, (bool, int, float, str)):
67 if isinstance(value, (int, float)) and not isinstance(value, bool):
68 finite_float(value, name="restart metadata")
69 return value if not isinstance(value, str) else value[:500]
70 return None
73def _attrs(group: Any, names: tuple[str, ...]) -> dict[str, Any]:
74 result: dict[str, Any] = {}
75 for name in names:
76 if name not in group.attrs:
77 continue
78 identifier = group.attrs.get_id(name)
79 if identifier.shape != ():
80 continue
81 value = _safe_scalar(group.attrs[name])
82 if value is not None:
83 result[name] = value
84 return result
87def _find_state_file(path: Path) -> Path:
88 if path.is_file():
89 return path
90 candidates = [
91 item
92 for item in iter_local_files(
93 path, suffixes=_STATE_SUFFIXES, max_files=256, recursive=False
94 )
95 if item.name.casefold().startswith("state_phys")
96 ]
97 if not candidates:
98 raise ToolError("no state_phys .nc/.h5 file found in restart directory")
99 ranked = []
100 for candidate in candidates:
101 state = load_hdf5_state(candidate, digest_limit=0)["state"]
102 time, iteration = state["time"], state["iteration"]
103 if not isinstance(time, (int, float)) or not math.isfinite(time):
104 raise ToolError("checkpoint time is missing/nonfinite; select and inspect an explicit file")
105 if not isinstance(iteration, int) or isinstance(iteration, bool) or iteration < 0:
106 raise ToolError("checkpoint iteration is missing/invalid")
107 ranked.append((time, iteration, candidate))
108 ranked.sort()
109 if len(ranked) > 1 and ranked[-1][:2] == ranked[-2][:2]:
110 raise ToolError("multiple checkpoints share the latest time/iteration; select an explicit file")
111 return ranked[-1][2]
114def _hard_child(group: Any, name: str, h5py: Any) -> Any:
115 """Access one local object without following a soft/external HDF5 link."""
116 link = group.get(name, getlink=True)
117 if link is None:
118 return None
119 if not isinstance(link, h5py.HardLink):
120 raise ToolError("restart metadata must not contain soft or external links")
121 return group[name]
124def load_hdf5_state(path: Path, *, digest_limit: int) -> dict[str, Any]:
125 """Read restart metadata and dataset names without reading field arrays."""
127 try:
128 import h5py # Lazy optional dependency.
129 except ImportError as exc:
130 raise ToolError(
131 "restart HDF5 inspection requires optional h5py from the pinned environment"
132 ) from exc
134 try:
135 with h5py.File(path, "r") as handle:
136 state_group = _hard_child(handle, "state_phys", h5py)
137 if not isinstance(state_group, h5py.Group):
138 raise ToolError("restart file has no /state_phys group")
139 datasets = {}
140 for count, name in enumerate(state_group, start=1):
141 if count > 128:
142 raise ToolError("restart dataset-count limit exceeded")
143 obj = _hard_child(state_group, name, h5py)
144 if isinstance(obj, h5py.Dataset):
145 if obj.is_virtual or obj.external:
146 raise ToolError("restart state uses external or virtual dataset storage")
147 datasets[name] = {"shape": list(obj.shape), "dtype": str(obj.dtype), "numeric": obj.dtype.kind in "iufc"}
148 state_params = _hard_child(handle, "state_params", h5py)
149 forcing_state = _hard_child(state_params, "forcing", h5py) if isinstance(state_params, h5py.Group) else None
150 forcing_attrs = _attrs(forcing_state, ("seed0", "seed1", "t_last_change")) if isinstance(forcing_state, h5py.Group) else {}
151 state = {
152 "datasets": sorted(datasets),
153 "dataset_metadata": datasets,
154 "iteration": _attrs(state_group, ("it",)).get("it"),
155 "state_parameters_present": isinstance(state_params, h5py.Group),
156 "forcing_state_complete": (
157 len(forcing_attrs) == 3
158 and all(isinstance(forcing_attrs[key], int) and not isinstance(forcing_attrs[key], bool) and 0 <= forcing_attrs[key] < 2**32 for key in ("seed0", "seed1"))
159 and isinstance(forcing_attrs["t_last_change"], (int, float))
160 ),
161 "time": _attrs(state_group, ("time",)).get("time"),
162 }
163 parameters: dict[str, Any] = {}
164 solver = None
165 source_version = None
166 info = _hard_child(handle, "info_simul", h5py)
167 if isinstance(info, h5py.Group):
168 params_group = _hard_child(info, "params", h5py)
169 if isinstance(params_group, h5py.Group):
170 parameters.update(_attrs(params_group, _PHYSICS_KEYS))
171 for child, keys in (
172 ("oper", _COMPATIBLE_OPER_KEYS + ("coef_dealiasing", "type_fft")),
173 (
174 "time_stepping",
175 ("t_end", "it_end", "type_time_scheme"),
176 ),
177 ("forcing", ("enable", "type", "forcing_rate")),
178 ):
179 child_group = _hard_child(params_group, child, h5py)
180 if isinstance(child_group, h5py.Group):
181 parameters[child] = _attrs(child_group, keys)
182 solver_group = _hard_child(info, "solver", h5py)
183 if isinstance(solver_group, h5py.Group):
184 solver_attrs = _attrs(
185 solver_group,
186 ("module_name", "short_name", "version", "fluidsim"),
187 )
188 module_name = solver_attrs.get("module_name")
189 solver = _MODULE_TO_SOLVER.get(str(module_name), None)
190 source_version = solver_attrs.get(
191 "fluidsim", solver_attrs.get("version")
192 )
193 except OSError as exc:
194 raise ToolError("restart file is not readable HDF5/netCDF4") from exc
196 return {
197 "parameters": parameters,
198 "provenance": {
199 "fluidfft": None,
200 "fluidsim": source_version,
201 "state_sha256": sha256_file(path, max_bytes=digest_limit),
202 "digest_origin": "computed",
203 },
204 "solver": solver,
205 "state": state,
206 }
209def load_manifest(path: Path) -> dict[str, Any]:
210 document = load_json(path)
211 if not isinstance(document, Mapping):
212 raise ToolError("restart manifest must be a JSON object")
213 validate_keys(
214 document,
215 allowed={"schema_version", "solver", "parameters", "state", "provenance"},
216 required={"schema_version", "solver", "parameters", "state", "provenance"},
217 context="restart manifest",
218 )
219 if document["schema_version"] != "1.1":
220 raise ToolError("restart manifest schema_version must be '1.1'")
221 if not isinstance(document["solver"], str) or document["solver"] not in SOLVER_IMPORTS:
222 raise ToolError("restart manifest solver is unknown")
223 parameters = document["parameters"]
224 state = document["state"]
225 provenance = document["provenance"]
226 if not all(isinstance(item, Mapping) for item in (parameters, state, provenance)):
227 raise ToolError("restart manifest parameters/state/provenance must be objects")
228 validate_keys(
229 state,
230 allowed={"datasets", "dataset_metadata", "forcing_state_complete", "iteration", "state_parameters_present", "time"},
231 required={"datasets", "iteration", "state_parameters_present", "time"},
232 context="restart manifest state",
233 )
234 validate_keys(
235 provenance,
236 allowed={"fluidfft", "fluidsim", "state_sha256"},
237 required={"fluidfft", "fluidsim", "state_sha256"},
238 context="restart manifest provenance",
239 )
240 if not isinstance(state["datasets"], list) or not all(
241 isinstance(item, str) and 0 < len(item) <= 128 for item in state["datasets"]
242 ):
243 raise ToolError("restart manifest datasets must be bounded strings")
244 if len(state["datasets"]) > 128 or len(set(state["datasets"])) != len(state["datasets"]):
245 raise ToolError("restart manifest has too many or duplicate datasets")
246 if not isinstance(state["state_parameters_present"], bool):
247 raise ToolError("state_parameters_present must be a boolean")
248 finite_float(state["time"], name="state.time", minimum=0)
249 bounded_int(state["iteration"], name="state.iteration", minimum=0, maximum=2**63-1)
250 for key in ("oper", "time_stepping", "forcing"):
251 if key in parameters and not isinstance(parameters[key], Mapping):
252 raise ToolError(f"restart parameters.{key} must be an object")
253 digest = provenance["state_sha256"]
254 if digest is not None and not re.fullmatch(r"[0-9a-f]{64}", str(digest)):
255 raise ToolError("restart manifest state_sha256 must be lowercase SHA-256")
256 return {
257 "parameters": dict(parameters),
258 "provenance": {**provenance, "digest_origin": "manifest"},
259 "solver": document["solver"],
260 "state": dict(state),
261 }
264def _same(left: Any, right: Any) -> bool:
265 if isinstance(left, (int, float)) and isinstance(right, (int, float)):
266 try:
267 a, b = float(left), float(right)
268 except OverflowError:
269 return False
270 return math.isfinite(a) and math.isfinite(b) and math.isclose(a, b, rel_tol=1e-12, abs_tol=0.0)
271 return left == right
274def compare(source: dict[str, Any], target: dict[str, Any]) -> dict[str, Any]:
275 blockers: list[dict[str, str]] = []
276 warnings: list[dict[str, str]] = []
278 def add(target_list: list[dict[str, str]], code: str, message: str) -> None:
279 target_list.append({"code": code, "message": message})
281 if source.get("solver") is None:
282 add(blockers, "missing_solver", "source solver metadata is unavailable")
283 elif source["solver"] != target["solver"]:
284 add(blockers, "solver_mismatch", "source and target solver keys differ")
286 source_params = source.get("parameters", {})
287 target_params = target["parameters"]
288 source_oper = source_params.get("oper", {})
289 target_oper = target_params.get("oper", {})
290 for key in _COMPATIBLE_OPER_KEYS:
291 if key not in target_oper:
292 continue
293 if key not in source_oper:
294 add(blockers, "missing_grid_metadata", f"source oper.{key} is unavailable")
295 elif not _same(source_oper[key], target_oper[key]):
296 add(
297 blockers,
298 "grid_or_domain_mismatch",
299 f"oper.{key} differs; use the reviewed resolution-change workflow",
300 )
301 for key in ("coef_dealiasing", "type_fft"):
302 if key in source_oper and key in target_oper and not _same(
303 source_oper[key], target_oper[key]
304 ):
305 add(
306 warnings,
307 "operator_change",
308 f"oper.{key} changes across restart and requires justification",
309 )
311 for key in _PHYSICS_KEYS:
312 if key in source_params and key in target_params and not _same(
313 source_params[key], target_params[key]
314 ):
315 add(
316 warnings,
317 "physics_parameter_change",
318 f"{key} changes across restart; reassess equations and budgets",
319 )
321 source_state = source.get("state", {})
322 if not source_state.get("datasets"):
323 add(blockers, "empty_state", "source contains no state datasets")
324 needed = PROFILES[target["solver"]]["required_state"]
325 metadata = source_state.get("dataset_metadata", {})
326 shape = [target_oper[f"n{axis}"] for axis in ("z", "y", "x") if f"n{axis}" in target_oper]
327 for key in needed:
328 if key not in source_state.get("datasets", []):
329 add(blockers, "missing_state_variable", f"required state variable {key} is absent")
330 entry = metadata.get(key) if isinstance(metadata, Mapping) else None
331 if not isinstance(entry, Mapping) or entry.get("shape") != shape or entry.get("numeric") is not True:
332 add(blockers, "state_shape_or_dtype", f"{key} must have a verified numeric dataset with shape {shape}")
333 source_time = source_state.get("time")
334 target_end = target_params["time_stepping"].get("t_end")
335 if isinstance(source_time, (int, float)) and math.isfinite(source_time) and isinstance(target_end, (int, float)):
336 if float(target_end) <= float(source_time):
337 add(blockers, "nonadvancing_end_time", "target t_end does not exceed state time")
338 else:
339 add(blockers, "missing_time", "source time or target t_end is unavailable/nonfinite")
341 init = target_params.get("init_fields", {})
342 if init.get("type") != "from_file":
343 add(
344 blockers,
345 "target_not_from_file",
346 "target init_fields.type must be 'from_file' for this restart plan",
347 )
349 source_provenance = source.get("provenance", {})
350 source_version = source_provenance.get("fluidsim")
351 if source_version is None:
352 add(
353 warnings,
354 "missing_source_version",
355 "source FluidSim version is unavailable; do not assume migration compatibility",
356 )
357 elif source_version != FLUIDSIM_VERSION:
358 add(
359 warnings,
360 "version_migration",
361 "source version differs from 0.9.0; review release notes and merge_missing_params",
362 )
364 target_restart = target["provenance"].get("restart")
365 if not isinstance(target_restart, Mapping):
366 add(blockers, "missing_restart_provenance", "target provenance.restart is required")
367 else:
368 if target_restart.get("path") != init.get("from_file", {}).get("path"):
369 add(blockers, "restart_path_mismatch", "provenance.restart.path and init_fields.from_file.path differ")
370 expected_digest = target_restart.get("sha256")
371 observed_digest = source_provenance.get("state_sha256")
372 if observed_digest is None:
373 add(
374 blockers,
375 "digest_not_checked",
376 "source state exceeded the hash bound or lacks a manifest digest",
377 )
378 elif expected_digest != observed_digest:
379 add(blockers, "digest_mismatch", "restart SHA-256 does not match target provenance")
381 forcing = target_params.get("forcing", {})
382 if (
383 forcing.get("enable")
384 and forcing.get("type", "").startswith("tcrandom")
385 and source_state.get("forcing_state_complete") is not True
386 ):
387 add(
388 blockers,
389 "forcing_state_missing",
390 "time-correlated forcing restart lacks seed0/seed1/t_last_change in /state_params/forcing",
391 )
393 return {
394 "blockers": blockers,
395 "compatible_for_mechanical_restart": not blockers,
396 "numerical_convergence_established": False,
397 "ok": not blockers,
398 "physical_validity_established": False,
399 "source": {
400 "dataset_count": len(source_state.get("datasets", [])),
401 "fluidfft": source_provenance.get("fluidfft"),
402 "fluidsim": source_version,
403 "iteration": source_state.get("iteration"),
404 "sha256_checked": source_provenance.get("digest_origin") == "computed" and source_provenance.get("state_sha256") is not None,
405 "digest_origin": source_provenance.get("digest_origin", "manifest"),
406 "solver": source.get("solver"),
407 "state_parameters_present": source_state.get(
408 "state_parameters_present"
409 ),
410 "time": source_time,
411 },
412 "target": {
413 "fluidfft": target["provenance"]["fluidfft"],
414 "fluidsim": target["provenance"]["fluidsim"],
415 "solver": target["solver"],
416 "t_end": target_end,
417 },
418 "tool": TOOL,
419 "warnings": warnings,
420 }
423def build_parser() -> argparse.ArgumentParser:
424 parser = argparse.ArgumentParser(
425 description=(
426 "Compare local restart metadata with validated target JSON. Field arrays "
427 "are never loaded; no restart, MPI launch, or job submission occurs."
428 )
429 )
430 parser.add_argument("--source", required=True, help="State file, run directory, or manifest.")
431 parser.add_argument("--target-config", required=True)
432 parser.add_argument("--root", default=".")
433 parser.add_argument(
434 "--hash-max-gib",
435 type=int,
436 default=1,
437 help="Hash source state only up to this many GiB, 0..8 (default: 1).",
438 )
439 return parser
442def main(argv: list[str] | None = None) -> int:
443 args = build_parser().parse_args(argv)
444 try:
445 hash_gib = bounded_int(
446 args.hash_max_gib, name="hash_max_gib", minimum=0, maximum=8
447 )
448 root = checked_root(args.root)
449 source_path = checked_input(args.source, root=root, kind="any")
450 target_path = checked_input(
451 args.target_config,
452 root=root,
453 kind="file",
454 suffixes={".json"},
455 )
456 target = normalized_copy(load_json(target_path))
457 if source_path.is_file() and source_path.name.casefold().endswith(".json"):
458 source = load_manifest(source_path)
459 selected = source_path
460 else:
461 selected = _find_state_file(source_path)
462 source = load_hdf5_state(
463 selected,
464 digest_limit=hash_gib * GIB,
465 )
466 report = compare(source, target)
467 report.update(
468 {
469 "arrays_loaded": False,
470 "commands_executed": False,
471 "network_used": False,
472 "selected_source": relative_display(selected, root),
473 }
474 )
475 emit_json(report)
476 return 0 if report["ok"] else 2
477 except (OSError, ToolError, UnicodeError) as exc:
478 return fail_json(TOOL, exc)
481if __name__ == "__main__":
482 raise SystemExit(main())