access-profiling 0.1__py3-none-any.whl

access/profiling/esmf_parser.py

@@ -0,0 +1,183 @@
+# Copyright 2025 ACCESS-NRI and contributors. See the top-level COPYRIGHT file for details.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Parser for ESMF profiling text summary data, such as output by nuopc.
+The data to be parsed is written in the following form:
+
+Region                                       PETs   PEs    Count    Mean (s)    Min (s)     Min PET Max (s)     Max PET
+  [ESMF]                                     1664   1664   1        2558.5684   2555.1450   279     2559.5801   817
+    [ensemble] RunPhase1                     1664   1664   1        1879.7292   1872.5078   376     1905.4939   1
+      [ESM0001] RunPhase1                    1664   1664   1        1879.7286   1872.5059   858     1905.4937   1
+        [OCN] RunPhase1                      1300   1300   960      1850.4170   1848.3905   1023    1858.6404   364
+        [OCN-TO-MED] RunPhase1               1664   1664   960      365.9532    0.1688      405     1673.2255   18
+        [ICE] RunPhase1                      364    364    960      155.8202    154.7637    94      160.2443    0
+          cice_run_total                     364    364    960      155.4648    154.3687    94      159.8980    0
+            cice_run_import                  364    364    960      3.7565      3.5426      218     8.3892      1
+              cice_imp_halo                  364    364    1920     2.1386      1.5405      111     6.9782      0
+              cice_imp_t2u                   364    364    960      0.9242      0.5578      355     1.5026      111
+              cice_imp_atm                   364    364    960      0.0579      0.0399      331     0.0834      18
+              cice_imp_ocn                   364    364    960      0.0499      0.0419      50      0.0632      12
+            cice_run_export                  364    364    960      1.1015      0.8846      361     1.3607      194
+        [MED-TO-OCN] RunPhase1               1664   1664   960      16.7498     0.4588      256     23.2875     1023
+        [MED] med_phases_restart_write       364    364    960      31.8234     31.8123     203     33.0513     363
+          MED:(med_phases_restart_write)     364    364    960      31.7651     31.7559     203     32.9919     363
+...
+
+Where indentation depth indicates depth in the call-stack. Usually there is a header like
+********
+A warning about identifying load-imbalance.
+********
+Note that the profiling summary stats may contain identical region names
+e.g. [ATM-TO-MED] RunPhase1 is found twice in OM3.
+"""
+
+import os
+from pathlib import Path
+
+from pint import Unit
+
+from access.profiling.metrics import (
+    ProfilingMetric,
+    count,
+    pemax,
+    pemin,
+    tavg,
+    tmax,
+    tmin,
+)
+from access.profiling.parser import ProfilingParser, _read_text_file
+
+pets = ProfilingMetric("PETs", Unit("dimensionless"), "ESMF Virtual Machine Persistent Execution Threads")
+pes = ProfilingMetric("PEs", Unit("dimensionless"), "Processing Elements")
+
+
+class ESMFSummaryProfilingParser(ProfilingParser):
+    """ESMF text summary profiling output parser."""
+
+    hierarchical: bool  # whether the call-stack hierarchy is parsed.
+
+    def __init__(self, hierarchical: bool = False):
+        """Instantiate ESMF profiling parser.
+
+        Args:
+            hierarchical (bool): Whether call-stack hierarchy is parsed.
+        """
+        super().__init__()
+
+        self.hierarchical = hierarchical
+
+        # ESMF provides the following metrics
+        self._metrics = [pets, pes, count, tavg, tmin, pemin, tmax, pemax]
+
+    def parse(self, file_path: str | Path | os.PathLike) -> dict:
+        stream = _read_text_file(file_path)
+
+        lines = stream.strip().split("\n")
+
+        if self.hierarchical:
+            result = {}
+            stack = [(result, -1)]  # (current_dict, indent_level)
+        else:
+            result = {m: [] for m in self._metrics}
+            result["region"] = []
+
+        for line in lines:
+            # Split the line into region name and statistics
+            parts = line.split()
+            if len(parts) < (1 + len(self._metrics)):  # Need at least 1 region + the number of stat columns
+                continue
+
+            # Extract region name and statistics
+            region = " ".join(parts[:-8])
+            stats = parts[-8:]
+
+            # Validate that all statistics can be parsed correctly
+            try:
+                stats_dict = {
+                    pets: int(stats[0]),
+                    pes: int(stats[1]),
+                    count: int(stats[2]),
+                    tavg: float(stats[3]),
+                    tmin: float(stats[4]),
+                    pemin: int(stats[5]),
+                    tmax: float(stats[6]),
+                    pemax: int(stats[7]),
+                }
+            except (ValueError, IndexError):
+                # Skip lines that don't match the expected format
+                continue
+
+            if self.hierarchical:
+                # Calculate indentation level (each level is 2 spaces)
+                indent = len(line) - len(line.lstrip())
+                indent_level = indent // 2
+
+                # Pop stack until we find the parent level
+                while stack and stack[-1][1] >= indent_level:
+                    stack.pop()
+
+                # Get parent dictionary
+                parent_dict = stack[-1][0]
+
+                # Create entry for this region
+                if region not in parent_dict:
+                    parent_dict[region] = {}
+
+                # Add statistics to this region
+                parent_dict[region].update(stats_dict)
+
+                # Push this level onto stack for potential children
+                stack.append((parent_dict[region], indent_level))
+            else:
+                _update_flat_result(result, stats_dict, region)
+
+        # fewer if statements to pass ruff checks
+        if (self.hierarchical and not result) or (not self.hierarchical and len(result["region"]) == 0):
+            raise ValueError("No ESMF summary profiling data found")
+
+        return result
+
+
+def _update_flat_result(result: dict, stats_dict: dict, region: str):
+    """Helper function to update flat result.
+
+    Besides appending results, this function also checks whether the region already exists
+    and aggregates the metric values in result if possible.
+
+    Args:
+        result (dict): The flat result dictionary to update.
+        stats_dict (dict): The stats to update the result with.
+        region (str): The region to append the results to.
+
+    Raises:
+        NotImplementedError: If a stats_dict["region"] is already in result["region"],
+                             but the PETs or PEs value aren't the same.
+    """
+    # Flat structure: just use region name as key
+    try:
+        idx = result["region"].index(region)
+        # only update existing region if PETs and PEs are same
+        if (
+            result[pets][idx] == stats_dict[pets]
+            and result[pes][idx] == stats_dict[pes]
+            and stats_dict[pets] == stats_dict[pes]
+        ):
+            # new avg is weighted average using count as the weight
+            result[tavg][idx] = (result[tavg][idx] * result[count][idx] + stats_dict[tavg] * stats_dict[count]) / (
+                result[count][idx] + stats_dict[count]
+            )
+            result[count][idx] += stats_dict[count]
+            if stats_dict[tmin] < result[tmin][idx]:
+                result[tmin][idx] = stats_dict[tmin]
+                result[pemin][idx] = stats_dict[pemin]
+            if stats_dict[tmax] > result[tmax][idx]:
+                result[tmax][idx] = stats_dict[tmax]
+                result[pemax][idx] = stats_dict[pemax]
+        else:
+            raise NotImplementedError(
+                "I don't know what to do with multiple regions with same name, but different PETs/PEs."
+            )
+    except ValueError:
+        result["region"].append(region)
+        for k, v in stats_dict.items():
+            result[k].append(v)

access/profiling/experiment.py

@@ -0,0 +1,263 @@
+# Copyright 2025 ACCESS-NRI and contributors. See the top-level COPYRIGHT file for details.
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+import tarfile
+import tempfile
+from contextlib import contextmanager
+from enum import Enum
+from pathlib import Path
+
+import xarray as xr
+
+from access.profiling.parser import ProfilingParser, flatten_hierarchical
+
+logger = logging.getLogger(__name__)
+
+
+def _make_unique_region_names(regions: list[object]) -> list[object]:
+    """Return region names with deterministic suffixes for duplicates."""
+
+    counts: dict[object, int] = {}
+    unique_regions: list[object] = []
+    for region in regions:
+        count = counts.get(region, 0) + 1
+        counts[region] = count
+        unique_regions.append(region if count == 1 else f"{region}_{count}")
+    return unique_regions
+
+
+class ProfilingLog:
+    """Represents a profiling log file.
+
+    Args:
+        filepath (Path): Path to the log file.
+        parser (ProfilingParser): Parser to use for this log file.
+        optional (bool): Whether this log might be missing or does not contain parsable data. If True, no error should
+        be raised if the log is missing or unparsable. Defaults to False.
+    """
+
+    filepath: Path  # Path to the log file
+    parser: ProfilingParser  # Parser to use for this log file
+    _optional: bool = False  # Whether this log might not be present
+
+    def __init__(self, filepath: Path, parser: ProfilingParser, optional: bool = False) -> None:
+        self.filepath = filepath
+        self.parser = parser
+        self._optional = optional
+
+    @property
+    def optional(self) -> bool:
+        """bool: Whether this log might not be present."""
+        return self._optional
+
+    def parse(self) -> xr.Dataset:
+        """Parses the log file and returns the profiling data as an xarray Dataset.
+
+        Accepts all three parser output formats (see parser.py module docstring):
+
+        - **Flat**: standard 1D Dataset over the ``region`` dimension.
+        - **Hierarchical nested dict**: automatically flattened via
+          :func:`flatten_hierarchical` before building the Dataset.
+        - **Per-PE**: produces a 2D Dataset with both ``region`` and ``pe`` dimensions.
+          Use :func:`aggregate_pe_data` on the result to compute summary statistics.
+
+        Returns:
+           xr.Dataset: Parsed profiling data.
+        """
+        data = self.parser.parse(self.filepath)
+
+        # Flatten hierarchical (nested dict) format if needed
+        if "region" not in data:
+            data = flatten_hierarchical(data, self.parser.metrics)
+
+        has_pe = "pe" in data
+        dims = ["region", "pe"] if has_pe else ["region"]
+        coords: dict = {"region": _make_unique_region_names(list(data["region"]))}
+        if has_pe:
+            coords["pe"] = data["pe"]
+
+        return xr.Dataset(
+            data_vars=dict(
+                zip(
+                    self.parser.metrics,
+                    [xr.DataArray(data[m], dims=dims).pint.quantify(m.units) for m in self.parser.metrics],
+                    strict=True,
+                )
+            ),
+            coords=coords,
+        )
+
+
+class ProfilingExperimentStatus(Enum):
+    """Enumeration representing the status of a profiling experiment."""
+
+    NEW = 1  # Experiment has been created but not started
+    RUNNING = 2  # Experiment is running or is queued
+    DONE = 3  # Experiment has finished
+    ARCHIVED = 4  # Experiment has been archived
+
+
+def experiment_directory_walker(path: Path, arcname: Path, root: Path, follow_symlinks: bool = False):
+    """Walks through the experiment directory, yielding files and corresponding names in the archive.
+
+    Symlinks are treated in a special manner.
+        - if the target is inside the experiment directory, the symlink itself is always returned
+        - if follow_symlinks is True and the target is a directory, then the all contents in the target directory are
+            recursively iterated
+        - if follow_symlinks is True and the target is a file, then the target file name is returned, not the symlink
+        - if follow_symlinks is False, then the symlink itself is returned for both files and directories
+
+    Args:
+        path (Path): Path to walk through.
+        arcname (Path): Archive name for the current path.
+        follow_symlinks (bool): Whether to follow symlinks. Defaults to False.
+
+    Yields:
+        Tuple[Path, Path]: A tuple containing the file path and its archive name.
+    """
+    if path.is_symlink():
+        if not follow_symlinks:
+            # Add symlink itself without following
+            yield path, arcname
+        else:
+            target = path.resolve()
+            if target.is_dir():
+                # Recursively add target contents
+                for child in target.iterdir():
+                    yield from experiment_directory_walker(
+                        child, Path(arcname) / child.name, root, follow_symlinks=follow_symlinks
+                    )
+            elif target.absolute().is_relative_to(root.absolute()):
+                # Target is within the experiment directory, so add symlink as is
+                yield path, arcname
+            else:
+                # Target is outside the experiment directory, add the target file instead
+                yield target, arcname
+
+    elif path.is_dir():
+        # Recursively add directory contents
+        for child in path.iterdir():
+            yield from experiment_directory_walker(
+                child, Path(arcname) / child.name, root, follow_symlinks=follow_symlinks
+            )
+    else:
+        yield path, arcname
+
+
+class ProfilingExperiment:
+    """Represents a profiling experiment.
+
+    Args:
+        path (Path): Path to the experiment directory.
+        run_path (Path | None): Path to a separate runs directory. If None, runs are assumed to be
+            inside path. When provided, the runs directory is also traversed during archival.
+            path contents are stored under experiment/ and run_path contents under runs/.
+    """
+
+    path: Path  # Path to the experiment directory
+    run_path: Path | None  # Path to a separate runs directory, or None
+    status: ProfilingExperimentStatus = ProfilingExperimentStatus.NEW  # Status of the experiment
+
+    def __init__(self, path: Path, run_path: Path | None = None) -> None:
+        self.path = path
+        self.run_path = run_path
+        if self.path.suffixes == [".tar", ".gz"]:
+            self.status = ProfilingExperimentStatus.ARCHIVED
+
+    def __repr__(self) -> str:
+        """Returns a string representation of the ProfilingExperiment."""
+        if self.run_path is not None:
+            return f"{type(self).__name__}(path={self.path!r}, run_path={self.run_path!r}, status={self.status.name})"
+        return f"{type(self).__name__}(path={self.path!r}, status={self.status.name})"
+
+    @contextmanager
+    def directory(self):
+        """Context manager returning the experiment and runs directories.
+
+        If the experiment has been archived, it will be extracted to a temporary directory. Otherwise, the original
+        directory paths will be used. Note that after exiting the context, the temporary directory is removed.
+
+        Returns:
+            tuple[Path, Path | None]: The experiment directory path and optional runs directory path.
+        """
+        if self.path.suffixes == [".tar", ".gz"]:
+            with tempfile.TemporaryDirectory(prefix="access-profiling_", suffix="_data") as tmpdir:
+                with tarfile.open(self.path) as tar:
+                    tar.extractall(path=Path(tmpdir), filter="data")
+                path = Path(tmpdir) / "experiment"
+                run_path = Path(tmpdir) / "runs"
+                yield path, run_path if run_path.exists() else None
+        else:
+            yield self.path, self.run_path
+
+    def archive(
+        self,
+        archive_path: Path,
+        exclude_dirs: list[str] | None = None,
+        exclude_files: list[str] | None = None,
+        follow_symlinks: bool = False,
+        overwrite: bool = False,
+    ):
+        """Archives the experiment to the specified archive path.
+
+        Only experiments with status DONE will be archived. No error will be raised if the experiment is not DONE.
+
+        Symlinks to files and directories inside the experiment directory will be include as symlinks. Symlinks to files
+        and directories outside the experiment directory will be followed if follow_symlinks is True, otherwise they
+        will be included as symlinks. path contents are stored under experiment/ in the archive. If
+        run_path is set, its contents are stored under runs/ in the archive.
+
+        Args:
+            archive_path (Path): Path to the archive destination. This should include the file name, but without
+            the .tar.gz suffix.
+            exclude_dirs (list[str] | None): Directory patterns to exclude when archiving.
+            exclude_files (list[str] | None): File patterns to exclude when archiving.
+            follow_symlinks (bool): Whether to follow symlinks when archiving. Defaults to False.
+            overwrite (bool): Whether to overwrite existing archives. Defaults to False.
+
+        Raises:
+            FileExistsError: If the archive destination already exists and overwrite is False.
+            ValueError: If the experiment status is unknown.
+        """
+        if self.status == ProfilingExperimentStatus.NEW:
+            logger.warning(f"Experiment at {self.path} is not yet started. Skipping archiving.", stacklevel=2)
+            return
+        elif self.status == ProfilingExperimentStatus.RUNNING:
+            logger.warning(f"Experiment at {self.path} is still running. Skipping archiving.", stacklevel=2)
+            return
+        elif self.status == ProfilingExperimentStatus.DONE:
+            logger.info(f"Archiving experiment at {self.path} to {archive_path.with_suffix('.tar.gz')}")
+        elif self.status == ProfilingExperimentStatus.ARCHIVED:
+            logger.warning(f"Experiment at {self.path} is already archived. Skipping archiving.", stacklevel=2)
+            return
+
+        archive_file = archive_path.with_suffix(".tar.gz")
+        mode = "w:gz" if overwrite else "x:gz"
+        if not overwrite and archive_file.exists():
+            raise FileExistsError(f"Archive destination {archive_file} already exists.")
+
+        exclude_dirs = exclude_dirs or []
+        exclude_files = exclude_files or []
+
+        paths_to_walk = (
+            [(self.path, Path("experiment"))]
+            if self.run_path is None
+            else [(self.path, Path("experiment")), (self.run_path, Path("runs"))]
+        )
+
+        with tarfile.open(archive_file, mode) as tar:
+            for root, prefix in paths_to_walk:
+                for file, arcname in experiment_directory_walker(root, prefix, root, follow_symlinks=follow_symlinks):
+                    # Skip if file is inside an excluded directory pattern
+                    if any(any(parent.match(pat) for pat in exclude_dirs) for parent in file.parents):
+                        continue
+                    # Skip if the file itself matches an excluded filename pattern
+                    if any(file.match(pat) for pat in exclude_files):
+                        continue
+                    logger.debug(f"Archiving file: {file} as {arcname}")
+                    tar.add(file, arcname=arcname)
+
+        self.status = ProfilingExperimentStatus.ARCHIVED
+        self.path = archive_file
+        self.run_path = None

access/profiling/fms_parser.py

@@ -0,0 +1,80 @@
+# Copyright 2025 ACCESS-NRI and contributors. See the top-level COPYRIGHT file for details.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Parser for FMS profiling data, such as output by MOM5 and MOM6.
+The data to be parsed is written in the following form:
+
+                                   hits          tmin          tmax          tavg          tstd  tfrac grain pemin pemax
+Total runtime                        1    138.600364    138.600366    138.600365      0.000001  1.000     0     0    11
+Ocean Initialization                 2      2.344926      2.345701      2.345388      0.000198  0.017    11     0    11
+Ocean                               23     86.869466     86.871652     86.870450      0.000744  0.627     1     0    11
+Ocean dynamics                      96     43.721019     44.391032     43.957944      0.244785  0.317    11     0    11
+Ocean thermodynamics and tracers    72     27.377185     33.281659     29.950144      1.792324  0.216    11     0    11
+ MPP_STACK high water mark=          0
+"""
+
+import os
+import re
+from pathlib import Path
+
+from pint import Unit
+
+from access.profiling.metrics import ProfilingMetric, count, pemax, pemin, tavg, tfrac, tmax, tmin, tstd
+from access.profiling.parser import ProfilingParser, _convert_from_string, _read_text_file
+
+grain = ProfilingMetric("grain", Unit("dimensionless"), "Grain")
+
+
+class FMSProfilingParser(ProfilingParser):
+    """FMS profiling output parser."""
+
+    has_hits: bool  # whether FMS timings contains "hits" column.
+
+    def __init__(self, has_hits: bool = True):
+        """Instantiate FMS profiling parser.
+
+        Args:
+            has_hits (bool): whether FMS timings contains "hits" column.
+        """
+        super().__init__()
+
+        self.has_hits = has_hits
+        # FMS provides the following metrics:
+        self._metrics = [count] if self.has_hits else []
+        self._metrics += [tmin, tmax, tavg, tstd, tfrac, grain, pemin, pemax]
+
+    def parse(self, file_path: str | Path | os.PathLike) -> dict:
+        stream = _read_text_file(file_path)
+
+        labels = ["hits"] if self.has_hits else []
+        labels += ["tmin", "tmax", "tavg", "tstd", "tfrac", "grain", "pemin", "pemax"]
+
+        # Regular expression to extract the profiling section from the file
+        header = r"\s*" + r"\s*".join(labels) + r"\s*"
+        footer = r" MPP_STACK high water mark=\s*\d*"
+        profiling_section_p = re.compile(header + r"(.*)" + footer, re.DOTALL)
+
+        # Regular expression to parse the data for each region
+        profile_line = r"^\s*(?P<region>[a-zA-Z:()_/\-*&\s]+(?<!\s))"
+        for label in labels:
+            profile_line += r"\s+(?P<" + label + r">[0-9.]+)"
+        profile_line += r"$"
+        profiling_region_p = re.compile(profile_line, re.MULTILINE)
+
+        # Parse data
+        stats = {"region": []}
+        stats.update({m: [] for m in self.metrics})
+        match = profiling_section_p.search(stream)
+        if match is None:
+            raise ValueError("No FMS profiling data found")
+        else:
+            profiling_section = match.group(1)
+        for line in profiling_region_p.finditer(profiling_section):
+            stats["region"].append(line.group("region"))
+            for label, metric in zip(labels, self.metrics, strict=True):
+                stats[metric].append(_convert_from_string(line.group(label)))
+
+        # Convert time fraction to percentage
+        stats[tfrac] = [val * 100 for val in stats[tfrac]]
+
+        return stats