Skip to content

Simplify CPU bucket cache locking - #14

Open
TianyeGGBond wants to merge 2 commits into
rlops:zhenyu/m11-mvp-testfrom
TianyeGGBond:tianye/simplify-cache-locking
Open

Simplify CPU bucket cache locking#14
TianyeGGBond wants to merge 2 commits into
rlops:zhenyu/m11-mvp-testfrom
TianyeGGBond:tianye/simplify-cache-locking

Conversation

@TianyeGGBond

Copy link
Copy Markdown

Context

The F4 CPU bucket cache path currently has two layers of locking around the same state:

  • CPUBucketCache owns an internal threading.Lock for _buckets and _cache_ready_step.
  • MegatronTrainRayActor also creates _cache_lock and wraps both build_cpu_bucket_cache and run_sync_session with it.

In the current deployment model, train actors are created as default synchronous Ray actors without max_concurrency, and these methods are not async. That means build_cpu_bucket_cache and run_sync_session are already serialized by Ray on the same actor. The actor-level lock is therefore redundant, and its comments make the implementation look more concurrent than it is.

What changed

  • Removed MegatronTrainRayActor._cache_lock initialization.
  • Removed the outer actor-level lock around build_cpu_bucket_cache.
  • Removed the outer actor-level lock around run_sync_session.
  • Kept CPUBucketCache's internal lock so the cache object remains self-contained and protects its own state.
  • Shortened F4 cache docstrings and comments to describe the API behavior without carrying review-history or over-defensive concurrency rationale in code.

Why

This keeps the locking model to one layer while preserving behavior. The cache still publishes and reads _cache_ready_step through its own methods, but the train actor no longer adds a redundant critical section around synchronous Ray actor methods. This is closer to the style in miles main: rely on the actor execution model where it applies, and keep local state protection inside the small helper object.

Validation

  • python -m py_compile miles/backends/megatron_utils/actor.py miles/backends/megatron_utils/update_weight/cpu_bucket_cache.py
  • git diff --check
  • git grep -n _cache_lock -- miles/backends/megatron_utils/actor.py miles/backends/megatron_utils/update_weight/cpu_bucket_cache.py returns no matches

@zhenyulincs

Copy link
Copy Markdown

Reviewed — the premise checks out, but please hold this until #34 lands, then rebase. Not approving yet.

Verified the core claim: no creation site of MegatronTrainRayActor passes max_concurrency (checked the whole repo), so the actor is a default sync Ray actor and build_cpu_bucket_cache / run_sync_session are already serialized per-call by Ray. _cache_lock has no other users beyond the two sites this PR removes — it is genuinely uncontended today. The change also stays entirely inside port-added F4 code; no upstream miles lines are touched.

Why hold:

  1. Direct conflict with fix(miles): supersession guard in run_sync_session — sync latest when requested step is stale #34 (approved bugfix). fix(miles): supersession guard in run_sync_session — sync latest when requested step is stale #34 adds its supersession guard inside the with self._cache_lock: block in run_sync_session — the exact block this PR deletes and dedents. Whichever merges second hits a manual conflict. Bugfix should win the race: land fix(miles): supersession guard in run_sync_session — sync latest when requested step is stale #34 first, rebase this on top.

  2. Please make the serialization assumption explicit when you rebase. After removal, correctness rests on "no train-actor creation site ever adds max_concurrency" — an invariant that currently would fail silently if violated (the rlix-side coordinator and pipeline actors already use max_concurrency=4/32, so the pattern is nearby). A one-line note on the actor class or at the creation site — e.g. "train actors must remain default sync actors (no max_concurrency); F4 cache-method atomicity relies on Ray's per-actor call serialization" — keeps the invariant visible in code instead of living only in this PR's description. The cache's own internal lock stays, which is the right layer for the remaining state.

With those two addressed (rebase over #34 + written-down invariant), happy to approve — the cleanup itself is sound and the two-layer locking genuinely over-states the concurrency of this actor.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants