"""Reusable ADIOS FileReader helpers for finite BP and ACA reads."""
from __future__ import annotations
from pathlib import Path
import logging
import os
import re
import threading
import time
from typing import Mapping, Sequence
import adios2
import numpy as np
_LOG_LEVEL_NAME = os.environ.get("XGC_INPUT_GUI_MONITOR_LOG_LEVEL", "INFO").upper()
_LOG_LEVEL = getattr(logging, _LOG_LEVEL_NAME, logging.INFO)
if not logging.getLogger().handlers:
logging.basicConfig(level=_LOG_LEVEL)
LOGGER = logging.getLogger(__name__)
LOGGER.setLevel(_LOG_LEVEL)
[docs]
def read_adios_file_steps(
source_path: Path,
variables: Sequence[str],
adios_steps: Sequence[int],
*,
source_id: str | None = None,
file_reader=None,
qualify_with_source_id: bool = False,
) -> Mapping[int, Mapping[str, object]]:
"""
Read selected variables and ADIOS steps through ``adios2.FileReader``.
Parameters
----------
source_path : pathlib.Path
ADIOS-readable BP source. This path is opened when ``file_reader`` is
omitted.
variables : sequence[str]
Reader-facing variable names to read at each selected step.
adios_steps : sequence[int]
Physical ADIOS step ids to read.
source_id : str or None, optional
Source id used to qualify campaign variables, for example
``xgc.f3d.00010.bp/e_den``.
file_reader : adios2.FileReader or None, optional
Existing open FileReader handle. Supplying this avoids reopening
expensive campaign files for every read batch.
qualify_with_source_id : bool, optional
If True, qualify unqualified variable names as ``source_id/variable``.
Returns
-------
dict[int, dict[str, object]]
Values keyed by ADIOS step and original reader-facing variable name.
"""
requested_steps = sorted({int(step) for step in adios_steps})
if not requested_steps:
return {}
LOGGER.debug(
"ADIOS read_adios_file_steps start: source_path=%s source_id=%s variables=%s requested_steps=%s persistent_reader=%s",
source_path,
source_id,
list(variables),
list(requested_steps),
file_reader is not None,
)
started = time.perf_counter()
if file_reader is not None:
result = _read_steps_from_open_file_reader(
file_reader,
variables,
requested_steps,
source_id=source_id,
qualify_with_source_id=qualify_with_source_id,
)
LOGGER.debug(
"ADIOS read_adios_file_steps done: source_path=%s source_id=%s steps=%s elapsed=%.3fs persistent_reader=%s",
source_path,
source_id,
len(requested_steps),
time.perf_counter() - started,
True,
)
return result
with adios2.FileReader(str(Path(source_path))) as reader:
result = _read_steps_from_open_file_reader(
reader,
variables,
requested_steps,
source_id=source_id,
qualify_with_source_id=qualify_with_source_id,
)
LOGGER.debug(
"ADIOS read_adios_file_steps done: source_path=%s source_id=%s steps=%s elapsed=%.3fs persistent_reader=%s",
source_path,
source_id,
len(requested_steps),
time.perf_counter() - started,
False,
)
return result
[docs]
def available_step_count(reader) -> int:
"""
Return the maximum ``AvailableStepsCount`` advertised by an open FileReader.
Parameters
----------
reader : adios2.FileReader
Open FileReader handle.
"""
vars_info = reader.available_variables() or {}
return max([available_steps_count_for_variable(info) for info in vars_info.values()] or [0])
[docs]
def available_steps_count_for_variable(info) -> int:
"""
Return ``AvailableStepsCount`` from one ADIOS variable metadata record.
Parameters
----------
info : dict
Raw metadata for one ADIOS variable.
"""
try:
return int(info.get("AvailableStepsCount", 0) or 0)
except Exception:
return 0
[docs]
def available_variable_names(reader) -> list[str]:
"""
Return sorted variable names advertised by an open FileReader.
Parameters
----------
reader : adios2.FileReader
Open FileReader handle.
"""
return sorted((reader.available_variables() or {}).keys())
[docs]
class AdiosFileSourceReader:
"""
Callable source-reader backend for finite catalog reads.
The same callable shape is used for directory-backed local BP products and
campaign-backed ACA products. Directory catalogs construct this class
without an open reader, so each physical BP source is opened once per
batched read request. Campaign catalogs pass a persistent FileReader handle
and enable source-id qualification, avoiding repeated ACA opens.
"""
def __init__(
self,
file_reader=None,
*,
qualify_with_source_id: bool = False,
owns_reader: bool = False,
reopen_on_unopened_file: bool = False,
on_reopen=None,
):
"""
Store FileReader backend configuration.
Parameters
----------
file_reader : adios2.FileReader or None, optional
Persistent reader handle. If omitted, ``source_path`` is opened and
closed on each call.
qualify_with_source_id : bool, optional
If True, unqualified variable names are read as
``source_id/variable``.
owns_reader : bool, optional
Whether :meth:`close` should close ``file_reader``.
reopen_on_unopened_file : bool, optional
If True, an ADIOS remote ``FileNotOpen`` read failure closes and
reopens the persistent reader once before retrying. This is useful
for campaign-backed remote reads where a remote file handle can
become invalid while the ACA metadata handle remains usable.
on_reopen : callable or None, optional
Callback invoked with the new FileReader when a persistent handle is
reopened.
"""
self.file_reader = file_reader
self.qualify_with_source_id = bool(qualify_with_source_id)
self.owns_reader = bool(owns_reader)
self.reopen_on_unopened_file = bool(reopen_on_unopened_file)
self.on_reopen = on_reopen
self._lock = threading.RLock()
def __call__(
self,
source_path: Path,
variables: Sequence[str],
adios_steps: Sequence[int],
*,
source_id: str | None = None,
) -> Mapping[int, Mapping[str, object]]:
"""
Read one catalog source batch.
Parameters
----------
source_path : pathlib.Path
BP source path, or the campaign path when ``file_reader`` is
persistent.
variables : sequence[str]
Reader-facing variable names.
adios_steps : sequence[int]
Physical ADIOS step ids.
source_id : str or None, optional
Catalog source id used for campaign qualification.
"""
with self._lock:
LOGGER.debug(
"AdiosFileSourceReader call start: source_path=%s source_id=%s variables=%s adios_steps=%s persistent_reader=%s",
source_path,
source_id,
list(variables),
list(adios_steps),
self.file_reader is not None,
)
started = time.perf_counter()
try:
result = self._read_locked(source_path, variables, adios_steps, source_id=source_id)
LOGGER.debug(
"AdiosFileSourceReader call done: source_path=%s source_id=%s elapsed=%.3fs",
source_path,
source_id,
time.perf_counter() - started,
)
return result
except RuntimeError as exc:
if not self._should_retry_with_reopened_reader(exc):
LOGGER.debug(
"AdiosFileSourceReader call failed: source_path=%s source_id=%s elapsed=%.3fs error=%s",
source_path,
source_id,
time.perf_counter() - started,
exc,
)
raise
LOGGER.debug(
"AdiosFileSourceReader reopening persistent reader: source_path=%s source_id=%s error=%s",
source_path,
source_id,
exc,
)
self._reopen_locked(source_path)
result = self._read_locked(source_path, variables, adios_steps, source_id=source_id)
LOGGER.debug(
"AdiosFileSourceReader call done after reopen: source_path=%s source_id=%s elapsed=%.3fs",
source_path,
source_id,
time.perf_counter() - started,
)
return result
[docs]
def close(self) -> None:
"""Close the persistent FileReader handle when this reader owns it."""
if self.file_reader is None or not self.owns_reader:
return
with self._lock:
close = getattr(self.file_reader, "close", None)
if close is not None:
close()
self.file_reader = None
[docs]
@staticmethod
def qualified_variable_name(source_id: str | None, variable: str) -> str:
"""
Return ``source_id/variable`` unless the name is already qualified.
Parameters
----------
source_id : str or None
Catalog source id.
variable : str
Reader-facing variable name.
"""
if source_id is None or "/" in variable:
return variable
return f"{source_id}/{variable}"
def _read_locked(
self,
source_path: Path,
variables: Sequence[str],
adios_steps: Sequence[int],
*,
source_id: str | None,
) -> Mapping[int, Mapping[str, object]]:
"""
Read using the current backend state while ``self._lock`` is held.
Parameters are the same as :meth:`__call__`.
"""
return read_adios_file_steps(
source_path,
variables,
adios_steps,
source_id=source_id,
file_reader=self.file_reader,
qualify_with_source_id=self.qualify_with_source_id,
)
def _should_retry_with_reopened_reader(self, exc: RuntimeError) -> bool:
"""Return True for the ADIOS remote unopened-file failure mode."""
if self.file_reader is None or not self.reopen_on_unopened_file:
return False
message = _strip_ansi_codes(str(exc))
return "FileNotOpen" in message or "unopened file" in message
def _reopen_locked(self, source_path: Path) -> None:
"""
Replace the persistent FileReader while ``self._lock`` is held.
Parameters
----------
source_path : pathlib.Path
ACA/BP path used to create the replacement FileReader.
"""
old_reader = self.file_reader
close = getattr(old_reader, "close", None)
if close is not None:
try:
close()
except Exception:
pass
self.file_reader = adios2.FileReader(str(Path(source_path)))
if callable(self.on_reopen):
self.on_reopen(self.file_reader)
def _read_steps_from_open_file_reader(
reader,
variables: Sequence[str],
requested_steps: Sequence[int],
*,
source_id: str | None,
qualify_with_source_id: bool,
) -> Mapping[int, Mapping[str, object]]:
"""
Read selected steps from an already open FileReader.
Parameters
----------
reader : adios2.FileReader
Open FileReader handle.
variables : sequence[str]
Reader-facing variable names.
requested_steps : sequence[int]
ADIOS step ids to read.
source_id : str or None
Catalog source id used when qualification is enabled.
qualify_with_source_id : bool
Whether to qualify unqualified variables with ``source_id``.
"""
available_info = reader.available_variables() or {}
fallback_step_count = available_step_count(reader)
out = {int(step): {} for step in requested_steps}
# Phase 1: queue one deferred read per variable and contiguous step range.
# For remote sources every synchronous read() is a full network round
# trip; deferring lets the engine pipeline (or batch) all requests in a
# single read_complete() flush.
pending: list[dict] = []
for variable in variables:
read_name = (
AdiosFileSourceReader.qualified_variable_name(source_id, variable)
if qualify_with_source_id
else variable
)
_validate_read_request(
available_info,
read_name,
variable,
requested_steps,
fallback_step_count,
)
info = available_info[read_name]
for step_range in _contiguous_step_ranges(requested_steps):
range_start = int(step_range[0])
range_count = len(step_range)
LOGGER.debug(
"ADIOS FileReader.read defer: read_name=%s requested_name=%s source_id=%s step_start=%s step_count=%s",
read_name,
variable,
source_id,
range_start,
range_count,
)
try:
value = reader.read(
read_name,
step_selection=[range_start, range_count],
defer_read=True,
)
except RuntimeError as exc:
LOGGER.debug(
"ADIOS FileReader.read defer failed: read_name=%s requested_name=%s source_id=%s step_start=%s step_count=%s error=%s",
read_name,
variable,
source_id,
range_start,
range_count,
exc,
)
raise _augment_adios_read_error(
exc,
read_name=read_name,
requested_name=variable,
adios_step=range_start,
step_count=range_count,
source_id=source_id,
) from exc
pending.append(
{
"variable": variable,
"read_name": read_name,
"step_range": step_range,
"info": info,
"value": value,
}
)
# Phase 2: one flush performs all queued reads.
if pending:
flush_started = time.perf_counter()
try:
reader.read_complete()
except RuntimeError as exc:
names = sorted({str(entry["read_name"]) for entry in pending})
LOGGER.debug(
"ADIOS read_complete failed: source_id=%s reads=%s elapsed=%.3fs error=%s",
source_id,
len(pending),
time.perf_counter() - flush_started,
exc,
)
raise RuntimeError(
f"ADIOS deferred read flush failed for {len(pending)} reads "
f"(source_id={source_id!r}, variables={names}): "
f"{_strip_ansi_codes(str(exc))}"
) from exc
LOGGER.debug(
"ADIOS read_complete done: source_id=%s reads=%s elapsed=%.3fs",
source_id,
len(pending),
time.perf_counter() - flush_started,
)
# Phase 3: split the now-filled buffers into per-step values.
for entry in pending:
step_values = _split_step_range_values(entry["value"], entry["step_range"], entry["info"])
for adios_step, step_value in step_values.items():
out[int(adios_step)][entry["variable"]] = step_value
return out
def _contiguous_step_ranges(requested_steps: Sequence[int]) -> list[list[int]]:
"""
Return contiguous ADIOS step runs from a sorted step list.
Parameters
----------
requested_steps : sequence[int]
Sorted ADIOS step ids requested from one source.
"""
ranges: list[list[int]] = []
current: list[int] = []
for raw_step in requested_steps:
step = int(raw_step)
if not current or step == current[-1] + 1:
current.append(step)
continue
ranges.append(current)
current = [step]
if current:
ranges.append(current)
return ranges
def _split_step_range_values(
value,
requested_steps: Sequence[int],
info,
) -> dict[int, object]:
"""
Split one multi-step ADIOS read into per-step values.
Parameters
----------
value : object
Value returned by ``FileReader.read`` for ``step_selection=[start,count]``.
requested_steps : sequence[int]
Contiguous ADIOS steps represented by ``value``.
info : dict
ADIOS metadata record for the variable.
"""
if len(requested_steps) == 1:
only_step = int(requested_steps[0])
return {only_step: _normalize_scalar_value(value, info)}
arr = np.asarray(value)
payload_dims = _payload_shape_dims(info)
if _is_single_value(info):
flat = arr.reshape(len(requested_steps))
return {int(step): flat[i].item() for i, step in enumerate(requested_steps)}
values_per_step = int(np.prod(payload_dims)) if payload_dims else 1
expected_size = len(requested_steps) * values_per_step
if arr.size != expected_size:
raise ValueError(
f"ADIOS range read returned size={arr.size} for Shape={info.get('Shape', '')!r}; "
f"expected {expected_size} values across {len(requested_steps)} steps."
)
reshaped = arr.reshape((len(requested_steps), *payload_dims)) if payload_dims else arr.reshape(len(requested_steps))
return {int(step): np.asarray(reshaped[i]) for i, step in enumerate(requested_steps)}
def _augment_adios_read_error(
exc: RuntimeError,
*,
read_name: str,
requested_name: str,
adios_step: int,
step_count: int = 1,
source_id: str | None,
) -> RuntimeError:
"""
Return a contextual read error for ADIOS/HPC-Campaign failures.
Parameters
----------
exc : RuntimeError
Original ADIOS exception.
read_name : str
Fully qualified ADIOS variable name.
requested_name : str
Reader-facing variable name.
adios_step : int
Physical ADIOS step requested.
source_id : str or None
Catalog source id, when available.
"""
message = _strip_ansi_codes(str(exc))
if int(step_count) > 1:
adios_step_text = f"{int(adios_step)}..{int(adios_step) + int(step_count) - 1}"
else:
adios_step_text = str(int(adios_step))
details = [
f"ADIOS read failed for variable '{requested_name}' as '{read_name}'",
f"source_id={source_id!r}, adios_step={adios_step_text}",
message,
]
cache_path = _missing_campaign_cache_path(message)
if cache_path is not None:
details.append(_campaign_cache_diagnostic(cache_path))
return RuntimeError("\n".join(detail for detail in details if detail))
_ANSI_ESCAPE_RE = re.compile(r"\x1b\[[0-9;]*m")
_MISSING_CACHE_RE = re.compile(r"couldn't open file\s+(.+?), in call to POSIX open: errno = 2")
def _strip_ansi_codes(text: str) -> str:
"""Return ``text`` without terminal ANSI color escape sequences."""
return _ANSI_ESCAPE_RE.sub("", str(text))
def _missing_campaign_cache_path(message: str) -> Path | None:
"""Return the missing cache path from an ADIOS FilePOSIX error."""
match = _MISSING_CACHE_RE.search(message)
if match is None:
return None
path = Path(match.group(1))
if "xgc-campaign-cache" not in str(path):
return None
return path
def _campaign_cache_diagnostic(cache_path: Path) -> str:
"""
Return extra context for a missing HPC-Campaign payload cache file.
Parameters
----------
cache_path : pathlib.Path
Missing file reported by ADIOS, usually ``.../data.0``.
"""
cache_dir = cache_path.parent
info_path = cache_dir / "info.txt"
lines = [
"HPC-Campaign payload cache file is missing.",
f"missing_cache_file={cache_path}",
f"cache_dir={cache_dir}",
]
if info_path.exists():
try:
lines.append("cache_info:")
lines.extend(f" {line}" for line in info_path.read_text(encoding="utf-8", errors="replace").splitlines())
except OSError:
lines.append(f"cache_info={info_path} (could not read)")
lines.append(
"This usually means campaign metadata was cached but the payload was not fetched; "
"clear or remove this cache entry and retry with the hpc_campaign connector running."
)
return "\n".join(lines)
def _validate_read_request(
available_info,
read_name: str,
requested_name: str,
requested_steps: Sequence[int],
fallback_step_count: int,
) -> None:
"""
Validate one variable read before calling ADIOS ``read``.
Parameters
----------
available_info : dict
Variable metadata returned by ``FileReader.available_variables()``.
read_name : str
Actual ADIOS variable name, possibly source-qualified for campaigns.
requested_name : str
Reader-facing variable name used in returned dictionaries and errors.
requested_steps : sequence[int]
ADIOS step ids requested from the source.
fallback_step_count : int
Source-level step count used only when the per-variable metadata does
not advertise ``AvailableStepsCount``.
"""
if read_name not in available_info:
raise KeyError(f"Variable '{requested_name}' not found in ADIOS source as '{read_name}'.")
step_count = available_steps_count_for_variable(available_info[read_name]) or int(fallback_step_count or 0)
missing_steps = [int(step) for step in requested_steps if int(step) < 0 or int(step) >= step_count]
if missing_steps:
raise KeyError(
f"ADIOS step(s) unavailable for variable '{requested_name}' "
f"as '{read_name}': {', '.join(str(step) for step in missing_steps)}"
)
def _normalize_scalar_value(value, info):
"""
Convert ADIOS scalar reads to scalar values while preserving arrays.
Parameters
----------
value : object
Value returned by ``FileReader.read``.
info : dict
ADIOS metadata for the variable.
"""
if str(info.get("SingleValue", "")).lower() != "true":
return value
try:
return value.item()
except AttributeError:
return value
def _is_single_value(info) -> bool:
"""Return whether ADIOS reports the variable as scalar-valued."""
return str(info.get("SingleValue", "")).lower() == "true"
def _payload_shape_dims(info) -> list[int]:
"""
Return payload dimensions parsed from ADIOS ``Shape`` metadata.
ADIOS shape strings may include a leading ``"<steps> * ..."`` prefix in
some contexts. Only the per-step payload dimensions are needed here.
"""
shape_text = str(info.get("Shape", "") or "").strip()
match = re.match(r"^\s*(\d+)\s*\*\s*(.+)\s*$", shape_text)
if match:
shape_text = match.group(2).strip()
return [int(num) for num in re.findall(r"\d+", shape_text or "")]