Bench/zstd integration - #68
Conversation
… sync Emit grep-able per-stage timing (server_start_s, update_weights_from_disk_s) to stdout plus fire-and-forget monitoring gauges. Fully guarded so measurement can never break the sync path. Bench context: docs/goal.md Phase 2 S4 (miner-side weight sync latency measurement). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Completing a 57 GB multipart upload against a slow local S3 mock (moto, single-threaded in-memory assembly) exceeds the hardcoded 30 s read timeout. Default unchanged. Bench context: docs/goal.md Phase 2 S7. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
Warning Review limit reached
Next review available in: 35 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (28)
WalkthroughThe PR adds benchmark scripts and Slurm jobs covering sparse delta checkpoints, mocked S3 throughput, checkpoint-chain smoke tests, NCCL synchronization, disk-backed transfers, and distributed REAL-training delta workflows. It also adds synchronization metrics and configurable S3 read timeouts. ChangesBenchmark and checkpoint synchronization suite
Estimated code review effort: 5 (Critical) | ~120 minutes Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 8
🧹 Nitpick comments (3)
docs/scripts/s9_real_delta_bench.py (1)
122-122: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueHoist stdlib inline imports to the module top. Both scripts import standard-library modules inside functions/branches; move the movable ones to the top-level import block (stdlib group). The
sparse_codecimport viasys.pathis intentionally deferred and can stay inline.
docs/scripts/s9_real_delta_bench.py#L122-L122: moveimport functools(andpandasat L230,threadingat L398/L467) to the top.docs/scripts/s22_multistep_delta_bench.py#L120-L120: moveimport functools(andjson as _jsonat L185,pandasat L228,dequeat L275) to the top.As per coding guidelines: "Place imports at the top of Python files; order: stdlib, third-party, local; avoid inline imports".
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@docs/scripts/s9_real_delta_bench.py` at line 122, Move the movable inline imports to each script’s module-level import block, preserving stdlib-before-third-party ordering: in docs/scripts/s9_real_delta_bench.py move functools, threading, and pandas; in docs/scripts/s22_multistep_delta_bench.py move functools, json as _json, deque, and pandas. Leave the deferred sparse_codec sys.path import inline.Source: Coding guidelines
docs/scripts/s4_rotate_publisher.py (1)
45-47: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
current_window()formula duplicated verbatim across both rotator scripts. The window-boundary calculation ((block // WINDOW_LEN) * WINDOW_LEN) is copy-pasted identically in two files; since correctness of the whole delta/full rotation flow depends on rotators and the mocked miner agreeing on window boundaries, any future tweak to this formula (or toWINDOW_LEN) risks silently drifting between the two files if only one is updated.
docs/scripts/s4_rotate_publisher.py#L45-L47: extractcurrent_window()(and the sharedWINDOW_LEN/GENESIS/BLOCK_TIMEconstants) into a small shared bench-utils module imported by both rotators.docs/scripts/s7_rotate_delta_publisher.py#L57-L58: import the same shared helper instead of redefining it.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@docs/scripts/s4_rotate_publisher.py` around lines 45 - 47, The current_window formula and its WINDOW_LEN, GENESIS, and BLOCK_TIME constants are duplicated across both rotator scripts. In docs/scripts/s4_rotate_publisher.py lines 45-47, extract these shared constants and current_window() into a small bench-utils module; in docs/scripts/s7_rotate_delta_publisher.py lines 57-58, remove the duplicate definitions and import the shared helper so both rotators use the same window-boundary implementation.docs/scripts/s2_run_miner_mocked.py (1)
27-28: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winHardcoded per-user absolute path duplicated across 5 scripts. Every entry point independently hardcodes
sys.path.insert(0, "/storage/openpsi/users/pengzai.pyq/grail")instead of reading theGRAIL_SRC/WORKDIRvalue the sbatch wrappers already compute; some wrappers already exportPYTHONPATHfrom that same variable, making the Python-side hardcode both redundant and a source of stale imports if the path ever changes or another user runs these scripts.
docs/scripts/s2_run_miner_mocked.py#L27-L28: derive the path from an env var (e.g.GRAIL_SRC) instead of the literal string.docs/scripts/s4_rotate_publisher.py#L21-L21: same change.docs/scripts/s5_sglang_reload_bench.py#L27-L27: same change.docs/scripts/s6_disk_segments_bench.py#L29-L29: same change.docs/scripts/s7_rotate_delta_publisher.py#L23-L23: same change.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@docs/scripts/s2_run_miner_mocked.py` around lines 27 - 28, Replace the hardcoded user-specific source paths with an environment-derived path such as GRAIL_SRC, preserving each script’s existing sys.path setup. Apply this change in docs/scripts/s2_run_miner_mocked.py lines 27-28, docs/scripts/s4_rotate_publisher.py line 21, docs/scripts/s5_sglang_reload_bench.py line 27, docs/scripts/s6_disk_segments_bench.py line 29, and docs/scripts/s7_rotate_delta_publisher.py line 23, using the sbatch-provided environment value so all entry points resolve the current checkout consistently.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@docs/sbatch/s13_preflight.sbatch`:
- Around line 39-50: Update the singularity execution flow before the RC
assignment so its non-zero status is captured without triggering set -e, while
preserving the command and its arguments. Ensure the subsequent completion
message reports the captured status and the script exits with that same status.
In `@docs/scripts/probe_grail.py`:
- Around line 24-29: Move the local imports from lines 24-29, the third-party
imports from lines 34-38, and the third-party import from line 127 in
docs/scripts/probe_grail.py to the file’s import section, ordered as
standard-library, third-party, then local modules; remove the original inline
import statements while preserving existing usage.
- Around line 44-55: Refactor docs/scripts/probe_grail.py by replacing all print
calls with a module logger obtained via logging.getLogger(__name__), adding
explicit type annotations for variables and relevant function values to satisfy
mypy-strict, and moving module-level execution into single-responsibility
functions of roughly 5–20 lines. Preserve the existing probe behavior,
assertions, and output information while keeping the script entry point minimal.
- Line 21: Replace the hardcoded configuration values in
docs/scripts/probe_grail.py at lines 21-21, 129-131, and 135-135: derive the
grail source path from __file__ or load it from an environment variable, load
the target model path from an environment variable, and load the storage disk
path from an environment variable instead of /storage. Use the existing
probe_grail configuration flow and provide appropriate handling for unset
variables.
In `@docs/scripts/s13_transfer_bench.py`:
- Around line 225-227: Guard the directory-creation step before writing the
output in the transfer benchmark flow: only call os.makedirs when
os.path.dirname(args.out) is non-empty, while preserving the existing
file-writing behavior for both bare filenames and paths with directories.
In `@docs/scripts/s2_run_miner_mocked.py`:
- Around line 46-49: Update the task-waiting error handling around
asyncio.wait_for in the harness so asyncio.TimeoutError is handled separately,
while other exceptions are logged with shutdown context and re-raised. Preserve
task cancellation on timeout and ensure genuine miner failures are no longer
silently swallowed.
In `@docs/scripts/s5_sglang_reload_bench.py`:
- Around line 78-113: Wrap the server lifecycle and benchmark request loop in a
try/finally so manager.__aexit__(None, None, None) always runs after
manager.__aenter__() succeeds, including when start_server(), client.post(), or
resp.raise_for_status() raises. Preserve the existing startup, measurement, and
result-recording behavior while guaranteeing SGLang process cleanup.
In `@docs/scripts/s6_disk_segments_bench.py`:
- Around line 113-116: Validate the result of load_model_state_dict(MODEL)
before the `for k in state` loop, and raise a clear error when it is None or
otherwise indicates missing weights. Preserve the existing cloning and
RAM-materialization behavior for valid state dictionaries.
---
Nitpick comments:
In `@docs/scripts/s2_run_miner_mocked.py`:
- Around line 27-28: Replace the hardcoded user-specific source paths with an
environment-derived path such as GRAIL_SRC, preserving each script’s existing
sys.path setup. Apply this change in docs/scripts/s2_run_miner_mocked.py lines
27-28, docs/scripts/s4_rotate_publisher.py line 21,
docs/scripts/s5_sglang_reload_bench.py line 27,
docs/scripts/s6_disk_segments_bench.py line 29, and
docs/scripts/s7_rotate_delta_publisher.py line 23, using the sbatch-provided
environment value so all entry points resolve the current checkout consistently.
In `@docs/scripts/s4_rotate_publisher.py`:
- Around line 45-47: The current_window formula and its WINDOW_LEN, GENESIS, and
BLOCK_TIME constants are duplicated across both rotator scripts. In
docs/scripts/s4_rotate_publisher.py lines 45-47, extract these shared constants
and current_window() into a small bench-utils module; in
docs/scripts/s7_rotate_delta_publisher.py lines 57-58, remove the duplicate
definitions and import the shared helper so both rotators use the same
window-boundary implementation.
In `@docs/scripts/s9_real_delta_bench.py`:
- Line 122: Move the movable inline imports to each script’s module-level import
block, preserving stdlib-before-third-party ordering: in
docs/scripts/s9_real_delta_bench.py move functools, threading, and pandas; in
docs/scripts/s22_multistep_delta_bench.py move functools, json as _json, deque,
and pandas. Leave the deferred sparse_codec sys.path import inline.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 3e15c7cb-c816-4680-a7c9-74fb5a308a4c
📒 Files selected for processing (30)
docs/sbatch/bench_qwen3_30b_grail_delta.sbatchdocs/sbatch/preflight_grail.sbatchdocs/sbatch/s13_preflight.sbatchdocs/sbatch/s3_miner_smoke.sbatchdocs/sbatch/s4_sync_bench.sbatchdocs/sbatch/s5_sglang_30b_reload.sbatchdocs/sbatch/s6_disk_segments.sbatchdocs/sbatch/s7_delta_e2e.sbatchdocs/sbatch/s8_nccl.sbatchdocs/scripts/bench_qwen3_30b_grail_delta.pydocs/scripts/grail_chain_mock.pydocs/scripts/moto_helper.shdocs/scripts/parse_sync_metrics.pydocs/scripts/probe_grail.pydocs/scripts/s0_comms_roundtrip.pydocs/scripts/s13_preflight_moto.pydocs/scripts/s13_transfer_bench.pydocs/scripts/s1_publish_consume.pydocs/scripts/s22_multistep_delta_bench.pydocs/scripts/s2_run_miner_mocked.pydocs/scripts/s4_rotate_publisher.pydocs/scripts/s5_sglang_reload_bench.pydocs/scripts/s6_disk_segments_bench.pydocs/scripts/s7_rotate_delta_publisher.pydocs/scripts/s8_nccl_bench.pydocs/scripts/s8_nccl_bench_2n.pydocs/scripts/s8_nccl_bench_fsdp.pydocs/scripts/s9_real_delta_bench.pygrail/infrastructure/comms.pygrail/mining/weight_sync.py
| singularity exec --no-home -B /storage:/storage \ | ||
| --env PYTHONPATH=$WORKDIR/tools/moto_env312 \ | ||
| "$SIF" /app/.venv/bin/python \ | ||
| "$GRAIL_SRC/docs/scripts/s13_preflight_moto.py" \ | ||
| --port "$PORT" \ | ||
| --sizes 0.5,1,3 \ | ||
| --rounds 3 \ | ||
| --out "$JSON_OUT" | ||
| RC=$? | ||
|
|
||
| echo "===== S13 preflight done rc=$RC JSON=$JSON_OUT =====" | ||
| exit $RC |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
RC=$? is unreachable on failure under set -e.
With errexit enabled (Line 17), a non-zero exit from the singularity exec block terminates the script immediately (firing the EXIT trap), so Line 47 only runs on success and RC is always 0. The rc=$RC report never reflects a real failure.
🐛 Capture rc without tripping errexit
-singularity exec --no-home -B /storage:/storage \
+RC=0
+singularity exec --no-home -B /storage:/storage \
--env PYTHONPATH=$WORKDIR/tools/moto_env312 \
"$SIF" /app/.venv/bin/python \
"$GRAIL_SRC/docs/scripts/s13_preflight_moto.py" \
--port "$PORT" \
--sizes 0.5,1,3 \
--rounds 3 \
- --out "$JSON_OUT"
-RC=$?
+ --out "$JSON_OUT" || RC=$?📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| singularity exec --no-home -B /storage:/storage \ | |
| --env PYTHONPATH=$WORKDIR/tools/moto_env312 \ | |
| "$SIF" /app/.venv/bin/python \ | |
| "$GRAIL_SRC/docs/scripts/s13_preflight_moto.py" \ | |
| --port "$PORT" \ | |
| --sizes 0.5,1,3 \ | |
| --rounds 3 \ | |
| --out "$JSON_OUT" | |
| RC=$? | |
| echo "===== S13 preflight done rc=$RC JSON=$JSON_OUT =====" | |
| exit $RC | |
| RC=0 | |
| singularity exec --no-home -B /storage:/storage \ | |
| --env PYTHONPATH=$WORKDIR/tools/moto_env312 \ | |
| "$SIF" /app/.venv/bin/python \ | |
| "$GRAIL_SRC/docs/scripts/s13_preflight_moto.py" \ | |
| --port "$PORT" \ | |
| --sizes 0.5,1,3 \ | |
| --rounds 3 \ | |
| --out "$JSON_OUT" || RC=$? | |
| echo "===== S13 preflight done rc=$RC JSON=$JSON_OUT =====" | |
| exit $RC |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@docs/sbatch/s13_preflight.sbatch` around lines 39 - 50, Update the
singularity execution flow before the RC assignment so its non-zero status is
captured without triggering set -e, while preserving the command and its
arguments. Ensure the subsequent completion message reports the captured status
and the script exits with that same status.
| import sys | ||
| import time | ||
|
|
||
| sys.path.insert(0, "/storage/openpsi/users/pengzai.pyq/grail") |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Load configuration via environment variables.
As per coding guidelines, configuration and cluster-specific values must be loaded via environment variables (e.g., using dotenv or os.environ.get) rather than being hardcoded.
docs/scripts/probe_grail.py#L21-L21: Load thegrailsource path using an environment variable or a relative path derived from__file__.docs/scripts/probe_grail.py#L129-L131: Load the target model path using an environment variable instead of hardcoding it.docs/scripts/probe_grail.py#L135-L135: Load the storage disk path using an environment variable instead of hardcoding/storage.
📍 Affects 1 file
docs/scripts/probe_grail.py#L21-L21(this comment)docs/scripts/probe_grail.py#L129-L131docs/scripts/probe_grail.py#L135-L135
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@docs/scripts/probe_grail.py` at line 21, Replace the hardcoded configuration
values in docs/scripts/probe_grail.py at lines 21-21, 129-131, and 135-135:
derive the grail source path from __file__ or load it from an environment
variable, load the target model path from an environment variable, and load the
storage disk path from an environment variable instead of /storage. Use the
existing probe_grail configuration flow and provide appropriate handling for
unset variables.
Source: Coding guidelines
| from grail.infrastructure.delta_checkpoint import ( # noqa: E402 | ||
| apply_sparse_delta, | ||
| compute_sparse_delta, | ||
| compute_weights_hash, | ||
| estimate_sparse_size, | ||
| ) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Move all imports to the top of the file.
As per coding guidelines, avoid inline imports and place all imports at the top of the file, explicitly ordered by standard library, third-party, and local modules.
docs/scripts/probe_grail.py#L24-L29: Move these local imports to the top.docs/scripts/probe_grail.py#L34-L38: Move these third-party library imports to the top.docs/scripts/probe_grail.py#L127-L127: Move this third-party library import to the top.
📍 Affects 1 file
docs/scripts/probe_grail.py#L24-L29(this comment)docs/scripts/probe_grail.py#L34-L38docs/scripts/probe_grail.py#L127-L127
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@docs/scripts/probe_grail.py` around lines 24 - 29, Move the local imports
from lines 24-29, the third-party imports from lines 34-38, and the third-party
import from line 127 in docs/scripts/probe_grail.py to the file’s import
section, ordered as standard-library, third-party, then local modules; remove
the original inline import statements while preserving existing usage.
Source: Coding guidelines
| print("\n=== [3/6] tiny synthetic round-trip ===") | ||
| g = torch.Generator().manual_seed(1234) | ||
| base = {"w1": torch.randn(100, 200, generator=g), "w2": torch.randn(50, 50, generator=g)} | ||
| current = {k: v.clone() for k, v in base.items()} | ||
| current["w1"].view(-1)[torch.randperm(100 * 200, generator=g)[:100]] = 0.5 | ||
| sparse, shapes, stats = compute_sparse_delta(current, base) | ||
| print(f"stats: total={stats['total_params']} changed={stats['nonzero_params']}") | ||
| assert stats["nonzero_params"] == 100, f"expected 100 changed, got {stats['nonzero_params']}" | ||
| reconstructed = apply_sparse_delta(base, sparse, shapes) | ||
| ok = all(torch.equal(current[k].cpu(), reconstructed[k].cpu()) for k in current) | ||
| print("tiny reconstruct OK:", ok) | ||
| assert ok |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | 🏗️ Heavy lift
Refactor script for logging, type hints, and function encapsulation.
As per coding guidelines, please update the script to adhere to the project's structural standards:
- Replace
printstatements withlogging.getLogger(__name__). - Add full type hints to all variables for
mypy-strictcompatibility. - Encapsulate the module-level execution logic into single-responsibility functions (target 5–20 lines).
Apply these changes throughout the file.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@docs/scripts/probe_grail.py` around lines 44 - 55, Refactor
docs/scripts/probe_grail.py by replacing all print calls with a module logger
obtained via logging.getLogger(__name__), adding explicit type annotations for
variables and relevant function values to satisfy mypy-strict, and moving
module-level execution into single-responsibility functions of roughly 5–20
lines. Preserve the existing probe behavior, assertions, and output information
while keeping the script entry point minimal.
Source: Coding guidelines
| os.makedirs(os.path.dirname(args.out), exist_ok=True) | ||
| with open(args.out, "w") as f: | ||
| json.dump(out_obj, f, indent=2) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
os.makedirs crashes on a bare-filename --out.
If --out is a filename with no directory component, os.path.dirname(args.out) returns "" and os.makedirs("", exist_ok=True) raises FileNotFoundError, discarding results after the full run completes. Guard the empty case.
🐛 Guard empty dirname
- os.makedirs(os.path.dirname(args.out), exist_ok=True)
+ out_dir = os.path.dirname(args.out)
+ if out_dir:
+ os.makedirs(out_dir, exist_ok=True)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| os.makedirs(os.path.dirname(args.out), exist_ok=True) | |
| with open(args.out, "w") as f: | |
| json.dump(out_obj, f, indent=2) | |
| out_dir = os.path.dirname(args.out) | |
| if out_dir: | |
| os.makedirs(out_dir, exist_ok=True) | |
| with open(args.out, "w") as f: | |
| json.dump(out_obj, f, indent=2) |
🧰 Tools
🪛 ast-grep (0.44.1)
[info] 227-227: use jsonify instead of json.dumps for JSON output
Context: json.dumps(steady, indent=2)
Note: [CWE-116] Improper Encoding or Escaping of Output.
(use-jsonify)
[warning] 225-225: File path is request-/variable-derived; validate and normalize to prevent path traversal.
Context: open(args.out, "w")
Note: [CWE-22] Improper Limitation of a Pathname to a Restricted Directory ('Path Traversal').
(open-filename-from-request)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@docs/scripts/s13_transfer_bench.py` around lines 225 - 227, Guard the
directory-creation step before writing the output in the transfer benchmark
flow: only call os.makedirs when os.path.dirname(args.out) is non-empty, while
preserving the existing file-writing behavior for both bare filenames and paths
with directories.
| try: | ||
| await asyncio.wait_for(task, timeout=30) | ||
| except (asyncio.TimeoutError, Exception): | ||
| task.cancel() |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
Broad except silently discards real shutdown errors.
except (asyncio.TimeoutError, Exception): catches everything (the TimeoutError in the tuple is redundant) and swallows the exception without logging it. Since this harness exists specifically to catch miner crashes, a genuine error here disappears silently.
As per coding guidelines, "No silent failures; log with context and re-raise when you can't handle errors."
🐛 Proposed fix
+ import logging as _logging
try:
await asyncio.wait_for(task, timeout=30)
- except (asyncio.TimeoutError, Exception):
+ except asyncio.TimeoutError:
+ task.cancel()
+ except Exception:
+ _logging.getLogger(__name__).exception("Miner task raised during shutdown")
task.cancel()📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| try: | |
| await asyncio.wait_for(task, timeout=30) | |
| except (asyncio.TimeoutError, Exception): | |
| task.cancel() | |
| import logging as _logging | |
| try: | |
| await asyncio.wait_for(task, timeout=30) | |
| except asyncio.TimeoutError: | |
| task.cancel() | |
| except Exception: | |
| _logging.getLogger(__name__).exception("Miner task raised during shutdown") | |
| task.cancel() |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@docs/scripts/s2_run_miner_mocked.py` around lines 46 - 49, Update the
task-waiting error handling around asyncio.wait_for in the harness so
asyncio.TimeoutError is handled separately, while other exceptions are logged
with shutdown context and re-raised. Preserve task cancellation on timeout and
ensure genuine miner failures are no longer silently swallowed.
Source: Coding guidelines
| manager = SGLangServerManager( | ||
| config=ServerConfig(model_path=MODEL, timeout_s=1800.0, env=env), | ||
| eval_config=eval_config, | ||
| ) | ||
|
|
||
| print(f"[s5] starting SGLang with {MODEL}", flush=True) | ||
| t0 = time.perf_counter() | ||
| await manager.__aenter__() | ||
| await manager.start_server() | ||
| t_start = time.perf_counter() - t0 | ||
| print(f"[s5] server up in {t_start:.1f}s at {manager.base_url}", flush=True) | ||
|
|
||
| results: dict[str, Any] = {"model": MODEL, "server_start_s": t_start, "cold": [], "warm": []} | ||
|
|
||
| async with httpx.AsyncClient() as client: | ||
| for mode, iters in (("cold", COLD_ITERS), ("warm", WARM_ITERS)): | ||
| for i in range(iters): | ||
| if mode == "cold": | ||
| n_files, gb = drop_page_cache(MODEL) | ||
| print(f"[s5] fadvise dropped {n_files} shards ({gb:.1f} GB)", flush=True) | ||
| t0 = time.perf_counter() | ||
| resp = await client.post( | ||
| f"{manager.base_url}/update_weights_from_disk", | ||
| json={"model_path": MODEL}, | ||
| timeout=1800.0, | ||
| ) | ||
| resp.raise_for_status() | ||
| elapsed = time.perf_counter() - t0 | ||
| results[mode].append(elapsed) | ||
| print( | ||
| f"[s5] reload {mode} #{i}: {elapsed:.2f}s " | ||
| f"(resp: {resp.json() if resp.headers.get('content-type', '').startswith('application/json') else resp.status_code})", | ||
| flush=True, | ||
| ) | ||
|
|
||
| await manager.__aexit__(None, None, None) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Server/process leak if anything raises between start and __aexit__.
manager.__aenter__()/start_server() run with no try/finally guaranteeing manager.__aexit__(None, None, None) (line 113) executes. If client.post(...) inside the loop raises (timeout, HTTP error via raise_for_status()), the SGLang server process/GPU memory is orphaned — exactly the "escaped singularity children" failure class this PR documents hitting in production (see docs/sbatch/s7_delta_e2e.sbatch lines 42-49, "escaped processes from a previous run... SGLang held 126.9GB on GPU0").
🔒 Proposed fix: guarantee cleanup with try/finally
print(f"[s5] starting SGLang with {MODEL}", flush=True)
t0 = time.perf_counter()
await manager.__aenter__()
- await manager.start_server()
- t_start = time.perf_counter() - t0
- print(f"[s5] server up in {t_start:.1f}s at {manager.base_url}", flush=True)
-
- results: dict[str, Any] = {"model": MODEL, "server_start_s": t_start, "cold": [], "warm": []}
-
- async with httpx.AsyncClient() as client:
- for mode, iters in (("cold", COLD_ITERS), ("warm", WARM_ITERS)):
- ...
-
- await manager.__aexit__(None, None, None)
+ try:
+ await manager.start_server()
+ t_start = time.perf_counter() - t0
+ print(f"[s5] server up in {t_start:.1f}s at {manager.base_url}", flush=True)
+
+ results: dict[str, Any] = {
+ "model": MODEL, "server_start_s": t_start, "cold": [], "warm": [],
+ }
+
+ async with httpx.AsyncClient() as client:
+ for mode, iters in (("cold", COLD_ITERS), ("warm", WARM_ITERS)):
+ ...
+ finally:
+ await manager.__aexit__(None, None, None)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| manager = SGLangServerManager( | |
| config=ServerConfig(model_path=MODEL, timeout_s=1800.0, env=env), | |
| eval_config=eval_config, | |
| ) | |
| print(f"[s5] starting SGLang with {MODEL}", flush=True) | |
| t0 = time.perf_counter() | |
| await manager.__aenter__() | |
| await manager.start_server() | |
| t_start = time.perf_counter() - t0 | |
| print(f"[s5] server up in {t_start:.1f}s at {manager.base_url}", flush=True) | |
| results: dict[str, Any] = {"model": MODEL, "server_start_s": t_start, "cold": [], "warm": []} | |
| async with httpx.AsyncClient() as client: | |
| for mode, iters in (("cold", COLD_ITERS), ("warm", WARM_ITERS)): | |
| for i in range(iters): | |
| if mode == "cold": | |
| n_files, gb = drop_page_cache(MODEL) | |
| print(f"[s5] fadvise dropped {n_files} shards ({gb:.1f} GB)", flush=True) | |
| t0 = time.perf_counter() | |
| resp = await client.post( | |
| f"{manager.base_url}/update_weights_from_disk", | |
| json={"model_path": MODEL}, | |
| timeout=1800.0, | |
| ) | |
| resp.raise_for_status() | |
| elapsed = time.perf_counter() - t0 | |
| results[mode].append(elapsed) | |
| print( | |
| f"[s5] reload {mode} #{i}: {elapsed:.2f}s " | |
| f"(resp: {resp.json() if resp.headers.get('content-type', '').startswith('application/json') else resp.status_code})", | |
| flush=True, | |
| ) | |
| await manager.__aexit__(None, None, None) | |
| manager = SGLangServerManager( | |
| config=ServerConfig(model_path=MODEL, timeout_s=1800.0, env=env), | |
| eval_config=eval_config, | |
| ) | |
| print(f"[s5] starting SGLang with {MODEL}", flush=True) | |
| t0 = time.perf_counter() | |
| await manager.__aenter__() | |
| try: | |
| await manager.start_server() | |
| t_start = time.perf_counter() - t0 | |
| print(f"[s5] server up in {t_start:.1f}s at {manager.base_url}", flush=True) | |
| results: dict[str, Any] = {"model": MODEL, "server_start_s": t_start, "cold": [], "warm": []} | |
| async with httpx.AsyncClient() as client: | |
| for mode, iters in (("cold", COLD_ITERS), ("warm", WARM_ITERS)): | |
| for i in range(iters): | |
| if mode == "cold": | |
| n_files, gb = drop_page_cache(MODEL) | |
| print(f"[s5] fadvise dropped {n_files} shards ({gb:.1f} GB)", flush=True) | |
| t0 = time.perf_counter() | |
| resp = await client.post( | |
| f"{manager.base_url}/update_weights_from_disk", | |
| json={"model_path": MODEL}, | |
| timeout=1800.0, | |
| ) | |
| resp.raise_for_status() | |
| elapsed = time.perf_counter() - t0 | |
| results[mode].append(elapsed) | |
| print( | |
| f"[s5] reload {mode} #{i}: {elapsed:.2f}s " | |
| f"(resp: {resp.json() if resp.headers.get('content-type', '').startswith('application/json') else resp.status_code})", | |
| flush=True, | |
| ) | |
| finally: | |
| await manager.__aexit__(None, None, None) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@docs/scripts/s5_sglang_reload_bench.py` around lines 78 - 113, Wrap the
server lifecycle and benchmark request loop in a try/finally so
manager.__aexit__(None, None, None) always runs after manager.__aenter__()
succeeds, including when start_server(), client.post(), or
resp.raise_for_status() raises. Preserve the existing startup, measurement, and
result-recording behavior while guaranteeing SGLang process cleanup.
| state = load_model_state_dict(MODEL) | ||
| for k in state: | ||
| state[k] = state[k].clone() # break mmap so staging write reads RAM, not CPFS | ||
| say("state dict materialized in RAM") |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Add an explicit check for missing weights before iterating state.
load_model_state_dict can return None when no weights are found in MODEL; this iterates it unconditionally, so a misconfigured/missing model path fails with an opaque TypeError instead of a clear error.
🐛 Proposed fix
state = load_model_state_dict(MODEL)
+ if state is None:
+ raise FileNotFoundError(f"No model weights found under {MODEL}")
for k in state:
state[k] = state[k].clone() # break mmap so staging write reads RAM, not CPFS📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| state = load_model_state_dict(MODEL) | |
| for k in state: | |
| state[k] = state[k].clone() # break mmap so staging write reads RAM, not CPFS | |
| say("state dict materialized in RAM") | |
| state = load_model_state_dict(MODEL) | |
| if state is None: | |
| raise FileNotFoundError(f"No model weights found under {MODEL}") | |
| for k in state: | |
| state[k] = state[k].clone() # break mmap so staging write reads RAM, not CPFS | |
| say("state dict materialized in RAM") |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@docs/scripts/s6_disk_segments_bench.py` around lines 113 - 116, Validate the
result of load_model_state_dict(MODEL) before the `for k in state` loop, and
raise a clear error when it is None or otherwise indicates missing weights.
Preserve the existing cloning and RAM-materialization behavior for valid state
dictionaries.
- s9_real_delta_bench: v2 pathway (gather to rank0 GPU) + fp32_shadow prealloc option (S9_PREALLOC_OPTIM), optional grail production sparse_codec (v2/v3 zstd via S9_DISK_COMPRESS), miner unflatten for v2 flat-linear -> 2D COO, round-scoped memory peaks (reset each round for baseline delta measurement) - s22 multistep delta bench: K=1..4 change ratio measurement per round - s13 preflight + moto helper: local S3 mock (moto server) for pre-cluster tests - s0..s8 comms/publish/consume/rotate benches: infra pathway smoke tests - probe_grail / parse_sync_metrics: sync metric analysis helpers - Slurm sbatch launchers for AIS + on-prem clusters (docs/sbatch/) No changes to upstream tracked code. All additions live under docs/scripts/ and docs/sbatch/ and are opt-in bench/instrumentation.
61b3a52 to
8b8436a
Compare
1 participant
Summary by CodeRabbit
New Features
Bug Fixes