drbx 2.0.0.dev0__py3-none-any.whl

drbx/runtime/output.py

@@ -0,0 +1,374 @@
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Mapping
+
+import numpy as np
+
+
+@dataclass(frozen=True)
+class RestartBundle:
+    case_name: str
+    parity_mode: str
+    component_labels: tuple[str, ...]
+    current_time: float
+    completed_steps: int
+    configured_timestep: float
+    state_variables: dict[str, np.ndarray]
+
+
+def write_restart_bundle(
+    bundle: RestartBundle,
+    path: str | Path,
+) -> Path:
+    target = Path(path)
+    target.parent.mkdir(parents=True, exist_ok=True)
+    payload: dict[str, Any] = {
+        "__metadata__": np.asarray(
+            json.dumps(
+                {
+                    "case_name": bundle.case_name,
+                    "parity_mode": bundle.parity_mode,
+                    "component_labels": list(bundle.component_labels),
+                    "current_time": bundle.current_time,
+                    "completed_steps": bundle.completed_steps,
+                    "configured_timestep": bundle.configured_timestep,
+                    "state_variables": sorted(bundle.state_variables),
+                },
+                sort_keys=True,
+            ),
+            dtype=np.str_,
+        ),
+    }
+    for name, value in bundle.state_variables.items():
+        payload[f"state:{name}"] = np.asarray(value)
+    np.savez_compressed(target, **payload)
+    return target
+
+
+def load_restart_bundle(path: str | Path) -> RestartBundle:
+    source = Path(path)
+    with np.load(source, allow_pickle=False) as payload:
+        metadata = json.loads(str(np.asarray(payload["__metadata__"]).item()))
+        state_variables = {
+            key.removeprefix("state:"): np.asarray(payload[key])
+            for key in payload.files
+            if key.startswith("state:")
+        }
+    return RestartBundle(
+        case_name=str(metadata["case_name"]),
+        parity_mode=str(metadata["parity_mode"]),
+        component_labels=tuple(str(value) for value in metadata["component_labels"]),
+        current_time=float(metadata["current_time"]),
+        completed_steps=int(metadata["completed_steps"]),
+        configured_timestep=float(metadata["configured_timestep"]),
+        state_variables=state_variables,
+    )
+
+
+def build_run_log_payload(
+    *,
+    input_file: str | Path,
+    case_name: str,
+    parity_mode: str,
+    capability_tier: str,
+    component_labels: tuple[str, ...],
+    time_points: tuple[float, ...],
+    dimensions: Mapping[str, int],
+    compare_variables: tuple[str, ...],
+    restart_supported: bool,
+    outputs: Mapping[str, str],
+    variable_summaries: Mapping[str, Any],
+    run_configuration: Mapping[str, Any] | None = None,
+    restart_info: Mapping[str, Any] | None = None,
+    events: tuple[Mapping[str, Any], ...] | None = None,
+) -> dict[str, Any]:
+    payload = {
+        "input_file": str(input_file),
+        "case_name": case_name,
+        "parity_mode": parity_mode,
+        "capability_tier": capability_tier,
+        "component_labels": list(component_labels),
+        "time_points": list(time_points),
+        "dimensions": dict(dimensions),
+        "compare_variables": list(compare_variables),
+        "restart_supported": bool(restart_supported),
+        "outputs": dict(outputs),
+        "variable_summaries": dict(variable_summaries),
+    }
+    if run_configuration is not None:
+        payload["run_configuration"] = dict(run_configuration)
+    if restart_info is not None:
+        payload["restart_info"] = dict(restart_info)
+    if events is not None:
+        rendered_events = [dict(event) for event in events]
+        payload["events"] = rendered_events
+        payload["event_count"] = len(rendered_events)
+        payload["event_stages"] = [str(event.get("stage", "")) for event in rendered_events]
+    return payload
+
+
+def write_run_log_payload(payload: Mapping[str, Any], path: str | Path) -> Path:
+    target = Path(path)
+    target.parent.mkdir(parents=True, exist_ok=True)
+    target.write_text(json.dumps(dict(payload), indent=2, sort_keys=True), encoding="utf-8")
+    return target
+
+
+def format_run_log_text(payload: Mapping[str, Any]) -> str:
+    outputs = payload.get("outputs", {})
+    output_lines = "\n".join(f"  - {name}: {value}" for name, value in outputs.items()) if outputs else "  - (none)"
+    time_points = payload.get("time_points", [])
+    if time_points:
+        time_line = f"{time_points[0]} -> {time_points[-1]} ({len(time_points)} stored states)"
+    else:
+        time_line = "(none)"
+    components = ", ".join(payload.get("component_labels", [])) or "(none)"
+    compare_variables = ", ".join(payload.get("compare_variables", [])) or "(none)"
+    run_configuration = payload.get("run_configuration", {})
+    restart_info = payload.get("restart_info", {})
+    mesh = run_configuration.get("mesh", {})
+    solver = run_configuration.get("solver", {})
+    time_cfg = run_configuration.get("time", {})
+    runtime_cfg = run_configuration.get("runtime", {})
+    runtime_cfg = run_configuration.get("runtime", {})
+    restart_lines = []
+    if restart_info:
+        for key in ("loaded_from", "start_time", "input_completed_steps", "requested_additional_steps", "saved_completed_steps"):
+            if key in restart_info:
+                restart_lines.append(f"  - {key}: {restart_info[key]}")
+    restart_block = "\n".join(restart_lines) if restart_lines else "  - (fresh run)"
+    return (
+        f"Run Summary\n"
+        f"  input: {payload.get('input_file')}\n"
+        f"  case: {payload.get('case_name')}\n"
+        f"  mode: {payload.get('parity_mode')}\n"
+        f"  capability tier: {payload.get('capability_tier', '(unknown)')}\n"
+        f"  precision: {runtime_cfg.get('precision', '(default)')}\n"
+        f"  configured nout/timestep: {time_cfg.get('nout', '(unknown)')} / {time_cfg.get('timestep', '(unknown)')}\n"
+        f"  runtime: precision={runtime_cfg.get('precision', '(default)')}, backend={runtime_cfg.get('backend', '(unknown)')}, device={runtime_cfg.get('device', '(unknown)')}, elapsed={runtime_cfg.get('elapsed_seconds', '(unknown)')}\n"
+        f"  mesh: nx={mesh.get('nx', '(unknown)')}, ny={mesh.get('ny', '(unknown)')}, nz={mesh.get('nz', '(unknown)')}, file={mesh.get('file', '<analytic mesh>')}\n"
+        f"  solver: type={solver.get('type', '<native default>')}, mxstep={solver.get('mxstep', '(unknown)')}, rtol={solver.get('rtol', '(unknown)')}, atol={solver.get('atol', '(unknown)')}\n"
+        f"  components: {components}\n"
+        f"  compare variables: {compare_variables}\n"
+        f"  dimensions: {payload.get('dimensions')}\n"
+        f"  time: {time_line}\n"
+        f"  events: {payload.get('event_count', 0)}\n"
+        f"  restart supported: {'yes' if payload.get('restart_supported') else 'no'}\n"
+        f"  restart:\n{restart_block}\n"
+        f"  outputs:\n{output_lines}"
+    )
+
+
+def print_run_log(payload: Mapping[str, Any], *, verbosity: str = "summary") -> None:
+    try:
+        from rich.console import Console
+        from rich.panel import Panel
+        from rich.table import Table
+    except Exception:
+        print(format_run_log_text(payload))
+        return
+
+    console = Console()
+    run_configuration = payload.get("run_configuration", {})
+    restart_info = payload.get("restart_info", {})
+    mesh = run_configuration.get("mesh", {})
+    solver = run_configuration.get("solver", {})
+    time_cfg = run_configuration.get("time", {})
+    runtime_cfg = run_configuration.get("runtime", {})
+
+    summary = Table.grid(padding=(0, 2))
+    summary.add_column(style="bold cyan")
+    summary.add_column()
+    summary.add_row("input", str(payload.get("input_file")))
+    summary.add_row("case", str(payload.get("case_name")))
+    summary.add_row("mode", str(payload.get("parity_mode")))
+    summary.add_row("tier", str(payload.get("capability_tier", "(unknown)")))
+    summary.add_row("precision", str(runtime_cfg.get("precision", "(default)")))
+    summary.add_row("backend", str(runtime_cfg.get("backend", "(unknown)")))
+    summary.add_row("device", str(runtime_cfg.get("device", "(unknown)")))
+    summary.add_row(
+        "elapsed",
+        "(unknown)" if runtime_cfg.get("elapsed_seconds") is None else f"{float(runtime_cfg['elapsed_seconds']):.3f} s",
+    )
+    summary.add_row("cache", str(runtime_cfg.get("compilation_cache_dir", "(none)")))
+    summary.add_row("time", f"{time_cfg.get('nout', '(unknown)')} outputs, dt={time_cfg.get('timestep', '(unknown)')}")
+    summary.add_row(
+        "mesh",
+        f"nx={mesh.get('nx', '(unknown)')}, ny={mesh.get('ny', '(unknown)')}, nz={mesh.get('nz', '(unknown)')}, file={mesh.get('file', '<analytic mesh>')}",
+    )
+    summary.add_row(
+        "solver",
+        f"type={solver.get('type', '<native default>')}, mxstep={solver.get('mxstep', '(unknown)')}, rtol={solver.get('rtol', '(unknown)')}, atol={solver.get('atol', '(unknown)')}",
+    )
+    summary.add_row("components", ", ".join(payload.get("component_labels", [])) or "(none)")
+    summary.add_row("compare vars", ", ".join(payload.get("compare_variables", [])) or "(none)")
+    summary.add_row("dimensions", json.dumps(payload.get("dimensions", {}), sort_keys=True))
+    summary.add_row("events", str(payload.get("event_count", 0)))
+    summary.add_row("restart", "yes" if payload.get("restart_supported") else "no")
+
+    outputs = Table(title="Outputs", show_header=True, header_style="bold magenta")
+    outputs.add_column("artifact")
+    outputs.add_column("path")
+    for name, value in payload.get("outputs", {}).items():
+        outputs.add_row(str(name), str(value))
+
+    restart_table = Table(title="Restart Provenance", show_header=True, header_style="bold magenta")
+    restart_table.add_column("field")
+    restart_table.add_column("value")
+    if restart_info:
+        for key in ("loaded_from", "start_time", "input_completed_steps", "requested_additional_steps", "saved_completed_steps"):
+            if key in restart_info:
+                restart_table.add_row(str(key), str(restart_info[key]))
+    else:
+        restart_table.add_row("state", "fresh run")
+
+    variable_table = Table(title="Variable Summaries", show_header=True, header_style="bold magenta")
+    variable_table.add_column("variable")
+    variable_table.add_column("min", justify="right")
+    variable_table.add_column("max", justify="right")
+    variable_table.add_column("mean", justify="right")
+    variable_table.add_column("delta", justify="right")
+    for name, summary_payload in payload.get("variable_summaries", {}).items():
+        delta = summary_payload.get("max_abs_delta_last_first")
+        variable_table.add_row(
+            str(name),
+            f"{float(summary_payload.get('minimum', 0.0)):.6e}",
+            f"{float(summary_payload.get('maximum', 0.0)):.6e}",
+            f"{float(summary_payload.get('mean', 0.0)):.6e}",
+            "(n/a)" if delta is None else f"{float(delta):.6e}",
+        )
+
+    console.print(Panel(summary, title="Run Summary", border_style="cyan"))
+    console.print(outputs)
+    if verbosity == "detailed":
+        console.print(restart_table)
+        console.print(variable_table)
+
+
+def build_run_event(
+    *,
+    stage: str,
+    message: str,
+    elapsed_seconds: float | None = None,
+    details: Mapping[str, Any] | None = None,
+) -> dict[str, Any]:
+    event: dict[str, Any] = {
+        "stage": stage,
+        "message": message,
+    }
+    if elapsed_seconds is not None:
+        event["elapsed_seconds"] = elapsed_seconds
+    if details:
+        event["details"] = {str(key): _jsonify_event_value(value) for key, value in details.items()}
+    return event
+
+
+def print_run_event(event: Mapping[str, Any], *, verbosity: str = "summary") -> None:
+    try:
+        from rich.console import Console
+        from rich.panel import Panel
+        from rich.table import Table
+    except Exception:
+        elapsed = event.get("elapsed_seconds")
+        prefix = f"[{float(elapsed):8.3f}s] " if elapsed is not None else ""
+        suffix = _format_event_summary_details(event)
+        print(f"{prefix}{event.get('stage', 'event')}: {event.get('message', '')}{suffix}")
+        details = event.get("details", {})
+        if verbosity == "detailed" and isinstance(details, Mapping):
+            for key, value in details.items():
+                print(f"  - {key}: {value}")
+        return
+
+    console = Console()
+    if verbosity != "detailed":
+        elapsed = event.get("elapsed_seconds")
+        prefix = f"[{float(elapsed):8.3f}s] " if elapsed is not None else ""
+        suffix = _format_event_summary_details(event)
+        console.print(f"{prefix}[bold blue]{event.get('stage', 'event')}[/bold blue]: {event.get('message', '')}{suffix}")
+        return
+
+    details = event.get("details", {})
+    table = Table.grid(padding=(0, 2))
+    table.add_column(style="bold cyan")
+    table.add_column()
+    table.add_row("stage", str(event.get("stage", "")))
+    table.add_row("message", str(event.get("message", "")))
+    if event.get("elapsed_seconds") is not None:
+        table.add_row("elapsed", f"{float(event['elapsed_seconds']):.3f} s")
+    if isinstance(details, Mapping):
+        for key, value in details.items():
+            table.add_row(str(key), str(value))
+    console.print(Panel(table, border_style="blue"))
+
+
+def _format_event_summary_details(event: Mapping[str, Any]) -> str:
+    details = event.get("details")
+    if not isinstance(details, Mapping):
+        return ""
+    stage = str(event.get("stage", ""))
+    if stage != "progress":
+        return ""
+    segments: list[str] = []
+    interval_index = details.get("interval_index")
+    steps = details.get("steps")
+    if interval_index is not None and steps is not None:
+        segments.append(f"interval {_format_count(interval_index)}/{_format_count(steps)}")
+    fraction_complete = details.get("fraction_complete")
+    if fraction_complete is not None:
+        try:
+            segments.append(f"{100.0 * float(fraction_complete):.1f}%")
+        except (TypeError, ValueError):
+            pass
+    solver_mode = details.get("solver_mode")
+    if solver_mode is not None:
+        segments.append(f"mode {solver_mode}")
+    accepted_dt = details.get("accepted_dt")
+    if accepted_dt is not None:
+        try:
+            segments.append(f"dt {float(accepted_dt):.3g}")
+        except (TypeError, ValueError):
+            pass
+    if bool(details.get("live_progress", True)):
+        eta = details.get("estimated_remaining_seconds")
+        if eta is not None:
+            try:
+                segments.append(f"eta {_format_elapsed_seconds(float(eta))}")
+            except (TypeError, ValueError):
+                pass
+    if not segments:
+        return ""
+    return f" ({', '.join(segments)})"
+
+
+def _format_count(value: Any) -> str:
+    try:
+        numeric = float(value)
+    except (TypeError, ValueError):
+        return str(value)
+    if numeric.is_integer():
+        return str(int(numeric))
+    return f"{numeric:.3g}"
+
+
+def _format_elapsed_seconds(value: float) -> str:
+    seconds = max(float(value), 0.0)
+    if seconds < 60.0:
+        return f"{seconds:.1f}s"
+    minutes, rem = divmod(seconds, 60.0)
+    if minutes < 60.0:
+        return f"{int(minutes)}m {rem:04.1f}s"
+    hours, minutes = divmod(minutes, 60.0)
+    return f"{int(hours)}h {int(minutes)}m"
+
+
+def _jsonify_event_value(value: Any) -> Any:
+    if isinstance(value, Path):
+        return str(value)
+    if isinstance(value, Mapping):
+        return {str(key): _jsonify_event_value(item) for key, item in value.items()}
+    if isinstance(value, (list, tuple)):
+        return [_jsonify_event_value(item) for item in value]
+    return value

drbx/runtime/paths.py

@@ -0,0 +1,9 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+
+def repo_root() -> Path:
+    """Repository root (the directory containing ``src/``)."""
+
+    return Path(__file__).resolve().parents[3]

drbx/runtime/performance.py

@@ -0,0 +1,161 @@
+from __future__ import annotations
+
+import os
+import platform
+import shlex
+import sys
+from pathlib import Path
+from typing import Any
+
+from ..config.boutinp import BoutConfig
+
+_VALID_PRECISIONS = {"float32", "float64"}
+_HOST_DEVICE_FLAG_PREFIX = "--xla_force_host_platform_device_count="
+
+
+def resolve_runtime_precision(
+    *,
+    requested: str | None = None,
+    config: BoutConfig | None = None,
+) -> str:
+    candidate = requested
+    if candidate is None and config is not None and config.has_option("runtime", "precision"):
+        parsed = config.parsed("runtime", "precision")
+        candidate = str(parsed)
+    if candidate is None:
+        candidate = os.environ.get("DRBX_PRECISION", "float64")
+    normalized = str(candidate).strip().lower()
+    if normalized not in _VALID_PRECISIONS:
+        raise ValueError(f"Unsupported precision {candidate!r}; expected one of {sorted(_VALID_PRECISIONS)}")
+    return normalized
+
+
+def resolve_host_device_count(*, requested: int | str | None = None) -> int | None:
+    candidate = requested
+    if candidate is None:
+        candidate = os.environ.get("DRBX_HOST_DEVICE_COUNT")
+    if candidate is None:
+        return None
+    normalized = str(candidate).strip()
+    if not normalized or normalized == "0":
+        return None
+    try:
+        value = int(normalized)
+    except ValueError as exc:
+        raise ValueError(f"Unsupported host device count {candidate!r}; expected a positive integer") from exc
+    if value <= 0:
+        raise ValueError(f"Unsupported host device count {candidate!r}; expected a positive integer")
+    return value
+
+
+def configure_jax_runtime(
+    *,
+    precision: str | None = None,
+    host_device_count: int | str | None = None,
+) -> Path | None:
+    resolved_precision = resolve_runtime_precision(requested=precision)
+    os.environ["DRBX_PRECISION"] = resolved_precision
+    resolved_host_device_count = resolve_host_device_count(requested=host_device_count)
+    if resolved_host_device_count is not None:
+        os.environ["DRBX_HOST_DEVICE_COUNT"] = str(resolved_host_device_count)
+        _configure_host_device_count_xla_flags(resolved_host_device_count)
+    if os.environ.get("DRBX_DISABLE_COMPILATION_CACHE", "").strip().lower() in {"1", "true", "yes", "on"}:
+        cache_dir = None
+    else:
+        cache_dir = _compilation_cache_dir()
+        cache_dir.mkdir(parents=True, exist_ok=True)
+
+    import jax
+    jax.config.update("jax_enable_x64", resolved_precision == "float64")
+    if cache_dir is not None:
+        from jax.experimental.compilation_cache import compilation_cache as compilation_cache
+
+        jax.config.update("jax_enable_compilation_cache", True)
+        jax.config.update("jax_compilation_cache_dir", str(cache_dir))
+        if "DRBX_PERSISTENT_CACHE_MIN_COMPILE_TIME_SECS" in os.environ:
+            min_compile_time = float(os.environ["DRBX_PERSISTENT_CACHE_MIN_COMPILE_TIME_SECS"])
+        else:
+            min_compile_time = 0.0
+        if "DRBX_PERSISTENT_CACHE_MIN_ENTRY_SIZE_BYTES" in os.environ:
+            min_entry_size = int(os.environ["DRBX_PERSISTENT_CACHE_MIN_ENTRY_SIZE_BYTES"])
+        else:
+            min_entry_size = 0
+        jax.config.update("jax_persistent_cache_min_compile_time_secs", min_compile_time)
+        jax.config.update("jax_persistent_cache_min_entry_size_bytes", min_entry_size)
+        compilation_cache.set_cache_dir(str(cache_dir))
+    return cache_dir
+
+
+def runtime_numpy_dtype(*, precision: str | None = None) -> Any:
+    import numpy as np
+
+    resolved = resolve_runtime_precision(requested=precision)
+    return np.float32 if resolved == "float32" else np.float64
+
+
+def runtime_jax_dtype(*, precision: str | None = None) -> Any:
+    import jax.numpy as jnp
+
+    resolved = resolve_runtime_precision(requested=precision)
+    return jnp.float32 if resolved == "float32" else jnp.float64
+
+
+def runtime_parallel_summary() -> dict[str, Any]:
+    import jax
+
+    requested_host_devices = resolve_host_device_count()
+    xla_flags = os.environ.get("XLA_FLAGS", "")
+    configured_host_devices = _extract_host_device_count_from_flags(xla_flags)
+    return {
+        "backend": jax.default_backend(),
+        "cpu_count": os.cpu_count(),
+        "device_count": jax.device_count(),
+        "local_device_count": jax.local_device_count(),
+        "devices": [str(device) for device in jax.devices()],
+        "requested_host_device_count": requested_host_devices,
+        "configured_host_device_count": configured_host_devices,
+        "explicit_host_device_parallelism_enabled": (
+            jax.default_backend() == "cpu" and jax.local_device_count() > 1
+        ),
+        "xla_flags": xla_flags,
+    }
+
+
+def _compilation_cache_dir() -> Path:
+    override = os.environ.get("DRBX_CACHE_DIR")
+    if override:
+        return Path(override).expanduser()
+    return _default_user_cache_root() / "drbx" / "jax_compilation_cache"
+
+
+def _default_user_cache_root() -> Path:
+    xdg = os.environ.get("XDG_CACHE_HOME")
+    if xdg:
+        return Path(xdg).expanduser()
+    if platform.system() == "Darwin":
+        return Path.home() / "Library" / "Caches"
+    return Path.home() / ".cache"
+
+
+def _configure_host_device_count_xla_flags(host_device_count: int) -> None:
+    existing_flags = os.environ.get("XLA_FLAGS", "")
+    configured_count = _extract_host_device_count_from_flags(existing_flags)
+    if configured_count == host_device_count:
+        return
+    if "jax" in sys.modules:
+        raise RuntimeError(
+            "DRBX_HOST_DEVICE_COUNT must be set before importing jax/drbx so CPU devices can be configured."
+        )
+    tokens = [token for token in shlex.split(existing_flags) if not token.startswith(_HOST_DEVICE_FLAG_PREFIX)]
+    tokens.append(f"{_HOST_DEVICE_FLAG_PREFIX}{host_device_count}")
+    os.environ["XLA_FLAGS"] = " ".join(tokens)
+
+
+def _extract_host_device_count_from_flags(flags: str) -> int | None:
+    for token in shlex.split(flags):
+        if token.startswith(_HOST_DEVICE_FLAG_PREFIX):
+            try:
+                return int(token.split("=", 1)[1])
+            except ValueError:
+                return None
+    return None