CORE/TEAM_CACHE: Add LFU eviction policy - #1351
Draft
bwestheimer wants to merge 23 commits into
Draft
Conversation
The alloc path builds an id as i * 64 + pos, where i is the pool word index and pos is in [1, 64], so bit (pos - 1) of word i encodes that id. The release path instead computed map_pos = id / 64, which for any id that is a multiple of 64 names the next word: id 64 was released as word 1 bit 63, the slot for id 128. That leaks id 64 and hands out id 128 twice. Use map_pos = (id - 1) / 64 so release is the exact inverse of alloc. Expose the two pool bit helpers so the boundary behavior can be tested directly, and add a gtest that round-trips ids across every word boundary and checks no residue is left in the pool.
Move the ucc_assert(id >= 1) guard to the top of ucc_team_id_pool_set_bit so invalid ids are rejected before the (id - 1) subtraction; a value of 0 or INT_MIN would cause signed overflow UB before the assert fires.
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.
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.
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
1 participant
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.