access-profiling 0.1__py3-none-any.whl

access/profiling/parser.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
+
+"""Classes and utilities for reading and transforming profiling data.
+
+Data formats
+------------
+Parsers return a plain dict. Three shapes are supported:
+
+Flat (standard)
+    One list per metric, all the same length as 'region':
+
+        {'region': [...], metric_a: [...], metric_b: [...]}
+
+Hierarchical (nested dict)
+    Used when regions form a call-stack tree (e.g. ESMF). String keys are child
+    region names; ProfilingMetric keys are metric values. The two key types never
+    collide, so no separator is needed:
+
+        {'[ESMF]': {tavg: 2558.6, '[ICE] RunPhase1': {tavg: 155.8, ...}}}
+
+    Use flatten_hierarchical() to convert to the flat format.
+
+Per-PE
+    Each region has one measurement per processing element (MPI process). Metric
+    values are 2D lists of shape (n_regions, n_pes); a 'pe' key holds the PE IDs:
+
+        {'region': [...], 'pe': [0, 1, ..., N-1], metric_a: [[...], [...]]}
+
+    Use aggregate_pe_data() on the resulting xarray Dataset to reduce over PEs
+    and compute load-imbalance statistics (see metrics.py for the naming convention).
+
+Utilities
+---------
+flatten_hierarchical(data, metrics)
+    Converts a hierarchical nested dict to the standard flat dict (DFS pre-order).
+aggregate_pe_data(ds)
+    Reduces a per-PE Dataset along 'pe', producing variables named
+    {var}_{stat}_pe for each input variable and statistic.
+"""
+
+import os
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Any
+
+# Next import is required to register pint with xarray
+import pint_xarray  # noqa: F401
+import xarray as xr
+
+from access.profiling.metrics import ProfilingMetric
+
+
+class ProfilingParser(ABC):
+    """Abstract parser of profiling data.
+
+    The main purpose of a parser is to read profiling data from a file and return it
+    as a dict. Three output shapes are supported (see module docstring for full details):
+
+    Flat (standard)::
+
+        {
+            'region': ['region1', 'region2', ...],
+            metric_a: [val1a, val2a, ...],
+            metric_b: [val1b, val2b, ...],
+        }
+
+    Hierarchical (nested dict, no 'region' key)::
+
+        {
+            'root_region': {
+                metric_a: val,
+                'child_region': {metric_a: val, ...},
+            }
+        }
+
+    Per-PE (flat with an additional 'pe' key; metric values are 2D lists)::
+
+        {
+            'region': ['region1', 'region2', ...],
+            'pe': [0, 1, ..., N-1],
+            metric_a: [[pe0_val1, pe1_val1, ...], [pe0_val2, pe1_val2, ...]],
+        }
+    """
+
+    _metrics: list[ProfilingMetric]
+
+    @property
+    def metrics(self) -> list[ProfilingMetric]:
+        """list: Metrics available when using this parser."""
+        return self._metrics
+
+    @abstractmethod
+    def parse(self, file_path: str | Path | os.PathLike) -> dict:
+        """Parse the given file.
+
+        Args:
+            file_path (str | Path | os.PathLike): file to parse.
+
+        Returns:
+            dict: profiling data in one of the three formats described in the class
+                docstring (flat, hierarchical, or per-PE).
+
+        Raises:
+            ValueError: If no parsable text is found in file_path.
+            TypeError: If file_path cannot be converted to a valid Path object.
+            FileNotFoundError: If file_path doesn't exist or isn't a file.
+        """
+
+
+def flatten_hierarchical(data: dict, metrics: list[ProfilingMetric]) -> dict:
+    """Converts a hierarchical (nested dict) parser output into the standard flat format.
+
+    Traverses the nested dict depth-first (pre-order: parent before children).
+    At each node, string keys are treated as child region names and ProfilingMetric
+    keys as metric values — these types never collide so no separator is needed.
+
+    Args:
+        data (dict): Nested dict as returned by a hierarchical parser. At each level,
+            string keys are child region names and ProfilingMetric keys are metric values.
+        metrics (list[ProfilingMetric]): Metrics to extract from each node.
+
+    Returns:
+        dict: Standard flat profiling dict with a 'region' key and one list per metric.
+    """
+    result: dict = {"region": []}
+    for m in metrics:
+        result[m] = []
+
+    def _visit(node: dict, name: str) -> None:
+        result["region"].append(name)
+        for m in metrics:
+            result[m].append(node.get(m))
+        for key, value in node.items():
+            if isinstance(key, str):  # child region — ProfilingMetric keys are not str
+                _visit(value, key)
+
+    for region_name, region_data in data.items():
+        _visit(region_data, region_name)
+
+    return result
+
+
+def aggregate_pe_data(ds: xr.Dataset) -> xr.Dataset:
+    """Aggregates a per-PE profiling Dataset into summary statistics over the pe dimension.
+
+    For each data variable, computes the following reductions over the 'pe' dimension
+    and returns them as new string-named variables following the {var}_{stat}_pe
+    convention defined in metrics.py:
+
+    ========================  =============================================
+    Variable                  Description
+    ========================  =============================================
+    ``{var}_min_pe``          minimum across PEs
+    ``{var}_max_pe``          maximum across PEs
+    ``{var}_mean_pe``         mean across PEs
+    ``{var}_median_pe``       median across PEs
+    ``{var}_std_pe``          standard deviation across PEs
+    ``{var}_total_pe``        sum across PEs (total work done by all PEs)
+    ``{var}_argmin_pe``       PE index of the minimum value (dimensionless)
+    ``{var}_argmax_pe``       PE index of the maximum value (dimensionless)
+    ``{var}_imbalance_pe``    (max - min) / mean; 0 = perfectly balanced
+    ========================  =============================================
+
+    Args:
+        ds (xr.Dataset): Dataset with a 'pe' dimension, as returned by
+            ProfilingLog.parse() for per-PE parser output.
+
+    Returns:
+        xr.Dataset: New Dataset with 'pe' reduced, containing the derived variables
+            described above. All non-pe coordinates are preserved.
+
+    Raises:
+        ValueError: If ds does not have a 'pe' dimension.
+    """
+    if "pe" not in ds.dims:
+        raise ValueError("Dataset does not have a 'pe' dimension.")
+
+    result_vars = {}
+    for var in ds.data_vars:
+        da = ds[var]
+        base = str(var).replace(" ", "_")
+        vmin = da.min("pe")
+        vmax = da.max("pe")
+        vmean = da.mean("pe")
+        result_vars[f"{base}_min_pe"] = vmin
+        result_vars[f"{base}_max_pe"] = vmax
+        result_vars[f"{base}_mean_pe"] = vmean
+        result_vars[f"{base}_median_pe"] = da.median("pe")
+        result_vars[f"{base}_std_pe"] = da.std("pe")
+        result_vars[f"{base}_total_pe"] = da.sum("pe")
+        result_vars[f"{base}_argmin_pe"] = da.argmin("pe")
+        result_vars[f"{base}_argmax_pe"] = da.argmax("pe")
+        result_vars[f"{base}_imbalance_pe"] = xr.where(vmean != 0, (vmax - vmin) / vmean, 0)
+
+    coords = {k: v for k, v in ds.coords.items() if k != "pe"}
+    return xr.Dataset(result_vars, coords=coords)
+
+
+def _convert_from_string(value: str) -> Any:
+    """Tries to convert a string to the most appropriate numeric type. Leaves it unchanged if conversion does not
+    succeed.
+
+    Args:
+        value (str): string to convert.
+
+    Returns:
+        Any: the converted string or the original string.
+    """
+    for type_conversion in (int, float):
+        try:
+            return type_conversion(value)
+        except Exception:
+            continue
+    return value
+
+
+def _test_file(file_path: str | Path | os.PathLike) -> Path:
+    """Checks whether file_path is a valid path.
+
+    Args:
+        file_path (str | Path | os.PathLike): the path to check.
+
+    Returns
+        Path: file_path as a Path object.
+
+    Raises:
+        TypeError: if file_path cannot be turned into a Path object.
+        FileNotFoundError: if file_path can be converted into a Path, but the file dosen't exist.
+    """
+
+    try:
+        path = Path(file_path)
+    except TypeError as e:
+        raise TypeError(f"{file_path} is not a valid path.") from e
+
+    if not path.is_file():
+        raise FileNotFoundError(f"{file_path} is not a file or doesn't exist.")
+
+    return path
+
+
+def _read_text_file(file_path: str | Path | os.PathLike) -> str:
+    """Checks whether file_path is a valid path to a text file and tries to read it.
+
+    Args:
+        file_path (str | Path | os.PathLike): the path to check/read
+
+    Returns:
+        str: The text within the file.
+
+    Raises:
+        TypeError: if file_path is not a valid path
+        FileNotFoundError: if file_path is a path, but is not a file or doesn't exist.
+        ValueError: if file_path is a file, but cannot be read as a text file.
+    """
+
+    path = _test_file(file_path)
+
+    try:
+        return path.read_text()
+    except UnicodeDecodeError as e:
+        raise ValueError(f"{file_path} is not a text file.") from e

