diff --git a/conda/dev.yml b/conda/dev.yml index 85f64d32..e5a9b2a0 100644 --- a/conda/dev.yml +++ b/conda/dev.yml @@ -23,6 +23,9 @@ dependencies: # ======================= - pytest - pytest-cov + - matplotlib-base + - pandas + - numpy # Documentation # ================= # If versions are updated, also update in `.github/workflows/workflow.yml` diff --git a/tests/README.md b/tests/README.md index c4552471..3c825fac 100644 --- a/tests/README.md +++ b/tests/README.md @@ -109,3 +109,7 @@ GitHub Actions runs the tests according to `.github/workflows/build_workflow.yml python -m unittest tests/integration/python_tests/group_by_command/test_*.py python -m unittest tests/integration/python_tests/group_by_workflow/test_*.py ``` + +## Performance + +For performance profiling, see `tests/performance/README.md`. diff --git a/tests/performance/README.md b/tests/performance/README.md new file mode 100644 index 00000000..aa3670c5 --- /dev/null +++ b/tests/performance/README.md @@ -0,0 +1,100 @@ +# How to profile zstash's performance + +Performance profiling should be done on Perlmutter. We're keeping the performance records in `/global/homes/f/forsyth/zstash_performance_records`. (NOTE: this is currently user-specific. If we start having many other developers running performance profiling, we may try to find a more centralized location.) + +## Generate performance data.bash + +In `zstash/tests/performance/generate_performance_data.bash`, edit the run metadata: + +```bash +# The performance data will end up in `results_csv="${work_dir}${unique_id}/results.csv"` +# Use `/pscratch` since a lot of data will be transferred. +# The results csv alone will be copied to a long-term (i.e., non-scratch) directory at the end. +work_dir=/pscratch/sd/f/forsyth/zstash_performance/ +unique_id=performance_20260603 + +# This is the environment that zstash will be run in. +# Using Unified environment: +environment_commands="source /global/common/software/e3sm/anaconda_envs/load_latest_e3sm_unified_pm-cpu.sh" +# Example dev environment: +environment_commands="source /global/homes/f/forsyth/miniforge3/etc/profile.d/conda.sh; conda activate zstash-pr427-20260603" +``` + +These parameters you probably won't have a need to change: +```bash +# These are the directories to run `zstash create`, `zstash update`, and `zstash extract` on. +dir_to_copy_from=/global/cfs/cdirs/e3sm/forsyth/E3SMv2/v2.LR.historical_0201/ +subdir0=build/ +subdir1=run/ +subdir2=init/ + +# This specifies which `--hpss` settings will be run: +HPSS_OPTIONS=("none" "hpss" "globus") + +# This is what will be used for the "hpss" option: +dst_hpss_path=/home/f/forsyth/zstash_performance + +# This is what will be used the "globus" option: +fresh_globus=true # This will prompt a fresh Globus authentication +dst_endpoint_uuid=15288284-7006-4041-ba1a-6b52501e49f1 # This is LCRC's endpoint +dst_endpoint_archive_dir=/lcrc/group/e3sm/ac.forsyth2/zstash_performance_dst_dir/ +``` + +Once you have the parameters set up, run: +```bash +cd tests/performance/ +./generate_performance_data.bash +``` + +Results will be saved to `${results_csv}` (recall `results_csv="${work_dir}${unique_id}/results.csv`). To keep all records together in a non-scratch space, the results csv is also copied to: `/global/homes/f/forsyth/zstash_performance_records/${unique_id}_results.csv`. + +## Visualize performance + +In `zstash/tests/performance/visualize_performance.py`, edit the run metadata: + +```python +# The results to show in Fig. 1 and Fig. 3 (check). +# This should be the results.csv you just generated in the step above. +RESULTS_CSV: str = "/pscratch/sd/f/forsyth/zstash_performance/performance_20260414/results.csv" + +# The results to compare against in Fig. 2 and (if the baseline also contains +# check data) Fig. 3b. +# Set to None to skip those figures. +# This will typically be the second-to-oldest results.csv in the records space +BASELINE_RESULTS_CSV: Optional[str] = "/pscratch/sd/f/forsyth/zstash_performance/performance_20260402/results.csv" + +# Output path for the saved figures. +# Set to None to display interactively instead of saving. +# Make sure to put this on the web server path, +# i.e., /global/cfs/cdirs/e3sm/www/... +OUTPUT_PATH: Optional[str] = "/global/cfs/cdirs/e3sm/www/forsyth/zstash_performance/performance_pr427_20260603.png" +``` + +Once you have the parameters set up, run: +```bash +cd tests/performance/ +python visualize_performance.py +``` + +The script will print both the file path and the URL to access the plots. + +## For reference + +Records made before the long-term record space was made have been copied to it via: +```bash +SCRATCH_SPACE=/pscratch/sd/f/forsyth/zstash_performance +RECORDS_SPACE=/global/homes/f/forsyth/zstash_performance_records + +for unique_id in \ + performance_20260225 \ + performance_20260226_pr402 \ + performance_20260226_pr424 \ + performance_20260226_pr428 \ + performance_20260402 \ + performance_20260414 \ + performance_pr416_20260403 \ + performance_pr416_20260406 +do + cp "${SCRATCH_SPACE}/${unique_id}/results.csv" "${RECORDS_SPACE}/${unique_id}_results.csv" +done +``` diff --git a/tests/performance/generate_performance_data.bash b/tests/performance/generate_performance_data.bash new file mode 100755 index 00000000..cceb7321 --- /dev/null +++ b/tests/performance/generate_performance_data.bash @@ -0,0 +1,468 @@ +#!/bin/bash +set -e +set -o pipefail + +# Analogous to CI/CD matrix testing of Python versions, +# here we will do a matrix performance profiling +# by comparing runtimes for create/update/extract/check: +# - On multiple directories +# - With `--hpss=none`, with HPSS path, with Globus + +# We will also compare `zstash extract` in sequential-mode and parallel-mode + +############################################################################### +# Manually edit parameters here: + +# Run from Perlmutter, so that we can do both +# a direct transfer to HPSS & a Globus transfer to Chrysalis +work_dir=/pscratch/sd/f/forsyth/zstash_performance/ +unique_id=performance_20260611 +environment_commands="source /global/common/software/e3sm/anaconda_envs/load_latest_e3sm_unified_pm-cpu.sh" + +############################################################################### +# These parameters don't usually need to be changed, +# but can be changed for further customization. + +dir_to_copy_from=/global/cfs/cdirs/e3sm/forsyth/E3SMv2/v2.LR.historical_0201/ +subdir0=build/ +subdir1=run/ +subdir2=init/ +### +# For reference, these files have these sizes and number of files +# (Paths are from Chrysalis, but the data is identical on Perlmutter) + +# Analyzing: /lcrc/group/e3sm/ac.forsyth2/E3SMv2/v2.LR.historical_0201/build/ +# Total size: 1.2GiB +# Number of files: 7046 +# => Lots of small files + +# Analyzing: /lcrc/group/e3sm/ac.forsyth2/E3SMv2/v2.LR.historical_0201/run/ +# Total size: 11GiB +# Number of files: 111 + +# Analyzing: /lcrc/group/e3sm/ac.forsyth2/E3SMv2/v2.LR.historical_0201/init/ +# Total size: 6.9GiB +# Number of files: 14 +# => A few large files +### + + +# For `--hpss=...` +# Which HPSS options to run. Comment out any you want to skip. +# Options: "none" "hpss" "globus" +HPSS_OPTIONS=("none" "hpss" "globus") + +dst_hpss_path=/home/f/forsyth/zstash_performance + +# For `--hpss=globus...` +fresh_globus=true +# ENDPOINT UUIDS: +# LCRC_IMPROV_DTN_ENDPOINT=15288284-7006-4041-ba1a-6b52501e49f1 +# NERSC_PERLMUTTER_ENDPOINT=6bdc7956-fc0f-4ad2-989c-7aa5ee643a79 +# NERSC_HPSS_ENDPOINT=9cd89cfd-6d04-11e5-ba46-22000b92c6ec +# PIC_COMPY_DTN_ENDPOINT=68fbd2fa-83d7-11e9-8e63-029d279f7e24 +# GLOBUS_TUTORIAL_COLLECTION_1_ENDPOINT=6c54cade-bde5-45c1-bdea-f4bd71dba2cc +dst_endpoint_uuid=15288284-7006-4041-ba1a-6b52501e49f1 +dst_endpoint_archive_dir=/lcrc/group/e3sm/ac.forsyth2/zstash_performance_dst_dir/ + +############################################################################### +# Utility functions + +# Colors for output +RED='\033[0;31m' +GREEN='\033[0;32m' +YELLOW='\033[1;33m' +BLUE='\033[0;34m' +CYAN='\033[0;36m' +NC='\033[0m' # No Color + +# Functions to print colored messages +print_step() { + echo -e "${CYAN}[STEP]${NC} $1" +} + +print_info() { + echo -e "${BLUE}[INFO]${NC} $1" +} + +print_success() { + echo -e "${GREEN}[SUCCESS]${NC} $1" +} + +print_error() { + echo -e "${RED}[ERROR]${NC} $1" +} + +print_warning() { + echo -e "${YELLOW}[WARNING]${NC} $1" +} + +confirm() +{ + read -p "$1 (y/n): " -n 1 -r + echo + [[ $REPLY =~ ^[Yy]$ ]] +} + +validate_configuration() +{ + local dir_to_copy_from="${1}" + local subdir0="${2}" + local subdir1="${3}" + local subdir2="${4}" + + print_step "Validating configuration..." + + if [ ! -d "$dir_to_copy_from" ]; then + print_error "Source directory does not exist: $dir_to_copy_from" + exit 1 + fi + + if [ "$subdir0" != "none" ] && [ -n "$subdir0" ]; then + if [ ! -d "${dir_to_copy_from}${subdir0}" ]; then + print_error "subdir0 does not exist: ${dir_to_copy_from}${subdir0}" + exit 1 + fi + fi + + if [ "$subdir1" != "none" ] && [ -n "$subdir1" ]; then + if [ ! -d "${dir_to_copy_from}${subdir1}" ]; then + print_error "subdir1 does not exist: ${dir_to_copy_from}${subdir1}" + exit 1 + fi + fi + + if [ "$subdir2" != "none" ] && [ -n "$subdir2" ]; then + if [ ! -d "${dir_to_copy_from}${subdir2}" ]; then + print_error "subdir2 does not exist: ${dir_to_copy_from}${subdir2}" + exit 1 + fi + fi + + print_success "Configuration validated" +} + +refresh_globus() +{ + print_step "Setting up fresh Globus authentication..." + + # 1. Activate endpoints + echo "Go to https://app.globus.org/file-manager?two_pane=true > For 'Collection', choose the endpoints you're using, and authenticate if needed:" + echo "LCRC Improv DTN, NERSC Perlmutter, NERSC HPSS, pic#compy-dtn" + if ! confirm "Have you authenticated into the correct endpoints?"; then + exit 1 + fi + + # 2. Reset authentication token files + INI_PATH=${HOME}/.zstash.ini + TOKEN_FILE=${HOME}/.zstash_globus_tokens.json + + if [ -f "${INI_PATH}" ]; then + rm -f "${INI_PATH}" + print_info "Removed ${INI_PATH}" + fi + + if [ -f "${TOKEN_FILE}" ]; then + rm -f "${TOKEN_FILE}" + print_info "Removed ${TOKEN_FILE}" + fi + + # 3. Reset Globus consents + echo "https://auth.globus.org/v2/web/consents > Globus Endpoint Performance Monitoring > rescind all" + if ! confirm "Have you revoked Globus consents?"; then + exit 1 + fi + + print_success "Globus authentication reset complete" +} + +# Parse the real-time (wall clock) seconds from the output of `time`. +# `time` writes to stderr a block like: +# real 1m23.456s +# user 0m12.345s +# sys 0m 1.234s +# We capture both stdout+stderr into the log, then grep for the real line. +parse_elapsed_seconds() +{ + local log_file="${1}" + # Extract "Xm Y.ZZZs" and convert to total seconds + awk '/^real/ { + split($2, a, "m"); + mins = a[1]; + secs = substr(a[2], 1, length(a[2])-1); + printf "%.3f\n", mins*60 + secs + }' "${log_file}" +} + +############################################################################### +# Core functions + +run_create() +{ + local dir_to_copy_from="${1}" + local subdir="${2}" + local archive_dir="${3}" + local hpss_path="${4}" + local cache_dir="${5}" + local create_log="${6}" + + print_step "Starting CREATE operation..." + + print_info "Copying data from ${dir_to_copy_from}${subdir}" + cp -r "${dir_to_copy_from}${subdir}" "${archive_dir}${subdir}" + + print_info "Running zstash create..." + print_info "Command: zstash create --hpss=${hpss_path} --cache=${cache_dir} -v ${archive_dir}" + + # We must be outside archive_dir when running create + if { time zstash create --hpss="${hpss_path}" --cache="${cache_dir}" -v "${archive_dir}" ; } 2>&1 | tee "${create_log}"; then + print_success "zstash create completed successfully" + else + print_error "zstash create failed with exit code $?" + exit 1 + fi +} + +run_update() +{ + local dir_to_copy_from="${1}" + local subdir="${2}" + local archive_dir="${3}" + local hpss_path="${4}" + local cache_dir="${5}" + local update_log="${6}" + + print_step "Starting UPDATE operation..." + + print_info "Copying additional data from ${dir_to_copy_from}${subdir}" + cp -r "${dir_to_copy_from}${subdir}" "${archive_dir}${subdir}" + + print_info "Running zstash update..." + print_info "Command: zstash update --hpss=${hpss_path} --cache=${cache_dir} -v" + + # zstash update must be run from within the archive directory + pushd "${archive_dir}" > /dev/null + if { time zstash update --hpss="${hpss_path}" --cache="${cache_dir}" -v ; } 2>&1 | tee "${update_log}"; then + print_success "zstash update completed successfully" + else + print_error "zstash update failed with exit code $?" + popd > /dev/null + exit 1 + fi + popd > /dev/null +} + +run_extract() +{ + local extract_dir="${1}" + local hpss_path="${2}" + local num_workers="${3}" + local cache_dir="${4}" + local extract_log="${5}" + + print_step "Starting EXTRACT operation (workers=${num_workers})..." + + print_info "Running zstash extract..." + print_info "Command: zstash extract --hpss=${hpss_path} --workers=${num_workers} --cache=${cache_dir} -v" + + # zstash extract must be run from within the extraction directory + pushd "${extract_dir}" > /dev/null + if { time zstash extract --hpss="${hpss_path}" --workers="${num_workers}" --cache="${cache_dir}" -v ; } 2>&1 | tee "${extract_log}"; then + print_success "zstash extract completed successfully" + else + print_error "zstash extract failed with exit code $?" + popd > /dev/null + exit 1 + fi + popd > /dev/null +} + +run_check() +{ + local check_dir="${1}" + local hpss_path="${2}" + local num_workers="${3}" + local cache_dir="${4}" + local check_log="${5}" + + print_step "Starting CHECK operation (workers=${num_workers})..." + + print_info "Running zstash check..." + print_info "Command: zstash check --hpss=${hpss_path} --workers=${num_workers} --cache=${cache_dir} -v" + + # zstash check must be run from a new, empty directory (like extract). + # The cache_dir points back to the same archive built by create+update, + # so there is no need to rerun those operations. + pushd "${check_dir}" > /dev/null + if { time zstash check --hpss="${hpss_path}" --workers="${num_workers}" --cache="${cache_dir}" -v ; } 2>&1 | tee "${check_log}"; then + print_success "zstash check completed successfully" + else + print_error "zstash check failed with exit code $?" + popd > /dev/null + exit 1 + fi + popd > /dev/null +} + +############################################################################### +# Results tracking + +# CSV file to collect all runtimes for later visualization +results_csv="${work_dir}${unique_id}/results.csv" + +record_result() +{ + local test_label="${1}" # e.g. "01" + local create_subdir="${2}" + local update_subdir="${3}" + local hpss_label="${4}" # "none", "hpss", "globus" + local operation="${5}" # "create", "update", "extract_seq", "extract_par", "check_seq", "check_par" + local log_file="${6}" + + local elapsed + elapsed=$(parse_elapsed_seconds "${log_file}") + echo "${test_label},${create_subdir},${update_subdir},${hpss_label},${operation},${elapsed}" >> "${results_csv}" + print_info "Recorded: test=${test_label} op=${operation} hpss=${hpss_label} elapsed=${elapsed}s" +} + +############################################################################### +# Main script: + +# Make sure we're running from the correct environment. +# It might not necessarily be a dev environment built off this branch! +if [[ ! "${environment_commands}" =~ ^(source[^;]+)(;[[:space:]]*conda activate[^;]+)?$ ]]; then + print_error "environment_commands must only contain 'source' and optionally 'conda activate'" + exit 1 +fi +eval "${environment_commands}" + +validate_configuration "$dir_to_copy_from" "$subdir0" "$subdir1" "$subdir2" + +if [ "${fresh_globus}" == "true" ] && [[ " ${HPSS_OPTIONS[*]} " == *" globus "* ]]; then + refresh_globus +fi + +# Create the top-level results directory and CSV header +mkdir -p "${work_dir}${unique_id}" +echo "test_label,create_subdir,update_subdir,hpss_label,operation,elapsed_seconds" > "${results_csv}" +print_info "Results CSV: ${results_csv}" + +# Array of subdirectories +subdirs=("$subdir0" "$subdir1" "$subdir2") + +# Define the 6 possible permutations as test configurations. +# Each string contains two space-separated indices into the subdirs array: +# first index = subdir used for create +# second index = subdir used for update +declare -a test_configs=( + "0 1" + "0 2" + "1 0" + "1 2" + "2 0" + "2 1" +) +declare -a test_labels=("01" "02" "10" "12" "20" "21") + +# Loop through the 6 test configurations +for test_idx in 0 1 2 3 4 5; do + # Parse the configuration + read -r -a config <<< "${test_configs[$test_idx]}" + i=${config[0]} # index for create subdir + j=${config[1]} # index for update subdir + + # Get the subdirectories for this test + create_subdir="${subdirs[$i]}" + update_subdir="${subdirs[$j]}" + + test_label="${test_labels[$test_idx]}" + + print_step "==========================================" + print_step "Running Test ${test_label}" + print_step " Create subdir: $create_subdir" + print_step " Update subdir: $update_subdir" + print_step "==========================================" + + # Create unique work directories for this test + work_subdir="${work_dir}${unique_id}/test${test_label}/" + mkdir -p "${work_subdir}" + + log_dir="${work_subdir}logs/" + mkdir -p "${log_dir}" + + dst_globus_path="globus://${dst_endpoint_uuid}/${dst_endpoint_archive_dir}${unique_id}/test${test_label}/" + + # Iterate over the three HPSS modes + declare -A hpss_path_map=( + ["none"]="none" + ["hpss"]="${dst_hpss_path}" + ["globus"]="${dst_globus_path}" + ) + + for hpss_label in "${HPSS_OPTIONS[@]}"; do + hpss_path="${hpss_path_map[$hpss_label]}" + print_step "--- HPSS mode: ${hpss_label} (${hpss_path}) ---" + + # Each hpss mode gets its own subdirectories to avoid cross-contamination + mode_dir="${work_subdir}${hpss_label}/" + archive_dir="${mode_dir}archive_dir/" + cache_dir="${mode_dir}cache/" + mkdir -p "${archive_dir}" "${cache_dir}" + + create_log="${log_dir}create_${hpss_label}.log" + update_log="${log_dir}update_${hpss_label}.log" + + # --- CREATE --- + run_create "$dir_to_copy_from" "$create_subdir" "$archive_dir" "$hpss_path" "$cache_dir" "$create_log" + record_result "$test_label" "$create_subdir" "$update_subdir" "$hpss_label" "create" "$create_log" + + # --- UPDATE --- + run_update "$dir_to_copy_from" "$update_subdir" "$archive_dir" "$hpss_path" "$cache_dir" "$update_log" + record_result "$test_label" "$create_subdir" "$update_subdir" "$hpss_label" "update" "$update_log" + + # --- EXTRACT (sequential=1 worker, parallel=2 workers) --- + for num_workers in 1 2; do + extract_log="${log_dir}extract_${hpss_label}_${num_workers}workers.log" + extract_dir="${mode_dir}extract_${num_workers}workers/" + mkdir -p "${extract_dir}" + + run_extract "$extract_dir" "$hpss_path" "$num_workers" "$cache_dir" "$extract_log" + + if [ "$num_workers" -eq 1 ]; then + op_label="extract_seq" + else + op_label="extract_par" + fi + record_result "$test_label" "$create_subdir" "$update_subdir" "$hpss_label" "$op_label" "$extract_log" + done + + # --- CHECK (sequential=1 worker, parallel=2 workers) --- + # check operates on the same archive as extract; no need to rerun create/update. + # Each worker count gets its own empty directory, as required by zstash check. + for num_workers in 1 2; do + check_log="${log_dir}check_${hpss_label}_${num_workers}workers.log" + check_dir="${mode_dir}check_${num_workers}workers/" + mkdir -p "${check_dir}" + + run_check "$check_dir" "$hpss_path" "$num_workers" "$cache_dir" "$check_log" + + if [ "$num_workers" -eq 1 ]; then + op_label="check_seq" + else + op_label="check_par" + fi + record_result "$test_label" "$create_subdir" "$update_subdir" "$hpss_label" "$op_label" "$check_log" + done + done + + print_success "Test ${test_label} completed" + echo "" +done + +print_success "All tests completed. Results saved to: ${results_csv}" + +performance_archive_path=/global/homes/f/forsyth/zstash_performance_records/${unique_id}_results.csv +cp "${results_csv}" "${performance_archive_path}" +print_success "Results copied to: ${performance_archive_path}" + +print_info "Now edit IO paths and run: python visualize_performance.py" diff --git a/tests/performance/visualize_performance.py b/tests/performance/visualize_performance.py new file mode 100644 index 00000000..1639b043 --- /dev/null +++ b/tests/performance/visualize_performance.py @@ -0,0 +1,1365 @@ +#!/usr/bin/env python3 +""" +visualize_performance.py – Plot zstash performance profiling results. + +Usage: + python visualize_performance.py + +Edit the constants at the top of this file to point at the CSV(s) to plot. + +The CSV is produced by generate_performance_data.bash and has columns: + test_label, create_subdir, update_subdir, hpss_label, operation, elapsed_seconds + +Visualization strategy +---------------------- +Four dimensions: + 1. Operation : create | update | extract_seq | extract_par | check_seq | check_par + 2. Directory : build/ (many small) | run/ (medium) | init/ (few large) + 3. HPSS mode : none | hpss | globus + 4. Parallelism: already encoded in operation (extract_seq vs extract_par, etc.) + +Figure 1 – Performance overview: + Layout: 2×2 grid of subplots, one per operation. + Within each subplot: + - X-axis groups = directory processed (create_subdir or update_subdir) + for create/update; or (create_subdir, update_subdir) + archive config for extract_seq/extract_par. + - Bars = HPSS mode (none / hpss / globus), colour-coded + - Each test config contributes one bar per (directory, hpss_mode) cell; + if multiple configs share the same directory for an operation, their + runtimes are shown as individual dots and the bar shows the mean. + An additional 5th subplot compares extract_seq vs extract_par side-by-side + to make the parallelism speed-up immediately visible. + +Figure 2 – Baseline comparison (current branch vs main): + Produced only when BASELINE_RESULTS_CSV is set to a valid path. + Same 2×2 + comparison layout, but each cell shows two bars + (current = solid, baseline = hatched) with a ratio annotation + (current/baseline) above each pair. Ratio > 1 = regression (slower), + ratio < 1 = improvement (faster). + +Figure 3 – zstash check vs extract: + Always produced when check data is present in the CSV. + Layout: 3 rows × 2 cols + Row 0: check_seq | check_par (standalone check performance) + Row 1: check_seq vs extract_seq (direct apples-to-apples comparison) + Row 2: check_par vs extract_par (same for parallel mode) + Since check is essentially a dry run of extract (it downloads tars and + verifies md5 checksums but does not write extracted files to disk), these + plots make any overhead difference immediately visible. + If BASELINE_RESULTS_CSV is set, a Fig. 3b is also produced using the same + current-vs-baseline pairing as Fig. 2. +""" + +import argparse +import os +import sys +from pathlib import Path +from typing import Optional + +import matplotlib.patches as mpatches +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd + +# --------------------------------------------------------------------------- +# ← EDIT THESE for each new run +# --------------------------------------------------------------------------- + +# The results to show in Fig. 1 and Fig. 3 +RESULTS_CSV: str = ( + "/pscratch/sd/f/forsyth/zstash_performance/performance_20260611/results.csv" +) + +# The results to compare against in Fig. 2 and Fig. 3b. +# Set to None to skip those figures. +BASELINE_RESULTS_CSV: Optional[str] = None + +# Output path for the saved figures. +# Set to None to display interactively instead of saving. +# Fig. 2 and Fig. 3 paths are derived automatically from this path. +OUTPUT_PATH: Optional[str] = ( + "/global/cfs/cdirs/e3sm/www/forsyth/zstash_performance/performance_20260611.png" +) + +# --------------------------------------------------------------------------- +# Config +# --------------------------------------------------------------------------- + +HPSS_ORDER = ["none", "hpss", "globus"] +HPSS_COLORS = {"none": "#4C72B0", "hpss": "#DD8452", "globus": "#55A868"} +HPSS_LABELS = {"none": "No HPSS", "hpss": "Direct HPSS", "globus": "Globus"} + +OP_ORDER = ["create", "update", "extract_seq", "extract_par"] +OP_TITLES = { + "create": "zstash create", + "update": "zstash update", + "extract_seq": "zstash extract (sequential, 1 worker)", + "extract_par": "zstash extract (parallel, 2 workers)", + "check_seq": "zstash check (sequential, 1 worker)", + "check_par": "zstash check (parallel, 2 workers)", +} + +# Map an operation to the column that holds the "relevant directory". +# Extract and check are intentionally absent: they operate on the combined +# create+update archive, so both subdirs are needed and are handled separately. +OP_DIR_COL = { + "create": "create_subdir", + "update": "update_subdir", +} + +BAR_WIDTH = 0.22 +DOT_ALPHA = 0.55 +DOT_SIZE = 40 + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def load_data(csv_path: str) -> pd.DataFrame: + df = pd.read_csv(csv_path) + df.columns = df.columns.str.strip() + df["elapsed_seconds"] = pd.to_numeric(df["elapsed_seconds"], errors="coerce") + # Normalise subdir names: strip trailing slashes for display + for col in ("create_subdir", "update_subdir"): + df[col] = df[col].str.strip().str.rstrip("/") + df["hpss_label"] = df["hpss_label"].str.strip() + df["operation"] = df["operation"].str.strip() + return df + + +def dir_sort_key(name: str) -> int: + """Sort dirs in a consistent order: build, run, init.""" + order = {"build": 0, "run": 1, "init": 2} + return order.get(name.lower(), 99) + + +def _add_dir_annotation(ax, dirs, x_positions): + """ + Add a small file-count hint below each directory group label. + + Parameters + ---------- + ax : the Axes to annotate + dirs : list of directory names in display order + x_positions : list of x-axis data coordinates for each dir group centre. + These are passed in explicitly so the function works for both + Fig. 1 (groups at 0, 1, 2, …) and Fig. 2 (wider group_span). + """ + hints = { + "build": "many small files\n(~7k files, 1.2 GiB)", + "run": "mixed\n(~111 files, 11 GiB)", + "init": "few large files\n(14 files, 6.9 GiB)", + } + for x_centre, d in zip(x_positions, dirs): + if d in hints: + ax.annotate( + hints[d], + xy=(x_centre, 0), + xycoords=("data", "axes fraction"), + xytext=(0, -46), + textcoords="offset points", + ha="center", + va="top", + fontsize=6.5, + color="#555555", + annotation_clip=False, + ) + + +def plot_operation(ax, df_op: pd.DataFrame, operation: str, dirs: list[str]): + """Draw grouped bars for one operation subplot.""" + dir_col = OP_DIR_COL[operation] + n_dirs = len(dirs) + n_hpss = len(HPSS_ORDER) + + x_base = np.arange(n_dirs) + offsets = np.linspace(-(n_hpss - 1) / 2, (n_hpss - 1) / 2, n_hpss) * BAR_WIDTH + + for h_idx, hpss in enumerate(HPSS_ORDER): + df_h = df_op[df_op["hpss_label"] == hpss] + means, all_vals, xs = [], [], [] + + for d_idx, d in enumerate(dirs): + vals = df_h[df_h[dir_col] == d]["elapsed_seconds"].dropna().values + mean = vals.mean() if len(vals) > 0 else 0.0 + means.append(mean) + all_vals.append(vals) + xs.append(x_base[d_idx] + offsets[h_idx]) + + color = HPSS_COLORS[hpss] + ax.bar( + xs, + means, + width=BAR_WIDTH, + color=color, + alpha=0.85, + label=HPSS_LABELS[hpss], + zorder=2, + ) + # Overlay individual data points so scatter is visible + for x_pos, vals in zip(xs, all_vals): + if len(vals) > 1: + jitter = np.random.uniform( + -BAR_WIDTH * 0.25, BAR_WIDTH * 0.25, size=len(vals) + ) + ax.scatter( + x_pos + jitter, + vals, + color="white", + edgecolors=color, + s=DOT_SIZE, + zorder=3, + alpha=DOT_ALPHA, + linewidths=1.2, + ) + + ax.set_title(OP_TITLES[operation], fontsize=10, fontweight="bold", pad=6) + ax.set_xticks(x_base) + ax.set_xticklabels([d + "/" for d in dirs], fontsize=9) + ax.set_ylabel("Wall-clock time (s)", fontsize=8) + ax.set_xlabel("Directory processed", fontsize=8, labelpad=14) + ax.yaxis.grid(True, linestyle="--", alpha=0.5, zorder=0) + ax.set_axisbelow(True) + # Fig. 1: group centres are at integer positions 0, 1, 2, … + _add_dir_annotation(ax, dirs, list(x_base)) + + # Value labels on bars + for rect in ax.patches: + h = rect.get_height() + if h > 0: + ax.text( + rect.get_x() + rect.get_width() / 2, + h * 1.01, + f"{h:.0f}s", + ha="center", + va="bottom", + fontsize=6, + color="#333333", + ) + + +def _extract_configs(df: pd.DataFrame) -> list[tuple[str, str]]: + """ + Return the sorted list of (create_subdir, update_subdir) pairs that + actually appear in the extract rows of *df*. These represent the + combined archives that were extracted from. + """ + mask = df["operation"].isin(["extract_seq", "extract_par"]) + pairs = ( + df[mask][["create_subdir", "update_subdir"]] + .drop_duplicates() + .apply(tuple, axis=1) + .tolist() + ) + # Sort by create_subdir first, then update_subdir + return sorted(pairs, key=lambda p: (dir_sort_key(p[0]), dir_sort_key(p[1]))) + + +def _check_configs(df: pd.DataFrame) -> list[tuple[str, str]]: + """ + Return the sorted list of (create_subdir, update_subdir) pairs that + appear in check rows of *df*. + """ + mask = df["operation"].isin(["check_seq", "check_par"]) + pairs = ( + df[mask][["create_subdir", "update_subdir"]] + .drop_duplicates() + .apply(tuple, axis=1) + .tolist() + ) + return sorted(pairs, key=lambda p: (dir_sort_key(p[0]), dir_sort_key(p[1]))) + + +def _extract_tick_label(create_sub: str, update_sub: str) -> str: + """Short two-line tick label for a (create, update) archive config.""" + return f"create: {create_sub}/\nupdate: {update_sub}/" + + +def _plot_extract_single_op(ax, df: pd.DataFrame, operation: str): + """ + Draw grouped bars for one extract operation (extract_seq or extract_par). + X-axis groups are the combined (create_subdir, update_subdir) archive + configs, since extract operates on the archive built by both operations. + """ + configs = _extract_configs(df) + n_configs = len(configs) + n_hpss = len(HPSS_ORDER) + + x_base = np.arange(n_configs, dtype=float) + offsets = np.linspace(-(n_hpss - 1) / 2, (n_hpss - 1) / 2, n_hpss) * BAR_WIDTH + + for h_idx, hpss in enumerate(HPSS_ORDER): + means, all_vals, xs = [], [], [] + for c_idx, (create_sub, update_sub) in enumerate(configs): + vals = ( + df[ + (df["operation"] == operation) + & (df["hpss_label"] == hpss) + & (df["create_subdir"] == create_sub) + & (df["update_subdir"] == update_sub) + ]["elapsed_seconds"] + .dropna() + .values + ) + mean = vals.mean() if len(vals) > 0 else 0.0 + means.append(mean) + all_vals.append(vals) + xs.append(x_base[c_idx] + offsets[h_idx]) + + color = HPSS_COLORS[hpss] + ax.bar( + xs, + means, + width=BAR_WIDTH, + color=color, + alpha=0.85, + label=HPSS_LABELS[hpss], + zorder=2, + ) + for x_pos, vals in zip(xs, all_vals): + if len(vals) > 1: + jitter = np.random.uniform( + -BAR_WIDTH * 0.25, BAR_WIDTH * 0.25, size=len(vals) + ) + ax.scatter( + x_pos + jitter, + vals, + color="white", + edgecolors=color, + s=DOT_SIZE, + zorder=3, + alpha=DOT_ALPHA, + linewidths=1.2, + ) + + ax.set_title(OP_TITLES[operation], fontsize=10, fontweight="bold", pad=6) + ax.set_xticks(x_base) + ax.set_xticklabels([_extract_tick_label(c, u) for c, u in configs], fontsize=7) + ax.set_ylabel("Wall-clock time (s)", fontsize=8) + ax.set_xlabel("Archive contents (create → update)", fontsize=8, labelpad=6) + ax.yaxis.grid(True, linestyle="--", alpha=0.5, zorder=0) + ax.set_axisbelow(True) + + for rect in ax.patches: + h = rect.get_height() + if h > 0: + ax.text( + rect.get_x() + rect.get_width() / 2, + h * 1.01, + f"{h:.0f}s", + ha="center", + va="bottom", + fontsize=6, + color="#333333", + ) + + +def plot_extract_comparison(ax, df: pd.DataFrame): + """ + Extra subplot: sequential vs parallel extract, grouped by + (archive config, hpss). Each archive config is the *combined* + create+update directory pair, since extract operates on the full + archive assembled by both operations. + Uses a hatch pattern to distinguish seq/par within each hpss colour. + """ + configs = _extract_configs(df) + n_configs = len(configs) + ops = ["extract_seq", "extract_par"] + hatches = {"extract_seq": "", "extract_par": "////"} + n_bars = len(HPSS_ORDER) * len(ops) + + group_width = n_bars * BAR_WIDTH + 0.15 # total width per config group + x_base = np.arange(n_configs) * group_width + + for c_idx, (create_sub, update_sub) in enumerate(configs): + for h_idx, hpss in enumerate(HPSS_ORDER): + for op_idx, op in enumerate(ops): + df_cell = df[ + (df["operation"] == op) + & (df["hpss_label"] == hpss) + & (df["create_subdir"] == create_sub) + & (df["update_subdir"] == update_sub) + ] + vals = df_cell["elapsed_seconds"].dropna().values + mean = vals.mean() if len(vals) > 0 else 0.0 + bar_x = x_base[c_idx] + (h_idx * len(ops) + op_idx) * BAR_WIDTH + ax.bar( + bar_x, + mean, + width=BAR_WIDTH, + color=HPSS_COLORS[hpss], + hatch=hatches[op], + alpha=0.85, + zorder=2, + ) + + tick_positions = x_base + (n_bars / 2 - 0.5) * BAR_WIDTH + ax.set_xticks(tick_positions) + ax.set_xticklabels([_extract_tick_label(c, u) for c, u in configs], fontsize=7.5) + ax.set_ylabel("Wall-clock time (s)", fontsize=8) + ax.set_xlabel( + "Archive contents (create subdir → update subdir)", fontsize=8, labelpad=14 + ) + ax.set_title( + "Extract: Sequential vs Parallel (speed-up comparison)\n" + "Each group = archive built from create subdir + update subdir", + fontsize=10, + fontweight="bold", + pad=6, + ) + ax.yaxis.grid(True, linestyle="--", alpha=0.5, zorder=0) + ax.set_axisbelow(True) + + # Custom legend: colour = hpss, hatch = seq/par + hpss_patches = [ + mpatches.Patch(color=HPSS_COLORS[h], label=HPSS_LABELS[h]) for h in HPSS_ORDER + ] + seq_patch = mpatches.Patch( + facecolor="grey", hatch="", label="Sequential (1 worker)" + ) + par_patch = mpatches.Patch( + facecolor="grey", hatch="////", label="Parallel (2 workers)" + ) + ax.legend( + handles=hpss_patches + [seq_patch, par_patch], + fontsize=7, + loc="upper right", + ncol=2, + ) + + +# --------------------------------------------------------------------------- +# Baseline comparison figure +# --------------------------------------------------------------------------- + +# Ratio colouring thresholds +RATIO_REGRESSION = 1.10 # >= 10% slower → red +RATIO_IMPROVEMENT = 0.90 # <= 10% faster → green +RATIO_NEUTRAL_COLOR = "#333333" +RATIO_REGRESSION_COLOR = "#CC3311" +RATIO_IMPROVEMENT_COLOR = "#228833" + + +def _ratio_color(ratio: float) -> str: + if ratio >= RATIO_REGRESSION: + return RATIO_REGRESSION_COLOR + if ratio <= RATIO_IMPROVEMENT: + return RATIO_IMPROVEMENT_COLOR + return RATIO_NEUTRAL_COLOR + + +def plot_comparison_operation( + ax, + df_cur: pd.DataFrame, + df_bas: pd.DataFrame, + operation: str, + dirs: list[str], +): + """ + For one operation, draw paired bars (current vs baseline) per + (directory, hpss_mode) cell, with a ratio annotation above each pair. + """ + dir_col = OP_DIR_COL[operation] + n_dirs = len(dirs) + n_hpss = len(HPSS_ORDER) + + # Each hpss group occupies 2 bars (current + baseline) + a small gap + pair_width = BAR_WIDTH + gap = BAR_WIDTH * 0.3 + group_span = n_hpss * (2 * pair_width + gap) + 0.2 + x_base = np.arange(n_dirs) * group_span + + for h_idx, hpss in enumerate(HPSS_ORDER): + color = HPSS_COLORS[hpss] + pair_offset = h_idx * (2 * pair_width + gap) + + for d_idx, d in enumerate(dirs): + x_left = x_base[d_idx] + pair_offset # baseline bar + x_right = x_base[d_idx] + pair_offset + pair_width # current bar + + def mean_for(df, _op=operation, _h=hpss, _d=d): + v = ( + df[ + (df["operation"] == _op) + & (df["hpss_label"] == _h) + & (df[dir_col] == _d) + ]["elapsed_seconds"] + .dropna() + .values + ) + return v.mean() if len(v) > 0 else 0.0 + + cur_mean = mean_for(df_cur) + bas_mean = mean_for(df_bas) + + # Baseline bar (hatched, lighter) — left + ax.bar( + x_left, + bas_mean, + width=pair_width, + color=color, + alpha=0.40, + hatch="////", + zorder=2, + edgecolor=color, + ) + # Current bar (solid) — right + ax.bar( + x_right, + cur_mean, + width=pair_width, + color=color, + alpha=0.85, + zorder=2, + label=HPSS_LABELS[hpss] if d_idx == 0 else "", + ) + + # Ratio annotation + if bas_mean > 0 and cur_mean > 0: + ratio = cur_mean / bas_mean + top = max(cur_mean, bas_mean) + rat_color = _ratio_color(ratio) + arrow = ( + "▲" + if ratio >= RATIO_REGRESSION + else ("▼" if ratio <= RATIO_IMPROVEMENT else "") + ) + ax.text( + (x_left + x_right) / 2, + top * 1.03, + f"{arrow}{ratio:.2f}×", + ha="center", + va="bottom", + fontsize=6.5, + fontweight="bold", + color=rat_color, + zorder=4, + ) + + ax.set_title(OP_TITLES[operation], fontsize=10, fontweight="bold", pad=6) + # Tick at the centre of each directory's group of bars + group_centre_offset = (n_hpss * (2 * pair_width + gap) - gap) / 2 + x_ticks = x_base + group_centre_offset + ax.set_xticks(x_ticks) + ax.set_xticklabels([d + "/" for d in dirs], fontsize=9) + ax.set_ylabel("Wall-clock time (s)", fontsize=8) + ax.set_xlabel("Directory processed", fontsize=8, labelpad=14) + ax.yaxis.grid(True, linestyle="--", alpha=0.5, zorder=0) + ax.set_axisbelow(True) + # Pass actual tick x-positions so annotations align with tick labels + _add_dir_annotation(ax, dirs, list(x_ticks)) + + +def _plot_comparison_extract_single_op( + ax, + df_cur: pd.DataFrame, + df_bas: pd.DataFrame, + operation: str, +): + """ + Fig. 2 version of a single extract-op subplot (extract_seq or extract_par). + X-axis = combined (create, update) archive config; bars = current vs baseline + paired within each HPSS group. + """ + configs = _extract_configs(df_cur) + n_configs = len(configs) + n_hpss = len(HPSS_ORDER) + + pair_width = BAR_WIDTH + gap = BAR_WIDTH * 0.3 + group_span = n_hpss * (2 * pair_width + gap) + 0.2 + x_base = np.arange(n_configs) * group_span + + for h_idx, hpss in enumerate(HPSS_ORDER): + color = HPSS_COLORS[hpss] + pair_offset = h_idx * (2 * pair_width + gap) + for c_idx, (create_sub, update_sub) in enumerate(configs): + x_left = x_base[c_idx] + pair_offset + x_right = x_left + pair_width + + def mean_for(df, _op=operation, _h=hpss, _cs=create_sub, _us=update_sub): + v = ( + df[ + (df["operation"] == _op) + & (df["hpss_label"] == _h) + & (df["create_subdir"] == _cs) + & (df["update_subdir"] == _us) + ]["elapsed_seconds"] + .dropna() + .values + ) + return v.mean() if len(v) > 0 else 0.0 + + cur_mean = mean_for(df_cur) + bas_mean = mean_for(df_bas) + + # Baseline bar (hatched, lighter) — left + ax.bar( + x_left, + bas_mean, + width=pair_width, + color=color, + alpha=0.40, + hatch="////", + zorder=2, + edgecolor=color, + ) + # Current bar (solid) — right + ax.bar( + x_right, + cur_mean, + width=pair_width, + color=color, + alpha=0.85, + zorder=2, + label=HPSS_LABELS[hpss] if c_idx == 0 else "", + ) + + if bas_mean > 0 and cur_mean > 0: + ratio = cur_mean / bas_mean + top = max(cur_mean, bas_mean) + arrow = ( + "▲" + if ratio >= RATIO_REGRESSION + else ("▼" if ratio <= RATIO_IMPROVEMENT else "") + ) + ax.text( + (x_left + x_right) / 2, + top * 1.03, + f"{arrow}{ratio:.2f}×", + ha="center", + va="bottom", + fontsize=6.5, + fontweight="bold", + color=_ratio_color(ratio), + zorder=4, + ) + + group_centre_offset = (n_hpss * (2 * pair_width + gap) - gap) / 2 + x_ticks = x_base + group_centre_offset + ax.set_xticks(x_ticks) + ax.set_xticklabels([_extract_tick_label(c, u) for c, u in configs], fontsize=7) + ax.set_ylabel("Wall-clock time (s)", fontsize=8) + ax.set_xlabel("Archive contents (create → update)", fontsize=8, labelpad=6) + ax.set_title(OP_TITLES[operation], fontsize=10, fontweight="bold", pad=6) + ax.yaxis.grid(True, linestyle="--", alpha=0.5, zorder=0) + ax.set_axisbelow(True) + + +def plot_comparison_extract( + ax, + df_cur: pd.DataFrame, + df_bas: pd.DataFrame, +): + """ + Seq vs par extract comparison with current/baseline pairing, grouped by + the combined (create_subdir, update_subdir) archive config. + + Bar order within each HPSS × op cell (innermost grouping): + [current/seq] [baseline/seq] ‹op_gap› [current/par] [baseline/par] + """ + configs = _extract_configs(df_cur) + n_configs = len(configs) + ops = ["extract_seq", "extract_par"] + op_hatches = {"extract_seq": "", "extract_par": "xxxx"} + + pair_width = BAR_WIDTH + inner_gap = BAR_WIDTH * 0.15 # gap between current/baseline within a pair + op_gap = BAR_WIDTH * 0.55 # larger gap between seq-pair and par-pair + hpss_gap = BAR_WIDTH * 0.30 # gap between HPSS groups + + pair_span = 2 * pair_width + inner_gap + hpss_group_span = 2 * pair_span + op_gap + + group_span = len(HPSS_ORDER) * (hpss_group_span + hpss_gap) + 0.3 + x_base = np.arange(n_configs) * group_span + + for c_idx, (create_sub, update_sub) in enumerate(configs): + for h_idx, hpss in enumerate(HPSS_ORDER): + color = HPSS_COLORS[hpss] + hpss_origin = x_base[c_idx] + h_idx * (hpss_group_span + hpss_gap) + for op_idx, op in enumerate(ops): + hatch = op_hatches[op] + op_origin = hpss_origin + op_idx * (pair_span + op_gap) + x_bas = op_origin # baseline — left + x_cur = op_origin + pair_width + inner_gap # current — right + + def mean_for(df, _op=op, _h=hpss, _cs=create_sub, _us=update_sub): + v = ( + df[ + (df["operation"] == _op) + & (df["hpss_label"] == _h) + & (df["create_subdir"] == _cs) + & (df["update_subdir"] == _us) + ]["elapsed_seconds"] + .dropna() + .values + ) + return v.mean() if len(v) > 0 else 0.0 + + cur_mean = mean_for(df_cur) + bas_mean = mean_for(df_bas) + + # Baseline bar (left): op-hatch + //// to mark it as baseline + bas_hatch = hatch + "////" + ax.bar( + x_bas, + bas_mean, + width=pair_width, + color=color, + hatch=bas_hatch, + alpha=0.35, + zorder=2, + edgecolor=color, + ) + # Current bar (right): op-hatch only + ax.bar( + x_cur, + cur_mean, + width=pair_width, + color=color, + hatch=hatch, + alpha=0.85, + zorder=2, + ) + + if bas_mean > 0 and cur_mean > 0: + ratio = cur_mean / bas_mean + top = max(cur_mean, bas_mean) + arrow = ( + "▲" + if ratio >= RATIO_REGRESSION + else ("▼" if ratio <= RATIO_IMPROVEMENT else "") + ) + ax.text( + (x_cur + x_bas) / 2 + pair_width / 2, + top * 1.03, + f"{arrow}{ratio:.2f}×", + ha="center", + va="bottom", + fontsize=5.5, + fontweight="bold", + color=_ratio_color(ratio), + zorder=4, + ) + + group_total_bar_span = len(HPSS_ORDER) * (hpss_group_span + hpss_gap) - hpss_gap + x_ticks = x_base + group_total_bar_span / 2 + ax.set_xticks(x_ticks) + ax.set_xticklabels([_extract_tick_label(c, u) for c, u in configs], fontsize=7.5) + ax.set_ylabel("Wall-clock time (s)", fontsize=8) + ax.set_xlabel( + "Archive contents (create subdir → update subdir)", fontsize=8, labelpad=14 + ) + ax.set_title( + "Extract: Sequential vs Parallel — current vs baseline\n" + "Each group = archive built from create subdir + update subdir", + fontsize=10, + fontweight="bold", + pad=6, + ) + ax.yaxis.grid(True, linestyle="--", alpha=0.5, zorder=0) + ax.set_axisbelow(True) + + # Legend: colour=hpss, hatch=seq/par, alpha=current/baseline + hpss_patches = [ + mpatches.Patch(color=HPSS_COLORS[h], label=HPSS_LABELS[h]) for h in HPSS_ORDER + ] + seq_patch = mpatches.Patch( + facecolor="grey", hatch="", alpha=0.85, label="Sequential, current" + ) + seq_bas_patch = mpatches.Patch( + facecolor="grey", hatch="////", alpha=0.35, label="Sequential, baseline" + ) + par_patch = mpatches.Patch( + facecolor="grey", hatch="xxxx", alpha=0.85, label="Parallel, current" + ) + par_bas_patch = mpatches.Patch( + facecolor="grey", hatch="xxxx////", alpha=0.35, label="Parallel, baseline" + ) + ax.legend( + handles=hpss_patches + [seq_patch, seq_bas_patch, par_patch, par_bas_patch], + fontsize=6.5, + loc="upper right", + ncol=3, + ) + + +def build_comparison_figure( + df_cur: pd.DataFrame, + df_bas: pd.DataFrame, + all_dirs: list[str], + cur_label: str, + bas_label: str, +) -> plt.Figure: + """Build and return the full baseline-comparison figure.""" + fig = plt.figure(figsize=(16, 17)) + fig.suptitle( + f"zstash Performance: Current vs Baseline\n" + f"current = {cur_label} | baseline (main) = {bas_label}\n" + f"Ratio = current / baseline — " + f"▲ {RATIO_REGRESSION_COLOR_LABEL} ≥{RATIO_REGRESSION:.0%} slower " + f"▼ {RATIO_IMPROVEMENT_COLOR_LABEL} ≤{RATIO_IMPROVEMENT:.0%} faster " + f"= within ±10%", + fontsize=11, + fontweight="bold", + y=0.98, + ) + + gs = fig.add_gridspec( + 3, 2, hspace=0.58, wspace=0.35, top=0.92, bottom=0.07, left=0.07, right=0.97 + ) + + axes = { + "create": fig.add_subplot(gs[0, 0]), + "update": fig.add_subplot(gs[0, 1]), + "extract_seq": fig.add_subplot(gs[1, 0]), + "extract_par": fig.add_subplot(gs[1, 1]), + } + ax_cmp = fig.add_subplot(gs[2, :]) + + for op in OP_ORDER: + if op in OP_DIR_COL: + plot_comparison_operation(axes[op], df_cur, df_bas, op, all_dirs) + else: + _plot_comparison_extract_single_op(axes[op], df_cur, df_bas, op) + + # Shared legend for solid=current, hatched=baseline + cur_patch = mpatches.Patch(facecolor="grey", alpha=0.85, label="Current branch") + bas_patch = mpatches.Patch( + facecolor="grey", alpha=0.40, hatch="////", label="Baseline (main)" + ) + hpss_patches = [ + mpatches.Patch(color=HPSS_COLORS[h], label=HPSS_LABELS[h]) for h in HPSS_ORDER + ] + axes["create"].legend( + handles=[cur_patch, bas_patch] + hpss_patches, + fontsize=7, + loc="upper right", + ) + + plot_comparison_extract(ax_cmp, df_cur, df_bas) + return fig + + +# --------------------------------------------------------------------------- +# Figure 3 – zstash check vs extract +# --------------------------------------------------------------------------- + + +def _plot_check_vs_extract_pair( + ax, + df: pd.DataFrame, + check_op: str, + extract_op: str, +): + """ + Draw a grouped-bar subplot comparing check and extract for the same + worker count. X-axis = (create, update) archive configs. + Within each config group, bars are ordered: [check, extract] × HPSS mode, + distinguished by hatch (check = "////", extract = ""). + + This makes the overhead (or savings) of check relative to extract + immediately visible, since check is conceptually a dry-run of extract. + """ + configs = _check_configs(df) + if not configs: + ax.set_visible(False) + return + + n_configs = len(configs) + n_hpss = len(HPSS_ORDER) + ops = [check_op, extract_op] + op_hatches = {check_op: "////", extract_op: ""} + + n_bars = n_hpss * len(ops) + group_width = n_bars * BAR_WIDTH + 0.2 + x_base = np.arange(n_configs) * group_width + + for c_idx, (create_sub, update_sub) in enumerate(configs): + for h_idx, hpss in enumerate(HPSS_ORDER): + color = HPSS_COLORS[hpss] + for op_idx, op in enumerate(ops): + bar_x = x_base[c_idx] + (h_idx * len(ops) + op_idx) * BAR_WIDTH + vals = ( + df[ + (df["operation"] == op) + & (df["hpss_label"] == hpss) + & (df["create_subdir"] == create_sub) + & (df["update_subdir"] == update_sub) + ]["elapsed_seconds"] + .dropna() + .values + ) + mean = vals.mean() if len(vals) > 0 else 0.0 + ax.bar( + bar_x, + mean, + width=BAR_WIDTH, + color=color, + hatch=op_hatches[op], + alpha=0.85, + zorder=2, + ) + if mean > 0: + ax.text( + bar_x + BAR_WIDTH / 2, + mean * 1.01, + f"{mean:.0f}s", + ha="center", + va="bottom", + fontsize=5.5, + color="#333333", + ) + + tick_positions = x_base + (n_bars / 2 - 0.5) * BAR_WIDTH + ax.set_xticks(tick_positions) + ax.set_xticklabels([_extract_tick_label(c, u) for c, u in configs], fontsize=7) + ax.set_ylabel("Wall-clock time (s)", fontsize=8) + ax.set_xlabel("Archive contents (create → update)", fontsize=8, labelpad=6) + + workers = "1 worker" if check_op == "check_seq" else "2 workers" + ax.set_title( + f"check vs extract ({workers})\n" f"Hatch = check / Solid = extract", + fontsize=10, + fontweight="bold", + pad=6, + ) + ax.yaxis.grid(True, linestyle="--", alpha=0.5, zorder=0) + ax.set_axisbelow(True) + + hpss_patches = [ + mpatches.Patch(color=HPSS_COLORS[h], label=HPSS_LABELS[h]) for h in HPSS_ORDER + ] + check_patch = mpatches.Patch(facecolor="grey", hatch="////", label="check") + extract_patch = mpatches.Patch(facecolor="grey", hatch="", label="extract") + ax.legend( + handles=hpss_patches + [check_patch, extract_patch], + fontsize=7, + loc="upper right", + ncol=2, + ) + + +def build_check_figure(df: pd.DataFrame) -> Optional[plt.Figure]: + """ + Build Figure 3: zstash check standalone and check-vs-extract comparison. + + Layout (3 rows × 2 cols): + Row 0: check_seq (standalone) | check_par (standalone) + Row 1: check_seq vs extract_seq comparison + Row 2: check_par vs extract_par comparison + """ + if not df["operation"].isin(["check_seq", "check_par"]).any(): + return None + + fig = plt.figure(figsize=(15, 16)) + fig.suptitle( + "zstash check Performance\n" + "Top row: check standalone; " + "bottom rows: check vs extract (check ≈ dry-run extract)", + fontsize=13, + fontweight="bold", + y=0.98, + ) + + gs = fig.add_gridspec( + 3, 2, hspace=0.55, wspace=0.35, top=0.93, bottom=0.07, left=0.07, right=0.97 + ) + + # Row 0: standalone check subplots + ax_check_seq = fig.add_subplot(gs[0, 0]) + ax_check_par = fig.add_subplot(gs[0, 1]) + _plot_extract_single_op(ax_check_seq, df, "check_seq") + _plot_extract_single_op(ax_check_par, df, "check_par") + + # Row 1: check_seq vs extract_seq + ax_cmp_seq = fig.add_subplot(gs[1, :]) + _plot_check_vs_extract_pair(ax_cmp_seq, df, "check_seq", "extract_seq") + + # Row 2: check_par vs extract_par + ax_cmp_par = fig.add_subplot(gs[2, :]) + _plot_check_vs_extract_pair(ax_cmp_par, df, "check_par", "extract_par") + + # Legend for the standalone row + legend_handles = [ + mpatches.Patch(color=HPSS_COLORS[h], label=HPSS_LABELS[h]) for h in HPSS_ORDER + ] + ax_check_seq.legend(handles=legend_handles, fontsize=7, loc="upper right") + + return fig + + +def _plot_check_vs_extract_pair_comparison( + ax, + df_cur: pd.DataFrame, + df_bas: pd.DataFrame, + check_op: str, + extract_op: str, +): + """ + Fig. 3b version of the check-vs-extract subplot, with current/baseline + pairing overlaid. + + Bar order (innermost, per HPSS group): + [bas/check] [cur/check] ‹op_gap› [bas/extract] [cur/extract] + """ + configs = _check_configs(df_cur) + if not configs: + ax.set_visible(False) + return + + n_configs = len(configs) + ops = [check_op, extract_op] + op_hatches = {check_op: "////", extract_op: ""} + + pair_width = BAR_WIDTH + inner_gap = BAR_WIDTH * 0.15 + op_gap = BAR_WIDTH * 0.55 + hpss_gap = BAR_WIDTH * 0.30 + + pair_span = 2 * pair_width + inner_gap + hpss_group_span = 2 * pair_span + op_gap + group_span = len(HPSS_ORDER) * (hpss_group_span + hpss_gap) + 0.3 + x_base = np.arange(n_configs) * group_span + + for c_idx, (create_sub, update_sub) in enumerate(configs): + for h_idx, hpss in enumerate(HPSS_ORDER): + color = HPSS_COLORS[hpss] + hpss_origin = x_base[c_idx] + h_idx * (hpss_group_span + hpss_gap) + for op_idx, op in enumerate(ops): + base_hatch = op_hatches[op] + op_origin = hpss_origin + op_idx * (pair_span + op_gap) + x_bas_bar = op_origin + x_cur_bar = op_origin + pair_width + inner_gap + + def mean_for(df, _op=op, _h=hpss, _cs=create_sub, _us=update_sub): + v = ( + df[ + (df["operation"] == _op) + & (df["hpss_label"] == _h) + & (df["create_subdir"] == _cs) + & (df["update_subdir"] == _us) + ]["elapsed_seconds"] + .dropna() + .values + ) + return v.mean() if len(v) > 0 else 0.0 + + cur_mean = mean_for(df_cur) + bas_mean = mean_for(df_bas) + + ax.bar( + x_bas_bar, + bas_mean, + width=pair_width, + color=color, + hatch=base_hatch + "....", + alpha=0.35, + zorder=2, + edgecolor=color, + ) + ax.bar( + x_cur_bar, + cur_mean, + width=pair_width, + color=color, + hatch=base_hatch, + alpha=0.85, + zorder=2, + ) + + if bas_mean > 0 and cur_mean > 0: + ratio = cur_mean / bas_mean + top = max(cur_mean, bas_mean) + arrow = ( + "▲" + if ratio >= RATIO_REGRESSION + else ("▼" if ratio <= RATIO_IMPROVEMENT else "") + ) + ax.text( + (x_bas_bar + x_cur_bar) / 2 + pair_width / 2, + top * 1.03, + f"{arrow}{ratio:.2f}×", + ha="center", + va="bottom", + fontsize=5.5, + fontweight="bold", + color=_ratio_color(ratio), + zorder=4, + ) + + group_total_bar_span = len(HPSS_ORDER) * (hpss_group_span + hpss_gap) - hpss_gap + x_ticks = x_base + group_total_bar_span / 2 + ax.set_xticks(x_ticks) + ax.set_xticklabels([_extract_tick_label(c, u) for c, u in configs], fontsize=7) + ax.set_ylabel("Wall-clock time (s)", fontsize=8) + ax.set_xlabel("Archive contents (create → update)", fontsize=8, labelpad=6) + + workers = "1 worker" if check_op == "check_seq" else "2 workers" + ax.set_title( + f"check vs extract ({workers}) — current vs baseline\n" + f"Hatch = check / Solid = extract / Faded = baseline", + fontsize=10, + fontweight="bold", + pad=6, + ) + ax.yaxis.grid(True, linestyle="--", alpha=0.5, zorder=0) + ax.set_axisbelow(True) + + hpss_patches = [ + mpatches.Patch(color=HPSS_COLORS[h], label=HPSS_LABELS[h]) for h in HPSS_ORDER + ] + check_cur = mpatches.Patch( + facecolor="grey", hatch="////", alpha=0.85, label="check, current" + ) + check_bas = mpatches.Patch( + facecolor="grey", hatch="////....", alpha=0.35, label="check, baseline" + ) + ext_cur = mpatches.Patch( + facecolor="grey", hatch="", alpha=0.85, label="extract, current" + ) + ext_bas = mpatches.Patch( + facecolor="grey", hatch="....", alpha=0.35, label="extract, baseline" + ) + ax.legend( + handles=hpss_patches + [check_cur, check_bas, ext_cur, ext_bas], + fontsize=6.5, + loc="upper right", + ncol=3, + ) + + +def build_check_comparison_figure( + df_cur: pd.DataFrame, + df_bas: pd.DataFrame, + cur_label: str, + bas_label: str, +) -> Optional[plt.Figure]: + """ + Build Figure 3b: check vs extract with current/baseline pairing. + + Layout (3 rows × 2 cols): + Row 0: check_seq (cur vs bas) | check_par (cur vs bas) + Row 1: check_seq vs extract_seq (cur vs bas) + Row 2: check_par vs extract_par (cur vs bas) + """ + if not df_cur["operation"].isin(["check_seq", "check_par"]).any(): + return None + + fig = plt.figure(figsize=(16, 17)) + fig.suptitle( + f"zstash check Performance: Current vs Baseline\n" + f"current = {cur_label} | baseline = {bas_label}\n" + f"Ratio = current / baseline — " + f"▲ {RATIO_REGRESSION_COLOR_LABEL} ≥{RATIO_REGRESSION:.0%} slower " + f"▼ {RATIO_IMPROVEMENT_COLOR_LABEL} ≤{RATIO_IMPROVEMENT:.0%} faster " + f"= within ±10%", + fontsize=11, + fontweight="bold", + y=0.98, + ) + + gs = fig.add_gridspec( + 3, 2, hspace=0.58, wspace=0.35, top=0.92, bottom=0.07, left=0.07, right=0.97 + ) + + ax_check_seq = fig.add_subplot(gs[0, 0]) + ax_check_par = fig.add_subplot(gs[0, 1]) + _plot_comparison_extract_single_op(ax_check_seq, df_cur, df_bas, "check_seq") + _plot_comparison_extract_single_op(ax_check_par, df_cur, df_bas, "check_par") + + ax_cmp_seq = fig.add_subplot(gs[1, :]) + _plot_check_vs_extract_pair_comparison( + ax_cmp_seq, df_cur, df_bas, "check_seq", "extract_seq" + ) + + ax_cmp_par = fig.add_subplot(gs[2, :]) + _plot_check_vs_extract_pair_comparison( + ax_cmp_par, df_cur, df_bas, "check_par", "extract_par" + ) + + # Shared legend for the standalone row + cur_patch = mpatches.Patch(facecolor="grey", alpha=0.85, label="Current branch") + bas_patch = mpatches.Patch( + facecolor="grey", alpha=0.40, hatch="////", label="Baseline (main)" + ) + hpss_patches = [ + mpatches.Patch(color=HPSS_COLORS[h], label=HPSS_LABELS[h]) for h in HPSS_ORDER + ] + ax_check_seq.legend( + handles=[cur_patch, bas_patch] + hpss_patches, + fontsize=7, + loc="upper right", + ) + + return fig + + +# String labels used in the suptitle (avoids referencing undefined vars earlier) +RATIO_REGRESSION_COLOR_LABEL = "red" +RATIO_IMPROVEMENT_COLOR_LABEL = "green" + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def main(): + parser = argparse.ArgumentParser( + description="Visualise zstash performance results." + ) + parser.add_argument( + "--dpi", type=int, default=150, help="Output DPI (default: 150)" + ) + args = parser.parse_args() + + results_path = Path(RESULTS_CSV) + if not RESULTS_CSV or not results_path.is_file(): + print(f"ERROR: RESULTS_CSV not found: {RESULTS_CSV!r}", file=sys.stderr) + sys.exit(1) + df = load_data(str(results_path)) + if df.empty: + print( + f"ERROR: RESULTS_CSV is empty or could not be parsed: {RESULTS_CSV!r}", + file=sys.stderr, + ) + sys.exit(1) + + # Determine the sorted list of directories that actually appear in the data + all_dirs = sorted( + set(df["create_subdir"].dropna()) | set(df["update_subdir"].dropna()), + key=dir_sort_key, + ) + + # ----------------------------------------------------------------------- + # Figure layout: 3 rows × 2 cols + # Row 0: create | update + # Row 1: extract_seq | extract_par + # Row 2: extract seq-vs-par comparison (spans both columns) + # ----------------------------------------------------------------------- + fig = plt.figure(figsize=(15, 16)) + fig.suptitle( + "zstash Performance Profiling\n" + "(bars = mean over test configs; dots = individual runs)", + fontsize=13, + fontweight="bold", + y=0.98, + ) + + gs = fig.add_gridspec( + 3, 2, hspace=0.55, wspace=0.35, top=0.93, bottom=0.07, left=0.07, right=0.97 + ) + + axes = { + "create": fig.add_subplot(gs[0, 0]), + "update": fig.add_subplot(gs[0, 1]), + "extract_seq": fig.add_subplot(gs[1, 0]), + "extract_par": fig.add_subplot(gs[1, 1]), + } + ax_cmp = fig.add_subplot(gs[2, :]) + + # ----------------------------------------------------------------------- + # Draw the four single-operation subplots + # ----------------------------------------------------------------------- + legend_handles = None + for op in OP_ORDER: + ax = axes[op] + if op in OP_DIR_COL: + df_op = df[df["operation"] == op] + plot_operation(ax, df_op, op, all_dirs) + else: + # extract_seq / extract_par each get a dedicated single-op view + # that still uses the combined archive config as the x-axis. + _plot_extract_single_op(ax, df, op) + + if legend_handles is None: + legend_handles = [ + mpatches.Patch(color=HPSS_COLORS[h], label=HPSS_LABELS[h]) + for h in HPSS_ORDER + ] + ax.legend(handles=legend_handles, fontsize=7, loc="upper right") + + # ----------------------------------------------------------------------- + # Draw the sequential vs parallel comparison subplot + # ----------------------------------------------------------------------- + plot_extract_comparison(ax_cmp, df) + + # ----------------------------------------------------------------------- + # Baseline comparison figure (Figure 2) + # ----------------------------------------------------------------------- + fig_cmp = None + df_bas = None + cur_label = Path(RESULTS_CSV).parent.name + bas_label = None + + if BASELINE_RESULTS_CSV: + bas_path = Path(BASELINE_RESULTS_CSV) + if not bas_path.exists(): + print( + f"WARNING: BASELINE_RESULTS_CSV not found: {bas_path}", file=sys.stderr + ) + print("Skipping baseline comparison figures.", file=sys.stderr) + else: + df_bas = load_data(str(bas_path)) + bas_label = bas_path.parent.name + fig_cmp = build_comparison_figure( + df, df_bas, all_dirs, cur_label, bas_label + ) + + # ----------------------------------------------------------------------- + # Figure 3 – check performance and check-vs-extract + # ----------------------------------------------------------------------- + fig_check = build_check_figure(df) + fig_check_cmp = None + bas_has_check = ( + df_bas is not None + and df_bas["operation"].isin(["check_seq", "check_par"]).any() + ) + if bas_has_check and fig_check is not None: + fig_check_cmp = build_check_comparison_figure(df, df_bas, cur_label, bas_label) + + # ----------------------------------------------------------------------- + # Save or show + # ----------------------------------------------------------------------- + def save_or_show(figure, out_path_str, label): + if out_path_str: + out_path = Path(out_path_str) + out_path.parent.mkdir(parents=True, exist_ok=True) + figure.savefig(out_path, dpi=args.dpi, bbox_inches="tight") + print(f"{label} saved to: {out_path}") + os.chmod(out_path_str, 0o644) + web_path = str(out_path).replace( + "/global/cfs/cdirs/e3sm/www/", + "https://portal.nersc.gov/cfs/e3sm/", + ) + print(f" Accessible at: {web_path}") + else: + plt.show() + + if OUTPUT_PATH: + save_or_show(fig, OUTPUT_PATH, "Figure 1 (overview)") + if fig_cmp is not None: + p = Path(OUTPUT_PATH) + cmp_output = str(p.with_stem(p.stem + "_vs_baseline")) + save_or_show(fig_cmp, cmp_output, "Figure 2 (baseline comparison)") + if fig_check is not None: + p = Path(OUTPUT_PATH) + check_output = str(p.with_stem(p.stem + "_check")) + save_or_show(fig_check, check_output, "Figure 3 (check)") + if fig_check_cmp is not None: + p = Path(OUTPUT_PATH) + check_cmp_output = str(p.with_stem(p.stem + "_check_vs_baseline")) + save_or_show( + fig_check_cmp, check_cmp_output, "Figure 3b (check vs baseline)" + ) + else: + plt.show() + + +if __name__ == "__main__": + np.random.seed(42) # reproducible jitter + main()