"""HPC-Campaign-backed catalog discovery for XGC ADIOS outputs."""
from __future__ import annotations
from dataclasses import replace
from pathlib import Path
import re
import sqlite3
from time import perf_counter
from typing import Dict, Iterable, List, Tuple
import zlib
import adios2
import numpy as np
from .campaign_source_reader import CampaignSourceReader
from .directory_catalog import (
SimulationCatalog,
classify_product,
_layout_for_sources,
_merge_variables,
_variables_from_adios_info,
)
from .run_manifest import TextArtifact
from .types import CatalogDiscoveryTiming, DataProduct, MetadataReadTiming, SourceInfo, VariableInfo
_BP_SEQ_RE = re.compile(r"^(?P<base>.+)\.(?P<index>\d+)\.bp$")
_XGC_TIMESERIES_RE = re.compile(r"^xgc\.[A-Za-z0-9_.-]+$")
_LOGICAL_STEP_VARIABLE_NAMES = ("gstep", "step", "timestep")
_CAMPAIGN_SOURCE_ALIASES = {
"2d": "xgc.2d.bp",
"3d": "xgc.3d.bp",
"bfield": "xgc.bfield.bp",
"bfieldm": "xgc.bfieldm.bp",
"cnv_from_surf": "xgc.cnv_from_surf.bp",
"cnv_to_surf": "xgc.cnv_to_surf.bp",
"equil": "xgc.equil.bp",
"f0": "xgc.f0.bp",
"f0.mesh": "xgc.f0.mesh.bp",
"f2d": "xgc.f2d.bp",
"f3d": "xgc.f3d.bp",
"fsourcediag": "xgc.fsourcediag.bp",
"grad_rz": "xgc.grad_rz.bp",
"heatdiag2": "xgc.heatdiag2.bp",
"hyp_vis_rad": "xgc.hyp_vis_rad.bp",
"loading.0": "xgc.loading.0.bp",
"loading.1": "xgc.loading.1.bp",
"loop_vol": "xgc.loop_vol.bp",
"mesh": "xgc.mesh.bp",
"neutrals": "xgc.neutrals.bp",
"oneddiag": "xgc.oneddiag.bp",
"saved_volumes": "xgc.saved_volumes.bp",
"sheathdiag": "xgc.sheathdiag.bp",
"smooth_pol": "xgc.smooth_pol.bp",
"units": "xgc.units.bp",
"volumes": "xgc.volumes.bp",
}
[docs]
def open_campaign_catalog(
campaign_path: str | Path,
*,
collect_metadata: bool = True,
persistent_reads: bool = True,
) -> SimulationCatalog:
"""
Open an HPC-Campaign-backed XGC catalog.
Parameters
----------
campaign_path : str or pathlib.Path
Path to an ADIOS-readable campaign ``.aca`` file.
collect_metadata : bool, optional
Whether to inspect campaign variable metadata. The initial campaign
backend requires metadata collection to build useful product records.
persistent_reads : bool, optional
If True, reuse the catalog's open campaign ``FileReader`` for data
reads. This is the default because opening campaign archives can be
expensive, especially for remote datasets. Set False only as a
troubleshooting fallback; read-plan batching still prevents one open
per variable in that mode.
Returns
-------
xgc_analysis.data_access.catalog.SimulationCatalog
Catalog populated from campaign-qualified ADIOS variable names, with
archive-opening and metadata-read durations in ``discovery_timing``.
"""
discovery_start = perf_counter()
campaign_path = Path(campaign_path)
campaign_reader = adios2.FileReader(str(campaign_path))
discovery_timing = CatalogDiscoveryTiming()
products = discover_campaign_products(
campaign_path,
collect_metadata=collect_metadata,
file_reader=campaign_reader,
timing=discovery_timing,
metadata_start_s=discovery_start,
)
discovery_timing.total_s = perf_counter() - discovery_start
def refresh():
"""Rescan the archive and replace the catalog's discovery timing.
Returns
-------
list[DataProduct]
Freshly discovered campaign products.
"""
refreshed_timing = CatalogDiscoveryTiming()
refreshed_products = discover_campaign_products(
campaign_path,
collect_metadata=collect_metadata,
file_reader=getattr(source_reader, "file_reader", None) or campaign_reader,
timing=refreshed_timing,
)
catalog.discovery_timing = refreshed_timing
return refreshed_products
if persistent_reads:
def update_campaign_reader(reader):
"""Store a reopened campaign reader on the owning catalog.
``reader`` is the new ADIOS campaign reader handle.
Returns
-------
None
The owning catalog keeps the reopened handle.
"""
catalog.campaign_reader = reader
source_reader = CampaignSourceReader(campaign_reader, owns_reader=True, on_reopen=update_campaign_reader)
else:
source_reader = CampaignSourceReader()
def close_campaign_reader():
"""Close the one-off campaign reader when the catalog closes.
Returns
-------
None
The one-off archive reader is closed.
"""
close = getattr(campaign_reader, "close", None)
if close is not None:
close()
source_reader.close = close_campaign_reader
catalog = SimulationCatalog(
campaign_path.parent,
products,
refresh_func=refresh,
state_key_func=lambda: campaign_state_key(campaign_path),
source_reader=source_reader,
manifest=_campaign_text_manifest(campaign_path),
text_reader=lambda artifact_path: _read_campaign_text_with_lock(
campaign_path,
artifact_path,
source_reader,
),
discovery_timing=discovery_timing,
)
catalog.campaign_path = campaign_path
catalog.campaign_reader = campaign_reader
return catalog
def _read_campaign_text_with_lock(campaign_path: Path, artifact_path: str, source_reader) -> str:
"""
Read campaign text through the persistent reader guarded by source lock.
Parameters
----------
campaign_path : pathlib.Path
Campaign ``.aca`` file.
artifact_path : str
Text dataset name.
source_reader : CampaignSourceReader
Source reader whose lock serializes access to its persistent campaign
FileReader.
"""
lock = getattr(source_reader, "_lock", None)
campaign_reader = getattr(source_reader, "file_reader", None)
if lock is None:
return read_campaign_text(campaign_path, artifact_path, file_reader=campaign_reader)
with lock:
campaign_reader = getattr(source_reader, "file_reader", None)
return read_campaign_text(campaign_path, artifact_path, file_reader=campaign_reader)
[docs]
def read_campaign_text(campaign_path: str | Path, artifact_path: str, *, file_reader=None) -> str:
"""
Read one embedded text artifact from an HPC-Campaign archive.
Parameters
----------
campaign_path : str or pathlib.Path
Campaign ``.aca`` file.
artifact_path : str
Dataset name of the embedded text artifact, for example ``"input"`` or
``"input_dir/profile.dat"``.
file_reader : adios2.FileReader or None, optional
Existing campaign reader used if the artifact needs an ADIOS variable read.
Returns
-------
str
UTF-8 decoded text content.
Raises
------
KeyError
If the campaign archive does not contain the named text dataset.
RuntimeError
If the archive cannot be queried as an HPC-Campaign SQLite database.
"""
row = _campaign_text_row(campaign_path, artifact_path)
if row is not None:
compression, data = row
if data is not None:
raw = bytes(data)
if int(compression or 0) == 1:
raw = zlib.decompress(raw)
return raw.decode("utf-8")
return _read_campaign_text_variable(campaign_path, artifact_path, file_reader=file_reader)
def _campaign_text_manifest(campaign_path: Path) -> dict:
"""
Build a minimal manifest-like text-artifact section for a campaign archive.
Parameters
----------
campaign_path : pathlib.Path
Campaign ``.aca`` file to inspect.
"""
artifacts = [
TextArtifact(
path=name,
artifact_class=_campaign_text_class(name),
role=_campaign_text_role(name),
source="campaign_text_dataset",
).to_manifest()
for name in _campaign_text_names(campaign_path)
]
return {"text_artifacts": artifacts}
def _campaign_text_names(campaign_path: str | Path) -> List[str]:
"""
Return embedded text dataset names from an HPC-Campaign archive.
Parameters
----------
campaign_path : str or pathlib.Path
Campaign ``.aca`` file.
"""
try:
with sqlite3.connect(Path(campaign_path)) as connection:
rows = connection.execute(
"""
select distinct dataset.name
from dataset
where dataset.fileformat = ? and ifnull(dataset.deltime, 0) = 0
order by dataset.name
""",
("TEXT",),
).fetchall()
except sqlite3.Error:
return _campaign_text_variable_names(campaign_path)
return [str(row[0]) for row in rows]
def _campaign_text_row(campaign_path: str | Path, artifact_path: str):
"""
Return the newest embedded text blob row for one campaign dataset.
Parameters
----------
campaign_path : str or pathlib.Path
Campaign ``.aca`` file.
artifact_path : str
Text dataset name.
"""
try:
with sqlite3.connect(Path(campaign_path)) as connection:
return connection.execute(
"""
select file.compression, file.data
from dataset
join replica on dataset.rowid = replica.datasetid
join repfiles on replica.rowid = repfiles.replicaid
join file on repfiles.fileid = file.fileid
where dataset.name = ? and dataset.fileformat = ?
order by replica.modtime desc, file.modtime desc
limit 1
""",
(artifact_path, "TEXT"),
).fetchone()
except sqlite3.Error as exc:
raise RuntimeError(f"Could not query campaign text artifacts in '{campaign_path}'.") from exc
def _campaign_text_variable_names(campaign_path: str | Path) -> List[str]:
"""
Return ADIOS-readable text variable names from a campaign archive.
Parameters
----------
campaign_path : str or pathlib.Path
Campaign ``.aca`` file.
"""
try:
with adios2.FileReader(str(campaign_path)) as reader:
vars_info = reader.available_variables() or {}
except Exception:
return []
return sorted(
name for name, info in vars_info.items()
if str(info.get("Type", "")).lower() in {"string", "char", "unsigned char"}
or name in {"input", "fort.input.used", "adios2cfg.xml", "petsc.rc"}
or name.startswith("input_dir/")
)
def _read_campaign_text_variable(campaign_path: str | Path, artifact_path: str, *, file_reader=None) -> str:
"""
Read one text dataset exposed as an ADIOS variable by HPC-Campaign.
Parameters
----------
campaign_path : str or pathlib.Path
Campaign ``.aca`` file.
artifact_path : str
ADIOS variable name for the text dataset.
"""
try:
if file_reader is not None:
vars_info = file_reader.available_variables() or {}
if artifact_path not in vars_info:
raise KeyError(f"Unknown campaign text artifact: {artifact_path}")
value = file_reader.read(artifact_path)
else:
with adios2.FileReader(str(campaign_path)) as reader:
vars_info = reader.available_variables() or {}
if artifact_path not in vars_info:
raise KeyError(f"Unknown campaign text artifact: {artifact_path}")
value = reader.read(artifact_path)
except KeyError:
raise
except Exception as exc:
raise RuntimeError(
f"Could not read campaign text artifact '{artifact_path}' from '{campaign_path}'."
) from exc
return _decode_campaign_text_value(value)
def _decode_campaign_text_value(value) -> str:
"""
Decode a text payload returned by ADIOS into a Python string.
Parameters
----------
value : object
Value returned by ``adios2.FileReader.read`` for a text dataset.
"""
if isinstance(value, str):
return value
if isinstance(value, (bytes, bytearray)):
return bytes(value).decode("utf-8")
arr = np.asarray(value)
if arr.dtype.kind in {"S", "U"}:
if arr.ndim == 0:
return str(arr.item())
return "".join(str(item) for item in arr.reshape(-1))
if arr.dtype.kind in {"i", "u"}:
return bytes(int(item) for item in arr.reshape(-1)).decode("utf-8")
raise TypeError(f"Cannot decode campaign text value with dtype={arr.dtype} shape={arr.shape}")
def _campaign_text_class(name: str) -> str:
"""
Return a coarse artifact class for one campaign text dataset name.
Parameters
----------
name : str
Campaign text dataset name.
"""
if name == "input":
return "input.parameter"
if name in {"fort.input.used", "adios2cfg.xml", "petsc.rc"}:
return "input.config"
if name.endswith((".out", ".log")):
return "log.application"
return "input.support"
def _campaign_text_role(name: str) -> str:
"""
Return an optional role for one campaign text dataset name.
Parameters
----------
name : str
Campaign text dataset name.
"""
if name == "input":
return "xgc_input"
if name == "adios2cfg.xml":
return "adios_config"
return ""
[docs]
def discover_campaign_products(
campaign_path: Path,
*,
collect_metadata: bool = True,
file_reader=None,
timing: CatalogDiscoveryTiming | None = None,
metadata_start_s: float | None = None,
) -> List[DataProduct]:
"""
Discover XGC products advertised by a campaign ``.aca`` file.
Parameters
----------
campaign_path : pathlib.Path
Campaign file to inspect.
collect_metadata : bool, optional
If False, return no products. Campaign discovery currently depends on
the campaign's ADIOS metadata because there is no filesystem listing to
scan.
file_reader : adios2.FileReader or None, optional
Existing open campaign reader. Supplying this avoids reopening the ACA
file during catalog construction and refresh.
timing : CatalogDiscoveryTiming or None, optional
Mutable summary receiving one physical archive metadata-read duration
and the total discovery duration.
metadata_start_s : float or None, optional
``perf_counter`` timestamp taken before opening ``file_reader`` when
the caller opened it for this discovery. Omit for an already-open
reader, including refreshes. Without a supplied reader, this function
times its own archive open.
Returns
-------
list[DataProduct]
Products assembled from archive variable metadata, or an empty list
when ``collect_metadata`` is False.
"""
discovery_start = perf_counter() if metadata_start_s is None else metadata_start_s
campaign_path = Path(campaign_path)
metadata_succeeded = False
try:
if not collect_metadata:
return []
if file_reader is not None:
products = _discover_campaign_products_from_reader(campaign_path, file_reader)
else:
with adios2.FileReader(str(campaign_path)) as reader:
products = _discover_campaign_products_from_reader(campaign_path, reader)
metadata_succeeded = True
return products
finally:
if timing is not None:
elapsed_s = perf_counter() - discovery_start
timing.total_s = elapsed_s
if collect_metadata:
timing.metadata_reads.append(
MetadataReadTiming(
source_id=campaign_path.name,
path=campaign_path,
elapsed_s=elapsed_s,
succeeded=metadata_succeeded,
)
)
def _discover_campaign_products_from_reader(campaign_path: Path, reader) -> List[DataProduct]:
"""
Build campaign product records from an open FileReader.
Parameters
----------
campaign_path : pathlib.Path
Campaign file path used for source records and state metadata.
reader : adios2.FileReader
Open campaign FileReader handle.
"""
vars_info = reader.available_variables() or {}
attrs_info = reader.available_attributes() or {}
grouped = _group_campaign_metadata(vars_info, attrs_info)
sources_by_product = {}
for source_id, source_vars, source_attrs in grouped:
product_key, file_index = _campaign_product_key(source_id)
variables = _campaign_variables_from_info(source_vars, source_attrs)
adios_step_count = _campaign_step_count(source_vars)
step_values = []
time_values = []
tindex_values = []
# ADIOS 2.12.1 can crash when scalar step variables are read from
# campaign time-series aliases. Literal BP-backed campaign sources keep
# the richer logical-step metadata; aliases fall back to ADIOS ordinals.
if source_id.endswith(".bp"):
step_values = _read_campaign_scalar_steps(
reader,
source_id,
source_vars,
_LOGICAL_STEP_VARIABLE_NAMES,
adios_step_count,
int,
)
time_values = _read_campaign_scalar_steps(
reader,
source_id,
source_vars,
("time",),
adios_step_count,
float,
)
tindex_values = _read_campaign_scalar_steps(
reader,
source_id,
source_vars,
("tindex",),
adios_step_count,
int,
)
try:
mtime = campaign_path.stat().st_mtime
except OSError:
mtime = 0.0
sources_by_product.setdefault(product_key, []).append(
SourceInfo(
source_id=source_id,
path=campaign_path,
file_index=file_index,
mtime=mtime,
variables=variables,
step_values=step_values,
time_values=time_values,
tindex_values=tindex_values,
adios_step_count=adios_step_count,
)
)
products = []
for product_key in sorted(sources_by_product):
sources = sorted(
sources_by_product[product_key],
key=lambda source: (-1 if source.file_index is None else source.file_index, source.source_id),
)
product_type, family = classify_product(product_key)
products.append(
DataProduct(
key=product_key,
label=product_key,
product_type=product_type,
layout=_layout_for_sources(sources),
sources=sources,
variables=_merge_variables(sources),
product_family=family,
)
)
return products
[docs]
def campaign_state_key(campaign_path: Path):
"""
Return a cheap state key for one campaign file.
Parameters
----------
campaign_path : pathlib.Path
Campaign file path.
Returns
-------
tuple
Archive path, timestamp, and size signature; empty if unavailable.
"""
campaign_path = Path(campaign_path)
try:
stat = campaign_path.stat()
except OSError:
return tuple()
return ((campaign_path.name, "file", int(stat.st_mtime_ns), int(stat.st_size)),)
def _group_campaign_metadata(vars_info, attrs_info) -> Iterable[Tuple[str, Dict[str, dict], Dict[str, dict]]]:
"""
Group campaign-qualified metadata by source id.
Campaign variables are expected to use names like
``xgc.f2d.00010.bp/e_den`` or time-series aliases such as ``xgc.f2d/e_den``
and legacy aliases such as ``f2d/e_den``.
The returned variable and attribute mappings are stripped to reader-facing
names such as ``e_den`` and ``e_den/unit``.
"""
grouped_vars: Dict[str, Dict[str, dict]] = {}
grouped_attrs: Dict[str, Dict[str, dict]] = {}
for qualified_name, info in (vars_info or {}).items():
if "/" not in qualified_name:
continue
source_id, variable = qualified_name.split("/", 1)
if not _is_supported_campaign_source(source_id):
continue
if not variable:
continue
grouped_vars.setdefault(source_id, {})[variable] = info
for qualified_name, info in (attrs_info or {}).items():
if "/" not in qualified_name:
continue
source_id, attr_name = qualified_name.split("/", 1)
if not _is_supported_campaign_source(source_id):
continue
if not attr_name:
continue
grouped_attrs.setdefault(source_id, {})[attr_name] = info
for source_id in sorted(grouped_vars):
yield source_id, grouped_vars[source_id], grouped_attrs.get(source_id, {})
def _is_supported_campaign_source(source_id: str) -> bool:
"""
Return whether a campaign source prefix maps to an XGC data product.
Parameters
----------
source_id : str
Campaign source prefix before the first ``/`` in an ADIOS variable name.
Returns
-------
bool
True for literal BP names, XGC time-series names such as ``xgc.2d``,
and known legacy HPC-Campaign representation aliases.
"""
return (
source_id.endswith(".bp")
or _XGC_TIMESERIES_RE.match(source_id) is not None
or source_id in _CAMPAIGN_SOURCE_ALIASES
)
def _campaign_product_key(source_id: str) -> Tuple[str, int | None]:
"""
Return logical product key and optional file index for one campaign source id.
"""
if source_id in _CAMPAIGN_SOURCE_ALIASES:
return _CAMPAIGN_SOURCE_ALIASES[source_id], None
match = _BP_SEQ_RE.match(source_id)
if match:
return f"{match.group('base')}.bp", int(match.group("index"))
if source_id.endswith(".bp"):
return source_id, None
if _XGC_TIMESERIES_RE.match(source_id):
return f"{source_id}.bp", None
return source_id, None
def _campaign_variables_from_info(vars_info, attrs_info) -> Dict[str, VariableInfo]:
"""
Convert campaign variable metadata to unqualified ``VariableInfo`` records.
"""
variables = _variables_from_adios_info(vars_info, attrs_info)
for name, info in vars_info.items():
variables[name] = replace(
variables[name],
step_count=_available_steps_count(info),
)
return variables
def _campaign_step_count(vars_info) -> int:
"""
Return the maximum ``AvailableStepsCount`` advertised by a source.
"""
return max([_available_steps_count(info) for info in vars_info.values()] or [1])
def _available_steps_count(info) -> int:
"""
Parse ADIOS ``AvailableStepsCount`` metadata.
"""
try:
return int(info.get("AvailableStepsCount", 0) or 0)
except Exception:
return 0
def _read_campaign_scalar_steps(reader, source_id, vars_info, names, step_count, value_type):
"""
Read scalar coordinate values for one campaign source.
Parameters
----------
reader : adios2.FileReader
Open campaign reader.
source_id : str
Campaign source prefix.
vars_info : dict
Unqualified variable metadata for the source.
names : iterable[str]
Candidate unqualified variable names.
step_count : int
Number of ADIOS steps to inspect.
value_type : type
Conversion type for returned scalar values.
"""
for name in names:
info = vars_info.get(name)
if info is None:
continue
if str(info.get("SingleValue", "")).lower() not in {"true", "1", "yes"}:
continue
qualified_name = f"{source_id}/{name}"
values = []
for step in range(int(step_count or 0)):
try:
value = reader.read(qualified_name, step_selection=[step, 1])
except Exception:
return []
try:
values.append(value_type(np.asarray(value).reshape(-1)[0]))
except Exception:
return []
return values
return []