Optimize zstash update by using scandir - #420
Conversation
|
The list order issue found during performance profiling is now resolved. |
Added unit tests Perlmutter tests passed Chrysalis tests passed File gathering list order matches that of the main branch
8234b99 to
a704282
Compare
forsyth2
left a comment
There was a problem hiding this comment.
Self-review of code. I've confirmed the entire test suite passes and the file addition list is unchanged from main.
| @@ -0,0 +1,657 @@ | |||
| """ | |||
| Unit tests for the optimized file scanning implementation in zstash. | |||
There was a problem hiding this comment.
For this review, I'm mainly focused on the non-test changes. I did confirm these Claude-generated tests pass though.
| logger.debug("Keep local tar files : {}".format(keep)) | ||
|
|
||
| files: List[str] = get_files_to_archive(cache, args.include, args.exclude) | ||
| file_stats: Dict[str, Tuple[int, datetime]] = get_files_to_archive_with_stats( |
There was a problem hiding this comment.
Now, we call the optimized get_files_to_archive
| archived_files: Dict[str, Tuple[int, datetime]] = {} | ||
|
|
||
| cur.execute("SELECT name, size, mtime FROM files") | ||
| db_rows = cur.fetchall() |
There was a problem hiding this comment.
Fetch all the rows at once
| size: int = row[1] | ||
| mtime: datetime = row[2] | ||
|
|
||
| # If file appears multiple times, keep the one with latest mtime |
| # Get the stat info we already collected during filesystem walk | ||
| size_new, mdtime_new = file_stats[file_path] |
There was a problem hiding this comment.
Now we don't have to run os.lstat -- we already have the value in memory.
| # Normalize the path. | ||
| # By building from path + entry.name, | ||
| # we guarantee the argument to normpath is constructed | ||
| # identically to how os.walk did it in previous code iterations. | ||
| normalized_path = os.path.normpath(os.path.join(path, entry.name)) | ||
| self.file_stats[normalized_path] = (size, mtime) |
There was a problem hiding this comment.
Important for maintaining sort order
| logger.warning(f"Error accessing {entry.path}: {e}") | ||
| continue | ||
|
|
||
| # Handle empty directories BEFORE recursing into subdirs |
There was a problem hiding this comment.
Important for maintaining sort order
| # Note: broken symlinks are caught by not os.path.isfile() but not by os.path.isdir(). We want the latter. | ||
| if size == 0 and os.path.isdir(path): |
There was a problem hiding this comment.
Important for maintaining sort order
| # Sort on directory and filename like original | ||
| path_components.sort(key=lambda x: (x[0], x[1])) |
There was a problem hiding this comment.
Important for maintaining sort order
| LEGACY VERSION: Still used for `zstash create`. | ||
| Uses the optimized version but returns only the file list. | ||
| """ | ||
| file_stats = get_files_to_archive_with_stats(cache, include, exclude) |
There was a problem hiding this comment.
For zstash create, we just used reduced output from the optimized function
How to do a performance profilingI'm using Perlmutter throughout these directions. I'll use constants to refer to specific paths. Change these out to work on the machine you're on & to use your username. HPSS_PATH=/home/projects/e3sm/www/CoupledSystem/E3SMv3/LR/v3.LR.amip_0151
PROFILE_PATH=/pscratch/sd/f/forsyth/performance_profiling/
EXTRACTED_DIRECTORY_PATH=${PROFILE_PATH}extraction/
LOG_PATH=${PROFILE_PATH}logs/
CACHE_PATH=${PROFILE_PATH}cache/
INDEX_COPIES_PATH=${PROFILE_PATH}index_copies/Also note I'm using Extract an archiveSkip this step if you already have an archive ready to use. screen # This could run for a while, so let's use `screen`
screen -ls # We need to know what node this screen is on, so we can come back to it.
mkdir ${PROFILE_PATH}
mkdir ${CACHE_PATH}
mkdir ${EXTRACTED_DIRECTORY_PATH}
mkdir ${LOG_PATH}
cd ${EXTRACTED_DIRECTORY_PATH}
# We'll extract using the Unified environment's version of zstash
source /global/common/software/e3sm/anaconda_envs/load_latest_e3sm_unified_pm-cpu.sh # Activate E3SM Unified
zstash extract --hpss=${HPSS_PATH} --cache=${CACHE_PATH} -v 2>&1 | tee ${LOG_PATH}$/extract.logBe sure to save the index.db as it exists initiallycp ${CACHE_PATH}/index.db ${INDEX_COPIES_PATH}/index_original.dbProfile the before-caseProfile using zstash as it exists in the latest E3SM Unified environment. source /global/common/software/e3sm/anaconda_envs/load_latest_e3sm_unified_pm-cpu.sh # Activate E3SM Unified
before_case_log_file=${LOG_PATH}update_before_changes.log
cd ${EXTRACTED_DIRECTORY_PATH} # Must run update from the archived directory!
zstash update --hpss=${HPSS_PATH} --cache=${CACHE_PATH} -v --dry-run 2>&1 | tee ${before_case_log_file}
mv ${CACHE_PATH}index.db ${INDEX_COPIES_PATH}index_before_changes.db # Optional, only if you want to keep the index
cp ${INDEX_COPIES_PATH}/index_original.db ${CACHE_PATH}index.db # Reset the index!!
grep -n real ${before_case_log_file}
# This will show the real time spentProfile the after-caseProfile using zstash as it exists on this branch. # Skip this block if you already have a zstash repo:
cd ~
git clone git@github.com:E3SM-Project/zstash.git
cd zstash
git remote -v
# You should see:
# origin git@github.com:E3SM-Project/zstash.git (fetch)
# origin git@github.com:E3SM-Project/zstash.git (push)
# Now, get this branch:
git fetch upstream issue-409-optimize-update
git checkout -b issue-409-optimize-update upstream/issue-409-optimize-update
# Now, set up the conda environment
lcrc_conda # Activate conda. This is my function to do so. You may use a different method.
rm -rf build
conda clean --all --y
conda env create -f conda/dev.yml -n profile-optimized-zstash-update
conda activate profile-optimized-zstash-update
python -m pip install . # Install the branch's changes into this conda environment.
# Now, we can actually profile, just as we did with the before-case
after_case_log_file=${LOG_PATH}update_after_changes.log
cd ${EXTRACTED_DIRECTORY_PATH} # Must run update from the archived directory!
zstash update --hpss=${HPSS_PATH} --cache=${CACHE_PATH} -v --dry-run 2>&1 | tee ${after_case_log_file}
mv ${CACHE_PATH}index.db ${INDEX_COPIES_PATH}index_after_changes.db # Optional, only if you want to keep the index
cp ${INDEX_COPIES_PATH}/index_original.db ${CACHE_PATH}index.db # Reset the index!!
grep -n real ${after_case_log_file}
# This will show the real time spentComparing profiling resultsTo get the speedup, divide the real time spent: To check the file list accuracy, you can just compare the logs directly (assuming you used diff ${before_case_log_file} ${after_case_log_file}
# Should show no results (if --dry-run was used) |
|
@chengzhuzhang @golaz @wlin7 I've created this PR as a simplified code change since #414 had a lot of performance logging in it that won't be necessary on the @chengzhuzhang @golaz if you'd like to do a visual inspection of the code, you can review the file diffs and my explanatory notes above. The implementation is what we've been discussing -- the replacing of @wlin7 I think we are ready to have you do a production/large/real use case test. I've included directions above. Please let me know if you have any questions. Thanks all! |
chengzhuzhang
left a comment
There was a problem hiding this comment.
@forsyth2 I visually checked the code update and it looks good to me. I added to the PR summary, so that it provides more information for reviewers.
|
Thanks @chengzhuzhang! |
|
Hi @forsyth2 . This is the test results using a production simulation (
The |
wlin7
left a comment
There was a problem hiding this comment.
Test results verified the performance gain and the integrity of zstash update.
|
Fantastic, thank you for testing, @wlin7! |
|
I think this should be ready to merge. After that, I can try to set up a dev environment others can access. Completed:
|
|
@wlin7 @xuezhengllnl I've created an environment conda env create -f /home/ac.forsyth2/user_envs/zstash-user-env-20260122.ymlI typically only create Steps I used to set up this environmentChrysalis: cd ~/ez/zstash
git status
# nothing to commit, working tree clean
git fetch upstream main
git checkout main
git reset --hard upstream/main
git log --oneline | head -n 2
# 880b9fd Optimize zstash update by using scandir (#420)
# 01ffe34 Add test for tar deletion (#404)
# => Good, we've included the optimization.
lcrc_conda # My function to activate Conda.
rm -rf build
conda clean --all --y
conda env create -f conda/dev.yml -n zstash-user-env-20260122
conda activate zstash-user-env-20260122
pre-commit run --all-files # Optional since we didn't add any code
python -m pip install .
mkdir ~/user_envs
conda env export > ~/user_envs/zstash-user-env-20260122.yml
ls -l ~/user_envs/zstash-user-env-20260122.yml
# -rw-r--r-- ... |
|
Using #420 (comment), I was able to produce this plot:
with this script: Expandimport os
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import numpy as np
MACHINE_WWW_PATH = "/global/cfs/cdirs/e3sm/www/forsyth"
WWW_WEB_PATH = "https://portal.nersc.gov/cfs/e3sm/forsyth"
# ── Palette ────────────────────────────────────────────────────────────────
BEFORE_COLOR = "#4A6FA5" # muted steel blue
AFTER_COLOR = "#E07A5F" # warm coral
BG_COLOR = "#F7F7F2" # off-white
GRID_COLOR = "#DDDDD5"
TEXT_COLOR = "#2B2B2B"
ACCENT_COLOR = "#3D9970" # improvement annotation
# ── Data ───────────────────────────────────────────────────────────────────
metrics = ["File Gathering\n(time before creating first tar)\nTime (s)", "Total zstash\n(time until final tar)\nTime (min)"]
before_vals = [556, 78 + 18/60] # seconds | minutes
after_vals = [152, 72 + 28/60] # seconds | minutes
units = ["s", "min"]
# Reduction percentages
reductions = [
(before_vals[0] - after_vals[0]) / before_vals[0] * 100,
(before_vals[1] - after_vals[1]) / before_vals[1] * 100,
]
# ── Layout ─────────────────────────────────────────────────────────────────
fig, axes = plt.subplots(1, 2, figsize=(12, 6), facecolor=BG_COLOR)
fig.subplots_adjust(wspace=0.38, top=0.74, bottom=0.14, left=0.08, right=0.97)
bar_width = 0.42
x = np.array([0])
for ax, metric, cv, nv, unit, red in zip(
axes, metrics, before_vals, after_vals, units, reductions):
ax.set_facecolor(BG_COLOR)
for spine in ax.spines.values():
spine.set_visible(False)
# Gridlines (horizontal only, behind bars)
ax.yaxis.grid(True, color=GRID_COLOR, linewidth=0.8, zorder=0)
ax.set_axisbelow(True)
# Bars
b1 = ax.bar(x - bar_width/2, cv, width=bar_width,
color=BEFORE_COLOR, zorder=3, label="Before",
linewidth=0)
b2 = ax.bar(x + bar_width/2, nv, width=bar_width,
color=AFTER_COLOR, zorder=3, label="After",
linewidth=0)
# Value labels on bars
for bar, val in zip([b1[0], b2[0]], [cv, nv]):
label = f"{val:.0f} {unit}" if unit == "s" else f"{val:.1f} {unit}"
ax.text(
bar.get_x() + bar.get_width() / 2,
bar.get_height() + cv * 0.01,
label,
ha="center", va="bottom",
fontsize=12, fontweight="bold", color=TEXT_COLOR,
)
# Reduction annotation (bracket + text)
y_top = cv * 1.07
y_brace = cv * 1.18
# ax.annotate(
# "", xy=(x[0] + bar_width/2, y_top),
# xytext=(x[0] - bar_width/2, y_top),
# arrowprops=dict(arrowstyle="-", color=ACCENT_COLOR, lw=1.6),
# )
ax.text(
x[0], y_brace,
f"−{red:.0f}%",
ha="center", va="bottom",
fontsize=13, fontweight="bold", color=ACCENT_COLOR,
)
# Axes formatting
ax.set_xticks([])
ax.set_ylabel(f"Time ({unit})", fontsize=11, color=TEXT_COLOR, labelpad=8)
ax.set_title(metric, fontsize=13, fontweight="bold",
color=TEXT_COLOR, pad=10)
ax.tick_params(axis="y", labelsize=10, colors=TEXT_COLOR)
ax.set_xlim(-0.55, 0.55)
ax.set_ylim(0, cv * 1.38)
# ── Figure-level title & subtitle ──────────────────────────────────────────
fig.text(
0.5, 0.95,
"zstash Performance",
ha="center", va="top",
fontsize=17, fontweight="bold", color=TEXT_COLOR,
fontfamily="DejaVu Sans",
)
fig.text(
0.5, 0.88,
"Test case: v3.LR.amip_0101 (8.4 TB / 39K files archived) "
"+ v3.LR.amip_0151 add (810 GB / 11K files)",
ha="center", va="top",
fontsize=9.5, color="#555550",
)
# ── Legend ─────────────────────────────────────────────────────────────────
legend_handles = [
mpatches.Patch(color=BEFORE_COLOR, label="os.lstat (E3SM Unified 1.12.0, zstash v1.5.0)"),
mpatches.Patch(color=AFTER_COLOR, label="os.scandir (later included in E3SM Unified 1.13.0, zstash v1.6.0)"),
]
fig.legend(
handles=legend_handles,
loc="lower center",
ncol=2,
fontsize=10,
frameon=False,
bbox_to_anchor=(0.5, 0.01),
)
# ── Save & show ────────────────────────────────────────────────────────────
file_name: str = "visualize_pr420_results.png"
file_location: str = f"{MACHINE_WWW_PATH}/{file_name}"
plt.savefig(file_location,
dpi=200, bbox_inches="tight", facecolor=BG_COLOR)
plt.show()
os.chmod(file_location, 0o755)
print(f"Saved to {file_location}")
print(f"View at {WWW_WEB_PATH}/{file_name}") |
|
@chengzhuzhang More plots we could use for the poster: RevisionRevision of plot above, again using the data from the table in #420 (comment)
Generator scriptimport os
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import numpy as np
MACHINE_WWW_PATH = "/global/cfs/cdirs/e3sm/www/forsyth"
WWW_WEB_PATH = "https://portal.nersc.gov/cfs/e3sm/forsyth"
# ── Palette ────────────────────────────────────────────────────────────────
BEFORE_COLOR = "#4A6FA5"
AFTER_COLOR = "#E07A5F"
BG_COLOR = "#F7F7F2"
GRID_COLOR = "#DDDDD5"
TEXT_COLOR = "#2B2B2B"
ACCENT_COLOR = "#3D9970"
# ── Data ───────────────────────────────────────────────────────────────────
metric = "File Gathering\n(time before creating first tar)\nTime (s)"
before_val = 556
after_val = 152
unit = "s"
reduction = (before_val - after_val) / before_val * 100
# ── Layout ─────────────────────────────────────────────────────────────────
fig, ax = plt.subplots(1, 1, figsize=(7, 6), facecolor=BG_COLOR)
fig.subplots_adjust(top=0.70, bottom=0.18, left=0.14, right=0.92)
bar_width = 0.42
x = np.array([0])
ax.set_facecolor(BG_COLOR)
for spine in ax.spines.values():
spine.set_visible(False)
ax.yaxis.grid(True, color=GRID_COLOR, linewidth=0.8, zorder=0)
ax.set_axisbelow(True)
b1 = ax.bar(x - bar_width/2, before_val, width=bar_width,
color=BEFORE_COLOR, zorder=3, linewidth=0)
b2 = ax.bar(x + bar_width/2, after_val, width=bar_width,
color=AFTER_COLOR, zorder=3, linewidth=0)
# Value labels on bars
for bar, val in zip([b1[0], b2[0]], [before_val, after_val]):
ax.text(
bar.get_x() + bar.get_width() / 2,
bar.get_height() + before_val * 0.01,
f"{val:.0f} {unit}",
ha="center", va="bottom",
fontsize=15, fontweight="bold", color=TEXT_COLOR,
)
# Reduction annotation
y_brace = before_val * 1.18
ax.text(
x[0], y_brace,
f"−{reduction:.0f}%",
ha="center", va="bottom",
fontsize=16, fontweight="bold", color=ACCENT_COLOR,
)
ax.set_xticks([])
ax.set_ylabel(f"Time ({unit})", fontsize=14, color=TEXT_COLOR, labelpad=8)
ax.set_title(metric, fontsize=15, fontweight="bold", color=TEXT_COLOR, pad=10)
ax.tick_params(axis="y", labelsize=13, colors=TEXT_COLOR)
ax.set_xlim(-0.55, 0.55)
ax.set_ylim(0, before_val * 1.38)
# ── Figure-level title & subtitle ──────────────────────────────────────────
fig.text(
0.5, 0.95,
"zstash Performance",
ha="center", va="top",
fontsize=20, fontweight="bold", color=TEXT_COLOR,
fontfamily="DejaVu Sans",
)
fig.text(
0.5, 0.88,
"Test case: v3.LR.amip_0101 (8.4 TB / 39K files archived) "
"+ v3.LR.amip_0151 add (810 GB / 11K files)",
ha="center", va="top",
fontsize=11, color="#555550",
)
# ── Legend ─────────────────────────────────────────────────────────────────
legend_handles = [
mpatches.Patch(color=BEFORE_COLOR, label="os.lstat (E3SM Unified 1.12.0, zstash v1.5.0)"),
mpatches.Patch(color=AFTER_COLOR, label="os.scandir (later included in E3SM Unified 1.13.0, zstash v1.6.0)"),
]
fig.legend(
handles=legend_handles,
loc="lower center",
ncol=1,
fontsize=12,
frameon=False,
bbox_to_anchor=(0.5, 0.01),
)
# ── Save & show ────────────────────────────────────────────────────────────
file_name: str = "visualize_pr420_results3.png"
file_location: str = f"{MACHINE_WWW_PATH}/{file_name}"
plt.savefig(file_location,
dpi=200, bbox_inches="tight", facecolor=BG_COLOR)
plt.show()
os.chmod(file_location, 0o755)
print(f"Saved to {file_location}")
print(f"View at {WWW_WEB_PATH}/{file_name}")
New plotNew plot, made using data from #414 (comment)
Generator script"""
Performance comparison: os.lstat (before) vs os.scandir (after)
Data from zstash performance profiling run 2026-01-08.
Script developed by Claude from:
https://github.com/E3SM-Project/zstash/pull/414#issuecomment-3726277207
"""
import matplotlib.pyplot as plt
import numpy as np
import os
MACHINE_WWW_PATH = "/global/cfs/cdirs/e3sm/www/forsyth"
WWW_WEB_PATH = "https://portal.nersc.gov/cfs/e3sm/forsyth"
cases = ["Add few\nsmall files", "Add many\nsmall files", "Add few\nlarge files"]
# Combined (file gathering + db comparison) time in seconds
before = [1119.53, 1246.23, 1104.21] # lstat
after = [39.55, 54.61, 33.30] # scandir
x = np.arange(len(cases))
width = 0.35
fig, ax = plt.subplots(figsize=(7, 5))
ax.bar(x - width/2, before, width, label="lstat (before)", color="#4C72B0", alpha=0.9)
ax.bar(x + width/2, after, width, label="scandir (after)", color="#DD8452", alpha=0.9)
# Speedup annotation above each pair
for i, (b, a) in enumerate(zip(before, after)):
ax.text(x[i], max(b, a) + 20, f"{b/a:.0f}× faster",
ha="center", va="bottom", fontsize=11, fontweight="bold", color="#2ca02c")
ax.set_ylabel("Time (seconds)", fontsize=11)
ax.set_title("zstash update: lstat vs scandir\n(file gathering + database comparison)", fontsize=12)
ax.set_xticks(x)
ax.set_xticklabels(cases, fontsize=10)
ax.legend(fontsize=10)
ax.grid(axis="y", linestyle="--", alpha=0.4)
ax.set_ylim(0, 1400)
plt.tight_layout()
file_name: str = "visualize_pr420_results2.png"
file_location: str = f"{MACHINE_WWW_PATH}/{file_name}"
plt.savefig(file_location, dpi=150, bbox_inches="tight")
os.chmod(file_location, 0o755)
print(f"Saved to {file_location}")
print(f"View at {WWW_WEB_PATH}/{file_name}")
|



Summary
Objectives:
What changed:
Issue resolution:
Select one: This pull request is...
Small Change