You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
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.
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.
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.
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Context
The F4 CPU bucket cache path currently has two layers of locking around the same state:
CPUBucketCacheowns an internalthreading.Lockfor_bucketsand_cache_ready_step.MegatronTrainRayActoralso creates_cache_lockand wraps bothbuild_cpu_bucket_cacheandrun_sync_sessionwith it.In the current deployment model, train actors are created as default synchronous Ray actors without
max_concurrency, and these methods are notasync. That meansbuild_cpu_bucket_cacheandrun_sync_sessionare 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
MegatronTrainRayActor._cache_lockinitialization.build_cpu_bucket_cache.run_sync_session.CPUBucketCache's internal lock so the cache object remains self-contained and protects its own state.Why
This keeps the locking model to one layer while preserving behavior. The cache still publishes and reads
_cache_ready_stepthrough 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.pygit diff --checkgit grep -n _cache_lock -- miles/backends/megatron_utils/actor.py miles/backends/megatron_utils/update_weight/cpu_bucket_cache.pyreturns no matches