stormlog.jax.diagnose

Diagnostic bundle builder for the JAX Stormlog diagnose command.

Functions

build_diagnostic_summary([device_index])	Build diagnostic summary and risk flags from current state.
collect_environment([device_index])	Collect system, JAX backend, and device environment details.
run_diagnose(output, device_index, duration, ...)	Build the full diagnostic bundle and write all artifact files.
run_timeline_capture(device_index, ...)	Capture a timeline of memory metrics by running the tracker briefly.

stormlog.jax.diagnose.collect_environment(device_index=0)

Collect system, JAX backend, and device environment details.

Parameters:

device_index (int)

Return type:

Dict[str, Any]

stormlog.jax.diagnose.run_timeline_capture(device_index, duration_seconds, interval)

Capture a timeline of memory metrics by running the tracker briefly.

Returns timeline data in the shared Stormlog shape: timestamps, allocated (bytes), reserved (bytes).

Parameters:

• device_index (int)

• duration_seconds (float)

• interval (float)

Return type:

Dict[str, Any]

stormlog.jax.diagnose.build_diagnostic_summary(device_index=0)

Build diagnostic summary and risk flags from current state.

Returns (summary_dict, risk_detected). Summary schema matches the TensorFlow backend for downstream compatibility.

Parameters:

device_index (int)

Return type:

Tuple[Dict[str, Any], bool]

stormlog.jax.diagnose.run_diagnose(output, device_index, duration, interval, command_line)

Build the full diagnostic bundle and write all artifact files.

Returns (artifact_dir, exit_code). exit_code: 0 = success no risk, 1 = failure, 2 = success with memory risk.

Parameters:

• output (str | None)

• device_index (int)

• duration (float)

• interval (float)

• command_line (str)

Return type:

Tuple[Path, int]