Simplify CPU bucket cache locking - #14
Conversation
|
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 Why hold:
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. |
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
f47642b to
9615401
Compare
All done.
No actor-level lock guards build_cpu_bucket_cache / run_sync_session: train actors must stay default sync Ray actors (no max_concurrency), so Ray serializes their calls. The cache's own lock protects its internal state. |
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