access/profiling/payu_manager.py

@@ -0,0 +1,338 @@
+# Copyright 2025 ACCESS-NRI and contributors. See the top-level COPYRIGHT file for details.
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+from abc import ABC, abstractmethod
+from collections.abc import Callable
+from datetime import timedelta
+from pathlib import Path
+
+from access.config import YAMLParser
+from access.config.esm1p6_layout_input import LayoutSearchConfig
+from access.config.layout_config import LayoutTuple
+from experiment_generator.experiment_generator import ExperimentGenerator
+from experiment_runner.experiment_runner import ExperimentRunner
+
+from access.profiling.experiment import ProfilingLog
+from access.profiling.manager import ProfilingExperiment, ProfilingExperimentStatus, ProfilingManager
+from access.profiling.payujson_parser import PayuJSONProfilingParser
+
+logger = logging.getLogger(__name__)
+
+
+class PayuManager(ProfilingManager, ABC):
+    """Abstract base class to handle profiling of Payu configurations."""
+
+    _repository_directory: str = "config"  # Repository directory name needed by the experiment generator and runner.
+    _nruns: int = 1  # Number of repetitions for the Payu experiments.
+    _startfrom_restart: str = "cold"  # Restart option for the Payu experiments.
+
+    @abstractmethod
+    def get_component_logs(self, path: Path) -> dict[str, ProfilingLog]:
+        """Returns available profiling logs for the components in the configuration.
+
+        Args:
+            path (Path): Path to the output directory.
+        Returns:
+            dict[str, ProfilingLog]: Dictionary mapping component names to their ProfilingLog instances.
+        """
+
+    @property
+    @abstractmethod
+    def model_type(self) -> str:
+        """Returns the model type identifier, as defined in Payu."""
+
+    @abstractmethod
+    def generate_core_layouts_from_node_count(
+        self,
+        num_nodes: float,
+        cores_per_node: int,
+        layout_search_config: LayoutSearchConfig | None = None,
+    ) -> list:
+        """Generates core layouts from the given number of nodes.
+
+        Args:
+            num_nodes (float): Number of nodes.
+            cores_per_node (int): Number of cores per node.
+            layout_search_config (LayoutSearchConfig | None): Configuration for layout search.
+        """
+
+    @abstractmethod
+    def generate_perturbation_block(self, layout: LayoutTuple, branch_name_prefix: str) -> dict:
+        """Generates a perturbation block for the given layout to be passed to the experiment generator.
+
+        Args:
+            layout (LayoutTuple): Core layout tuple.
+            branch_name_prefix (str): Branch name prefix.
+        Returns:
+            dict: Perturbation block configuration.
+        """
+
+    @property
+    def nruns(self) -> int:
+        """Returns the number of repetitions for the Payu experiments.
+
+        Returns:
+            int: Number of repetitions.
+        """
+        return self._nruns
+
+    @nruns.setter
+    def nruns(self, value: int) -> None:
+        """Sets the number of repetitions for the Payu experiments.
+
+        Args:
+            value (int): Number of repetitions.
+        """
+        if value < 0:
+            raise ValueError("Number of runs must be at least 0.")
+        self._nruns = value
+
+    @property
+    def startfrom_restart(self) -> str:
+        """Returns the restart option for the Payu experiments.
+
+        Returns:
+            str: Restart option.
+        """
+        return self._startfrom_restart
+
+    @startfrom_restart.setter
+    def startfrom_restart(self, value: str) -> None:
+        """Sets the restart option for the Payu experiments.
+
+        Args:
+            value (str): Restart option.
+        """
+        self._startfrom_restart = value
+
+    def set_control(self, repository, commit) -> None:
+        """Sets the control experiment from an existing Payu configuration.
+
+        Args:
+            repository: Git repository URL or path.
+            commit: Git commit hash or identifier.
+        """
+        self._repository = repository
+        self._control_commit = commit
+
+    def generate_scaling_experiments(
+        self,
+        num_nodes_list: list[float],
+        control_options: dict,
+        cores_per_node: int,
+        tol_around_ctrl_ratio: float,
+        max_wasted_ncores_frac: float | Callable[[float], float],
+        walltime: float | Callable[[float], float],
+    ) -> None:
+        """Generates scaling experiments using the ExperimentGenerator.
+
+        Args:
+            num_nodes_list (list[int]): List of number of nodes to generate experiments for.
+            control_options (dict): Options for the control experiment.
+            cores_per_node (int): Number of cores per node.
+            tol_around_ctrl_ratio (float): Tolerance around control core ratio for layout generation.
+            max_wasted_ncores_frac (float | Callable[[float], float]): Maximum fraction of wasted cores allowed.
+            walltime (float | Callable[[float], float]): Walltime in hours for each experiment.
+        """
+
+        generator_config = {
+            "model_type": self.model_type,
+            "repository_url": self._repository,
+            "start_point": self._control_commit,
+            "test_path": str(self.work_dir),
+            "repository_directory": self._repository_directory,
+            "control_branch_name": "ctrl",
+            "Control_Experiment": control_options,
+        }
+
+        seen_layouts = set()
+        seqnum = 1
+        generator_config["Perturbation_Experiment"] = {}
+        for num_nodes in num_nodes_list:
+            mwf = max_wasted_ncores_frac(num_nodes) if callable(max_wasted_ncores_frac) else max_wasted_ncores_frac
+            layout_config = LayoutSearchConfig(tol_around_ctrl_ratio=tol_around_ctrl_ratio, max_wasted_ncores_frac=mwf)
+            layouts = self.generate_core_layouts_from_node_count(
+                num_nodes,
+                cores_per_node=cores_per_node,
+                layout_search_config=layout_config,
+            )
+            if not layouts:
+                logger.warning(f"No layouts found for {num_nodes} nodes")
+                continue
+
+            layouts = [x for x in layouts if x not in seen_layouts]
+            seen_layouts.update(layouts)
+            logger.info(f"Generated {len(layouts)} layouts for {num_nodes} nodes. Layouts: {layouts}")
+
+            # TODO: the branch name needs to be simpler and model agnostic
+            branch_name = f"layout-unused-cores-to-cice-{layout_config.allocate_unused_cores_to_ice}"
+            walltime_hrs = walltime(num_nodes) if callable(walltime) else walltime
+
+            for layout in layouts:
+                pert_config = self.generate_perturbation_block(layout=layout, branch_name_prefix=branch_name)
+                branch = pert_config["branches"][0]
+                pert_config["config.yaml"]["walltime"] = str(timedelta(hours=walltime_hrs))
+
+                generator_config["Perturbation_Experiment"][f"Experiment_{seqnum}"] = pert_config
+                self.experiments[branch] = ProfilingExperiment(path=self.work_dir / branch / self._repository_directory)
+
+                seqnum += 1
+
+        ExperimentGenerator(generator_config).run()
+
+    def run_experiments(self) -> None:
+        """Runs Payu experiments for profiling data generation."""
+
+        runner_config = {
+            "test_path": self.work_dir,
+            "repository_directory": self._repository_directory,
+            "running_branches": [],
+            "keep_uuid": True,
+            "nruns": [],
+            "startfrom_restart": [],
+        }
+
+        for path, exp in self.experiments.items():
+            if exp.status == ProfilingExperimentStatus.NEW:
+                runner_config["running_branches"].append(path)
+                runner_config["nruns"].append(self.nruns)
+                runner_config["startfrom_restart"].append(self.startfrom_restart)
+                exp.status = ProfilingExperimentStatus.RUNNING
+
+        # Run the experiment runner
+        if runner_config["running_branches"]:
+            ExperimentRunner(runner_config).run()
+        else:
+            logger.info("No new experiments to run. Will skip execution.")
+
+        # We are marking all running experiments as done here, but later this should be implemented properly
+        # so that an actual check is performed, probably somewhere else.
+        for exp in self.experiments.values():
+            if exp.status == ProfilingExperimentStatus.RUNNING:
+                exp.status = ProfilingExperimentStatus.DONE
+
+    def delete_experiments(
+        self,
+        experiments: list[str] | None = None,
+        all_experiments: bool = False,
+        dry_run: bool = False,
+        remove_repo_dir: bool = False,
+    ) -> None:
+        """Deletes Payu experiments from the work directory and remove them from the manager.
+
+        Args:
+            experiments (list[str] | None): List of experiments (branches) to delete.
+            all_experiments (bool): If True, deletes all experiments managed by this instance.
+            dry_run (bool): If True, performs a dry run without deleting files. Defaults to False.
+            remove_repo_dir (bool): If True, removes the base repository directory if no branches are using it.
+        """
+        # remove_repo_dir would already be forwarded to _delete_experiment via the base class **kwargs, but this
+        # override declares it explicitly so it stays a documented, discoverable and typo-checked argument of the
+        # public Payu API rather than a hidden keyword convention.
+        super().delete_experiments(
+            experiments=experiments,
+            all_experiments=all_experiments,
+            dry_run=dry_run,
+            remove_repo_dir=remove_repo_dir,
+        )
+
+    def _delete_experiment(self, name: str, dry_run: bool, remove_repo_dir: bool = False) -> None:
+        """Deletes a single Payu experiment (branch) via the experiment runner.
+
+        Args:
+            name (str): Name of the experiment (branch) to delete.
+            dry_run (bool): If True, performs a dry run without deleting files.
+            remove_repo_dir (bool): If True, removes the base repository directory if no branches are using it.
+        """
+        runner_config = {
+            "test_path": self.work_dir,
+            "repository_directory": self._repository_directory,
+        }
+
+        runner = ExperimentRunner(runner_config)
+
+        runner.delete_experiments(
+            branches=[name],
+            hard=True,
+            dry_run=dry_run,
+            remove_repo_dir=remove_repo_dir,
+        )
+
+    def archive_experiments(
+        self,
+        exclude_dirs: list[str] | None = None,
+        exclude_files: list[str] | None = None,
+        follow_symlinks: bool = True,
+        overwrite: bool = False,
+    ) -> None:
+        """Archives completed experiments to the specified archive path.
+
+        Args:
+            exclude_dirs (list[str] | None): Directory patterns to exclude when archiving experiments. Defaults to
+                [".git", "restart*"] if not provided.
+            exclude_files (list[str] | None): File patterns to exclude when archiving experiments. Defaults to
+                ["*.nc"] if not provided.
+            follow_symlinks (bool): Whether to follow symlinks when archiving experiments. Defaults to True.
+            overwrite (bool): Whether to overwrite existing archives. Defaults to False.
+        """
+        if exclude_dirs is None:
+            exclude_dirs = [".git", "restart*"]
+        if exclude_files is None:
+            exclude_files = ["*.nc"]
+        super().archive_experiments(
+            exclude_dirs=exclude_dirs, exclude_files=exclude_files, follow_symlinks=follow_symlinks, overwrite=overwrite
+        )
+
+    def parse_ncpus(self, path: Path, run_path: Path | None = None) -> int:
+        """Parses the number of CPUs used in a given Payu experiment.
+
+        Args:
+            path (Path): Path to the Payu experiment directory. Must contain a config.yaml file.
+            run_path (Path | None): Optional path to a separate runs directory. Unused for Payu experiments.
+        Returns:
+            int: Number of CPUs used in the experiment. If multiple submodels are defined, returns the sum of their
+                 ncpus.
+        """
+        config_path = path / "config.yaml"
+        payu_config = YAMLParser().parse(config_path.read_text())
+        if "submodels" in payu_config:
+            return sum(submodel["ncpus"] for submodel in payu_config["submodels"])
+        else:
+            return payu_config["ncpus"]
+
+    def profiling_logs(self, path: Path, run_path: Path | None = None) -> dict[str, dict[int, ProfilingLog]]:
+        """Returns all profiling logs from the specified path.
+
+        Payu can be asked to submit the same experiment several times, in which case each run produces its own
+        output directory and its own telemetry log. Payu numbers both after the same run counter, so the logs of
+        every run are returned, keyed by that number.
+
+        Args:
+            path (Path): Path to the experiment directory.
+            run_path (Path | None): Optional path to a separate runs directory. Unused for Payu experiments.
+        Returns:
+            dict[str, dict[int, ProfilingLog]]: Dictionary mapping log names to their logs, keyed by run number.
+        """
+        logs: dict[str, dict[int, ProfilingLog]] = {}
+
+        # Check archive directory exists
+        archive = path / "archive"
+        if not archive.is_dir():
+            raise FileNotFoundError(f"Directory {archive} does not exist!")
+
+        # Parse payu json profiling data if available. Payu names the directory holding each log after the run number.
+        for json_path in archive.glob("payu_jobs/*/run/*.json"):
+            logs.setdefault("payu", {})[int(json_path.parts[-3])] = ProfilingLog(json_path, PayuJSONProfilingParser())
+
+        # Get the logs of each component of every output directory. Payu names these outputNNN, NNN being the run
+        # number, so output003 holds the same run as payu_jobs/3.
+        output_dirs = sorted(archive.glob("output*"))
+        if not output_dirs:
+            raise FileNotFoundError(f"No output files found in {path}!")
+        for output_dir in output_dirs:
+            run = int(output_dir.name.removeprefix("output"))
+            for name, log in self.get_component_logs(output_dir).items():
+                logs.setdefault(name, {})[run] = log
+
+        return logs

