CORE: Add team cache LFU eviction policy - #1351
Draft
bwestheimer wants to merge 30 commits into
Draft
bwestheimer wants to merge 30 commits into
bwestheimer wants to merge 30 commits into
Conversation
bwestheimer
force-pushed
the
bwestheimer/pub-team-cache-lfu
branch
from
September 2, 2026 18:14
ff2e9a2 to
ba9b51c
Compare
Collaborator
|
🤖 CI Triage Agent — TL;DR: The Full analysisSummary: Compilation of Root cause: In Implicated commit: File: Suggested fix: Emit the debug message before freeing, so the pointer is not used after kh_destroy(ucc_team_cache_map, (ucc_team_cache_map_t *)cache->table);
ucc_spinlock_destroy(&cache->lock);
ucc_debug("ucc_team_cache destroyed: %p", cache_ptr);
ucc_free(cache);(The separate Related: none found. |
Collaborator
|
🤖 CI Triage Agent — TL;DR: The Full analysisSummary: Root cause: This PR ( Implicated commit: File:
Suggested fix: Update the CUDA TL to use the new accessor macro instead of the removed struct field. In all three locations replace: UCC_TL_CORE_TEAM(team)->topowith: UCC_TEAM_TOPO(UCC_TL_CORE_TEAM(team))e.g. Related: none (no existing issue/PR found in this repo referencing the migration). |
bwestheimer
changed the title
bwestheimer
force-pushed
the
bwestheimer/pub-team-cache-lfu
branch
from
September 2, 2026 20:41
ba9b51c to
6da5ec5
Compare
Collaborator
|
🤖 CI Triage Agent — TL;DR: The Full analysisSummary: clang-tidy-17 in the ROCm Lint job reported a Root cause: In the LFU/usage-based branch, Implicated commit: File: Suggested fix: Make the non-null nature of
Avoid a blanket Related: PR #1351 (
|
bwestheimer
force-pushed
the
bwestheimer/pub-team-cache-lfu
branch
2 times, most recently
from
September 3, 2026 14:04
dc22548 to
0c138f7
Compare
bwestheimer
force-pushed
the
bwestheimer/pub-team-cache-lfu
branch
from
September 3, 2026 14:39
0c138f7 to
a0b09b3
Compare
Add an 'embedded' flag to ucc_service_coll_req_t. When set, ucc_service_coll_finalize skips ucc_free(req) so the request can live in caller-owned storage instead of being heap allocated. The team agreement vote embeds its request directly in ucc_team_t and relies on this. Add ucc_service_allreduce_ctx, which runs an allreduce over a pre-materialized subset of context endpoints using the context-level service team. The existing ucc_service_allreduce resolves subset ranks through team->ctx_map, which is not yet built while a team is being created. The new entry point maps subset indices straight to context ranks and routes through ctx->service_team, so the agreement vote can run before ctx_map exists.
Introduce ucc_team_cache_identity_t, the normalized key type used to recognize a team by its membership, together with its build, compare and free helpers and the cacheability policy check. No cache container is added here. ucc_team_cache_identity_build materializes members[] by evaluating the ep_map, so an identity never aliases caller-owned storage such as an ep_map callback closure or a user array. Any ep_map style (CB, ARRAY, STRIDED, FULL) describing the same membership yields the same identity. The FNV-1a hash covers membership only (size, self_ep, members[]). ext_id is compared separately and deliberately not hashed, so teams with the same membership but different external communicator ids land in the same bucket; PR3a introduces the chaining that lets them coexist there. ucc_team_cache_identity_equal_membership provides the ext_id-agnostic compare that the derived-team path uses later. instance_cookie is 0 until the agreement vote stamps it in PR1f. ucc_team_cache_is_cacheable rejects teams that set optional behavioral params (ORDERING, OUTSTANDING_COLLS, SYNC_TYPE, P2P_CONN, MEM_PARAMS), since those are not part of the identity and a reuse could silently change semantics.
Add the team-cache container as a standalone data structure. It owns a khash uint64 -> ucc_team_t* bucket table plus four intrusive lists (live, dormant, reserved, pending_destroy), all protected by a single spinlock, together with the size/capacity bookkeeping and the hit/miss/insert/ eviction counters. The lifecycle API is refcount based: insert admits a team as DORMANT, get adopts a DORMANT team into LIVE with refcount++, and put releases a LIVE team back to DORMANT once the refcount reaches zero. Insert is a no-op at capacity, which leaves the team uncached but fully functional, and it also skips a team whose bucket is already occupied by a duplicate identity or a hash collision. Lookup only ever returns a DORMANT team, so a LIVE or RESERVED entry is never handed out twice. The eviction victim picker is FIFO, that is the dormant list head, which is the oldest insert; LFU selection is added later. The registry helpers do the list surgery for the state transitions and table_erase removes a team from the bucket table. ucc_team_t gains the fields the cache API needs: refcount, cache_identity, cache_link, cache_state, cache_pending_insert and cache_local_action. This commit adds no wiring into team create/destroy and no agreement vote; the lifecycle hooks and the vote machinery are separate changes.
Add the five UCC_TEAM_CACHE_* context config knobs, create and destroy the per-context cache, and hook the cache into ucc_team_create_post, ucc_team_create_test and ucc_team_destroy so a destroyed team is retained as DORMANT and re-adopted by a later create with identical membership. Each rank classifies its create as a hit or a miss locally, with no cross-rank agreement. EXACT_REUSE without that agreement is safe as long as team scopes never overlap, that is, no rank belongs to two simultaneously created teams with the same membership. The user guide documents this restriction. The agreement vote that lifts it lands in a follow-up, together with the handling for the UCC_TEAM_CACHE_AGREE and UCC_TEAM_CACHE_MISS_TEARDOWN states defined here. UCC_TEAM_CACHE_ENABLE defaults to n, so the feature is entirely opt-in.
Add gtest integration coverage that drives real create/destroy/recreate cycles through UccJob: dormant re-adoption, the disabled knob leaving no cache, eviction with team-id release, id-pool headroom under dormant teams, and the dump-stats and disable-linear-check knobs. Add a multi-rank ucc_test_mpi suite, gated on UCC_TEAM_CACHE_CORRECTNESS_TESTS, covering dormant-reuse hit counts, safety of a caller ep_map callback freed after the team is cached, and singleton teams. Document the team cache and its knobs in the user guide, including the non-overlapping team scope restriction for reuse without agreement.
Members of a cacheable team create classify the create as a cache hit or a miss from their own cache contents, and those contents can diverge, for example after an eviction on one rank only. A create where some ranks re-adopt a dormant team while others build a fresh one does not progress. Reconcile the per-rank action with a UCC_OP_BAND allreduce over a small vote buffer before any rank skips the address exchange. The buffer carries a prepared flag, (value, ~value) equality pairs for the action, key and instance cookies, and team rank 0's proposed cookie. Any disagreement degrades the result to MISS, so all members fall back to a fresh build. This replaces the UCC_TEAM_CACHE_AGREE and UCC_TEAM_CACHE_MISS_TEARDOWN stubs. A rejected reuse candidate is torn down and rebuilt in place, which keeps the handle the caller already holds valid. Add UCC_TEAM_CACHE_AGREEMENT (default y) to control the vote.
Add gtests for the vote helpers: a unanimous EXACT_REUSE agrees and distributes rank 0's cookie, a single non-preparing rank degrades the result to MISS, a cookie or parent-cookie mismatch also degrades to MISS, and next_cookie stays monotonic and never returns 0. Add two MPI tests. overlap_agreement builds overlapping subcommunicator sets that, with UCC_TEAM_CACHE_MAX_SIZE=2, evict divergently across ranks; without the vote this deadlocks, with it the members reconcile to a fresh build. nonblocking_create_post delays rank 0 and checks that the peers' ucc_team_create_post returns rather than blocking on the vote. Document UCC_TEAM_CACHE_AGREEMENT and drop the overlapping-scope restriction note, which the vote removes.
Add test/mpi/run_cache_equivalence.sh, which runs ucc_test_mpi twice over the same team set (world, half, odd_even, reverse) and collective set (barrier, allreduce, bcast, alltoall, allgather) -- once with UCC_TEAM_CACHE_ENABLE=y and once with n -- using the test's built-in per-collective correctness checks as the equivalence oracle rather than diffing outputs across runs. Both passes also set UCC_TEAM_CACHE_CORRECTNESS_TESTS=y so the cache-on pass exercises real reuse and derivation. That pass additionally asserts the correctness suite did not skip itself, so a silently inert cache cannot masquerade as equivalent by never touching the cache at all. Wire both the team-cache correctness suite and the equivalence pass into .ci/scripts/run_tests_ucc_mpi.sh, and distribute the script via EXTRA_DIST (CI invokes it directly from the source tree).
Introduce ucc_team_artifacts_t to hold ctx_map, ctx_ranks and topo behind a refcount and a spinlock, so a later change can share them between a cached team and the teams derived from it. Every team uses an embedded inline holder (heap=0, refcount=1), so there is no sharing and no behavioral change yet: - ucc_team_artifacts_init_inline zero-inits the holder, sets heap=0 and refcount=1, and initializes the spinlock. - ucc_team_artifacts_put decrements the refcount under the lock; at zero it releases topo and ctx_ranks, and frees the struct only when heap=1. This replaces the topo and ctx_ranks teardown in the destroy path.
Replace every direct read of core_team->ctx_map, core_team->ctx_ranks and core_team->topo in the TL and CL components with UCC_TEAM_CTX_MAP, UCC_TEAM_CTX_RANKS and UCC_TEAM_TOPO, so all access goes through the artifacts holder. Mechanical one-liners with no behavioral change.
ucc_topo_t fills its sbgp, socket, numa, node and node-leader state lazily on first use. That is fine while a topo belongs to exactly one team, but a later change lets a cached team share its topo with the teams derived from it, and concurrent first-touch fills from several teams under UCC_THREAD_MULTIPLE would race on those writes. Add ucc_topo_prepare_shared, which walks every sbgp type and the all_* socket/numa/node arrays (plus node leaders for multi-node teams) so the topo is fully built and thereafter read-only. Per-sbgp failures record a terminal status and are not retried, so only an allocation failure in the retryable all_* paths is treated as fatal. Call it from ucc_team_create_cls, but only for teams that are candidates for the cache; ordinary teams keep the existing lazy behavior. On failure the team is still created and usable -- it is simply dropped from the cache candidate set rather than being shared with a half-built topo, so no user-visible operation aborts.
Today a hash collision in the team cache silently leaves the second team uncached: the bucket is already occupied, so the insert is skipped and that communicator never benefits from reuse. Two teams with the same membership but different external ids collide by construction, since the cache key hashes membership only. Chain the teams that share a bucket through a new bucket_link ring on ucc_team_t. Insert appends to the chain in collective order and still refuses an exact-identity duplicate, erase unlinks the entry and promotes the next sibling when the head leaves, and lookup walks the chain so each team is independently reachable by its own external id. Also add the update_id vtable stub to ucc_base_team_iface_t, defaulted to NULL by UCC_BASE_IFACE_DECLARE. It is scaffolding for the team-cache re-seat path that re-seats a CL team id and tag domain in place.
Add a gtest that inserts three stub teams with identical membership and distinct external ids, so all three land in one hash bucket. It checks that every insert grows the cache, that each team is independently reachable by its own external id, and that erasing the chain head, a non-head sibling and the last entry each collapse the ring correctly.
Add the derived-team fast path so a create that duplicates an existing
LIVE team's membership borrows that team's shared artifacts holder
(ctx_map + topo) instead of rebuilding it, which is the MPI_Comm_dup
shape. The derived team draws its own team id and tag domain, and skips
the ADDR_EXCHANGE phase entirely.
Key additions:
- ucc_team_artifacts_{alloc,get}: heap holder lifecycle for the state a
derived team shares with its parent
- UCC_TEAM_CACHE_ACTION_DERIVED_FROM_LIVE and ucc_team_cache_lookup_live,
which matches on membership only since a child's ext_id differs
- ucc_team_{can_derive_from,init_derived}: derived team setup, in both
the agreement and the direct create paths
- UCC_TEAM_CACHE_DERIVED config knob, default on
A team that loses the agreement vote drops the parent pin and rebuilds
as an ordinary full team.
gtest: - lookup_live_returns_live_sibling: lookup_live skips the DORMANT team in a same-membership chain and returns the LIVE sibling, then NULL once that sibling goes dormant - derived_coexist_interleaved: a parent and its derived team, both LIVE with identical membership, run interleaved allreduce+bcast in opposite per-rank orders; distinct team ids keep the tag domains isolated MPI: - dup_coexist_derived, with and without external ids: the same coexistence check across a real job, asserting the derived path fired and that the two teams share one artifacts holder - derived_reuse: a dormant derived team is re-adopted by exact identity on every iteration while its parent stays live - derived_exact_rebuild: a dormant derived team that loses the cross-rank vote is rebuilt as a full team, with is_derived cleared and a working ctx_map The derived tests skip when UCC_TEAM_CACHE_DERIVED is off.
Add the knob to the team-cache table and a section describing when a create derives from a live same-membership team, what it borrows, and how the borrowed state is released.
A cached team is keyed on its membership *and* its external id, which for MPI is the communicator context id. Workloads that never reuse a context id draw a fresh one on every create/free cycle over the same ranks, so the id drifts and the exact-identity lookup misses every time, even though a perfectly good dormant team of that membership is sitting in the cache. UCC_TEAM_CACHE_RESEAT=y recovers reuse there. When the exact lookup misses, the cache is searched a second time for a dormant derived team of identical membership, ignoring the external id. A match is re-adopted and re-seated: its team id, and with it the tag and sequence-number domain of its service team and of every CL and TL team beneath it, moves to the caller's new external id. Only derived teams are eligible, since only they hold borrowed artifacts that make the move cheap relative to a full build. Ignoring ext_id in the lookup costs the agreement vote its key lane. That lane carried ext_id, which is what proved every rank had selected the same cached entry; two ranks holding different dormant teams of identical membership would now agree on a value that no longer distinguishes them. The identity therefore gains an instance_cookie, stamped when a team is adopted and carried in the vote in place of the id, so either all members re-seat the same team or none do. A reseat candidate is booked optimistically as a hit against a lookup that actually missed, so the rollback paths rebook it as a miss if the vote is lost or the post fails. The knob is experimental and off by default. It requires UCC_TEAM_CACHE_DERIVED=y; with derived teams off there is nothing eligible to re-seat.
Re-seating a dormant derived team moves its team id, so every component that embedded the old id in a tag or sequence domain has to be told. Add an update_id team hook, implement it for the CL layer, and call it from ucc_team_reseat_id for the service team and for each CL team. The id has to be pushed all the way down. A core team that answered to the new id while a TL team beneath it still addressed the retired one would put two logically distinct teams into a single tag domain, which is exactly the aliasing this path exists to avoid. The hook is optional on both paths. A component with no settable tag-domain concept simply does not provide it: TL/NCCL is the case that matters, and it is why NCCL-backed teams take EXACT_REUSE only and are excluded from RESEAT_DERIVED.
Cover ucc_team_reseat_id and the reuse path it serves. gtest: reseat_id_rewrites_all_ids drives the id rewrite directly and asserts that team->id, bp.id and cache_identity.ext_id all move to the new tagged id while the membership hash the bucket is keyed on does not. A stub team has no service team and no CL teams, so the test also exercises the NULL guards on both fanout hooks. MPI: derived_reuse[drift] recreates a derived team whose external id drifts every iteration, so only a membership-match re-adopt can hit. It self-skips unless UCC_TEAM_CACHE_RESEAT=y, so run_tests_ucc_mpi.sh gains a second correctness pass with the knob on. Without that leg the entire reseat path would ship untested. Document the knob in the user guide beside the derived-team section it depends on.
Extend the eviction policy enum with two new values: - UCC_TEAM_CACHE_EVICTION_LFU (2): evict the dormant team with the smallest seq_num (fewest collectives served), keeping hot teams alive. - UCC_TEAM_CACHE_EVICTION_LRU (3): accepted alias for LFU; UCC tracks collective count rather than wall-clock recency. Add UCC_TEAM_CACHE_EVICTION_IS_USAGE_BASED() helper macro and update ucc_team_cache_pick_lru_victim() to walk the dormant list and select the min-seq_num entry under both LFU/LRU, with tie-break by list position (earlier/oldest wins). The FIFO path is unchanged. Update ucc_team_cache_eviction_names[] with "lfu" and "lru" entries and update the TEAM_CACHE_EVICTION config-table description to document them.
Expand evict_victim_selection to run FIFO, LFU, and LRU sub-cases in a single loop. For FIFO the oldest-inserted team is the victim regardless of seq_num; for LFU/LRU the min-seq_num team wins, with a tie-break case that confirms the earlier dormant-list entry is preferred. Add vote_reseat_different_cookie_misses: two ranks vote EXACT_REUSE for the same membership key but with different per-instance cookies, which breaks the cookie equality lane and forces a global MISS.
ucc_service_allreduce_ctx asserted that the context service team was non-NULL. That team is only created when the context carries an OOB and UCC_INTERNAL_OOB permits it, and its creation is deliberately non-fatal, so NULL is a supported runtime state rather than a caller error - ucc_service_coll_req_init already treats it as one and falls back. Release builds define NDEBUG, so the assert compiled out and left a NULL dereference in UCC_TL_TEAM_IFACE. Reproduced with UCC_INTERNAL_OOB=0 and UCC_TEAM_CACHE_ENABLE=y as a SIGSEGV in ucc_service_allreduce_ctx called from ucc_team_create_post. Return UCC_ERR_NOT_SUPPORTED instead. There is deliberately no fallback to the per-team service team: this collective addresses context ranks directly, which a team-scoped service team cannot do. Also realign the assignment block that req->embedded widened, and order the initialized declaration in ucc_service_coll_finalize first. (cherry picked from commit 4d564b1400fc46ab1cda4449b1a4eb9ab75e3010)
The cross-rank agreement vote runs over the context service team. That team is created after the cache is initialized, and its creation is not fatal, so a context can reach team creation with caching enabled, agreement requested, and no team to vote over. Drop the cache in that case rather than continuing without agreement. Reuse without the vote is safe only when team scopes never overlap, so disabling agreement automatically would trade a performance feature for a correctness risk the user never accepted, and the resulting failure is a hang rather than a slowdown. The condition derives from configuration alone, so every rank reaches the same decision and the gate stays uniform. UCC_TEAM_CACHE_AGREEMENT=n still caches without the vote, which the warning names; that path is unaffected by this check. (cherry picked from commit 4ec16e4cb57b5c7f034724b4fb0365daa17cfd8e)
A rank whose team cache failed to allocate was left running without one, while its peers kept theirs. Such a rank skips the agreement vote in ucc_team_create_post and its peers wait on a member-scoped allreduce that never completes. Fail the context instead, so whether a rank caches is decided by configuration alone, which every rank shares.
ucc_team_destroy_single finalized bp.params.oob whenever the context had a service team. When ucc_internal_oob_init itself failed, that field still held the caller's OOB, and the destroy on the create error path freed the caller's coll_info. Record ownership in team->internal_oob after a successful init and finalize only what UCC installed.
When the agreement allreduce returned an error the team stayed in UCC_TEAM_CACHE_AGREE: a RESERVED candidate never returned to the dormant list, a derived-from-live parent pin was never released, ucc_team_destroy rejected the handle, and a repeated ucc_team_create_test re-tested the finalized request. Split the post-time rollback into a reusable release of the reserved candidate and use it from the vote path too. A shell this create allocated moves to a terminal UCC_TEAM_CREATE_FAILED state that ucc_team_destroy accepts; a handle that named a cached team is handed back and destroy refuses it, as it does for any dormant or reserved team. Destroying such a shell also reached ucc_coll_score_free_map with a NULL map, which the post-time failure path already did; guard it.
A cached team whose teardown failed terminally had its id returned to the pool while the CL/TL teams that were not destroyed could still match traffic on it. Leak the id together with the rest of the state, so the failure stays a bounded leak instead of letting two teams share a tag domain.
Evicted teams keep their team id until their asynchronous destroy completes, but the pending-destroy list was only progressed from the next admission, eviction or context drain. A teardown that stayed in progress could hold ids inside the pool headroom with nothing driving it. Register a throttled context progress callback that drives the list when it is non-empty, and progress it once before a team posts its id-pool allreduce so finished evictions return their ids first.
The callback-lifetime test freed the poisoned cb_ctx box and immediately allocated a same-sized replacement with a valid magic. An allocator that returns the same address would make a retained stale pointer look valid. Keep the poisoned box allocated until the reuse checks finish.
The dormant-derived lookup matched on membership only, so a create without UCC_TEAM_PARAM_FIELD_ID could re-seat a candidate to id 0, and a derived team holding a pool id could be re-seated to an external id, orphaning the pool id. Re-seating moves a team between two external ids: decline a request without one and skip candidates whose id came from the pool.
bwestheimer
force-pushed
the
bwestheimer/pub-team-cache-lfu
branch
from
September 28, 2026 17:38
bf39db4 to
31f213c
Compare
This branch has not been deployed
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
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.
What
Add
lfuandlruas values forUCC_TEAM_CACHE_EVICTION. Under LFU/LRU, the eviction victim is the dormant team with the lowestseq_num(fewest collectives served) rather than the oldest insertion. LRU is an accepted alias since UCC has no wall-clock recency.Why
FIFO eviction can evict a frequently-reused team in favour of a rarely-used one. Usage-count based eviction keeps hot teams alive when the cache is at capacity.
Depends on #1350.