access/profiling/payujson_parser.py

@@ -0,0 +1,75 @@
+# Copyright 2025 ACCESS-NRI and contributors. See the top-level COPYRIGHT file for details.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Parser for payu JSON walltime data generated by payu.
+The data to be parsed is written in the following form:
+
+{
+    "scheduler_job_id": "149764665.gadi-pbs",
+    "scheduler_type": "pbs",
+    # ... many more fields ...
+    "timings": {
+        "payu_start_time": "2025-09-16T08:52:50.748807",
+        "payu_setup_duration_seconds": 47.73822930175811,
+        "payu_model_run_duration_seconds": 6776.044810215011,
+        "payu_run_duration_seconds": 6779.385873348918,
+        "payu_archive_duration_seconds": 8.063649574294686,
+        "payu_finish_time": "2025-09-16T10:46:48.974451",
+        "payu_total_duration_seconds": 6838.225644
+    },
+    # ... more fields
+}
+"""
+
+import json
+import os
+from pathlib import Path
+
+from access.profiling.metrics import tmax
+from access.profiling.parser import ProfilingParser, _read_text_file
+
+
+class PayuJSONProfilingParser(ProfilingParser):
+    """Payu JSON job output profiling parser."""
+
+    _metrics = [tmax]
+
+    def parse(self, file_path: str | Path | os.PathLike) -> dict:
+        """Implements "parse" abstract method to parse a JSON file generated by Payu.
+
+        Args:
+            file_path (str | Path | os.PathLike): file to parse.
+
+        Returns:
+            dict: Parsed timing information.
+
+        Raises:
+            ValueError: when input stream is not valid Payu JSON output.
+        """
+
+        stream = _read_text_file(file_path)
+
+        errmsg = "No Payu profiling data found"
+
+        try:
+            timings = json.loads(stream)["timings"]
+        except Exception as e:
+            raise ValueError(errmsg) from e
+
+        # remove known keys not relevant to profiling
+        for unwanted_key in ("payu_start_time", "payu_finish_time"):
+            if unwanted_key in timings:
+                del timings[unwanted_key]
+
+        # error if no relevant keys in timings
+        if not timings:
+            raise ValueError(errmsg)
+
+        result = {"region": [], tmax: []}
+
+        # transpose dict to be consistent with other profiling parsers.
+        for k, v in timings.items():
+            result["region"].append(k)
+            result[tmax].append(v)
+
+        return result