diff --git a/.ci/scripts/run_tests_ucc_mpi.sh b/.ci/scripts/run_tests_ucc_mpi.sh index 4fd42cf1ecf..5f78d733926 100755 --- a/.ci/scripts/run_tests_ucc_mpi.sh +++ b/.ci/scripts/run_tests_ucc_mpi.sh @@ -202,6 +202,25 @@ for MT in "" "-T"; do echo "INFO: UCC MPI unit tests (CL/HIER+2step bcast) ... DONE" done +# Team-cache correctness tests: exercise dormant reuse and eviction under +# pressure. UCC_TEAM_CACHE_MAX_SIZE=2 keeps the cache tiny so the overlapping +# subcommunicator test forces divergent per-rank eviction, the case the +# cross-rank agreement must reconcile. +echo "INFO: UCC team-cache correctness tests (world,half,odd_even) ..." +# shellcheck disable=SC2046,SC2086 # MPI argument fragments intentionally word-split. +cache_args=" -x UCC_TEAM_CACHE_ENABLE=y -x UCC_TEAM_CACHE_MAX_SIZE=2 -x UCC_TEAM_CACHE_CORRECTNESS_TESTS=y " +mpirun $(mpi_params $PPN) $ucx_tls_no_cuda_ipc $cache_args $EXE -c barrier -t world,half,odd_even +echo "INFO: UCC team-cache correctness tests (world,half,odd_even) ... DONE" + +# Team-cache enabled-vs-disabled equivalence pass. The script runs ucc_test_mpi +# with cache on, then off, using its built-in per-collective correctness checks +# as the equivalence oracle. +echo "INFO: UCC team-cache enabled-vs-disabled equivalence test (8 ranks) ..." +# shellcheck disable=SC2086 +MPIRUN="$(command -v mpirun)" EXE="${EXE%% *}" \ + bash "${SCRIPT_DIR}/../../test/mpi/run_cache_equivalence.sh" 8 +echo "INFO: UCC team-cache enabled-vs-disabled equivalence test (8 ranks) ... DONE" + end=`date +%s` echo Tests took $((end - start)) seconds diff --git a/docs/user_guide.md b/docs/user_guide.md index dbaccbd90f0..84030fbe590 100644 --- a/docs/user_guide.md +++ b/docs/user_guide.md @@ -338,6 +338,101 @@ $ UCC_COLL_TRACE=INFO srun ./c/mpi/collective/osu_allreduce -i 1 -x 0 -d cuda -m [1678205653.810705] [node_name:903 :0] ucc_coll.c:255 UCC_COLL INFO coll_init: Barrier; CL_BASIC {TL_UCP}, team_id 32768 ``` +## Team Cache (Experimental) + +UCC supports an optional per-context communicator team cache that retains +`ucc_team_t` objects after `ucc_team_destroy` so that a subsequent +`ucc_team_create_post` with identical membership can re-adopt the same built +team instead of rebuilding it from scratch. + +### Configuration knobs + +| Environment variable | Default | Description | +|---|---|---| +| `UCC_TEAM_CACHE_ENABLE` | `n` | Enable the team cache. Off by default; opt-in. | +| `UCC_TEAM_CACHE_MAX_SIZE` | `128` | Maximum number of teams retained in the cache. Also clamped by `UCC_TEAM_IDS_POOL_SIZE`. | +| `UCC_TEAM_CACHE_EVICTION` | `fifo` | Eviction policy when the cache is full. `none`: never evict (new teams stay uncached). `fifo`: evict the oldest dormant entry (default). | +| `UCC_TEAM_CACHE_DISABLE_LINEAR_CHECK` | `n` | Trust the 64-bit membership hash alone in lookup, skipping the exact rank-array compare. Faster but unsafe on hash collision. | +| `UCC_TEAM_CACHE_DUMP_STATS` | `n` | Log hit/miss/eviction counters at context destroy. | +| `UCC_TEAM_CACHE_AGREEMENT` | `y` | Agree on the reuse decision across the members of every cacheable team create. Makes reuse safe for overlapping team scopes, at the cost of one small allreduce per create. | + +### Cross-rank agreement + +Each rank classifies a create as a cache hit or a miss from its own cache +contents, and those contents can diverge - for example when an eviction happens +on one rank only. Without agreement, the members of a single create could then +disagree on whether to re-adopt a dormant team or build a fresh one, and a create +where some ranks re-adopt while others rebuild does not make progress. + +`UCC_TEAM_CACHE_AGREEMENT` (on by default) reconciles that with a small +`UCC_OP_BAND` allreduce over the members before any rank skips the address +exchange. Reuse happens only when every member independently classified the +create the same way; otherwise all members fall back to a fresh build. This makes +reuse safe even when team scopes overlap. + +Disable the agreement only when team scopes never overlap - that is, when no rank +belongs to two simultaneously created teams with the same membership - and the +per-create allreduce is measurably too expensive. Applications that build only +disjoint or strictly nested communicators, such as a fixed set of row/column +communicators recreated over and over, satisfy that condition. Single-rank teams +never vote, since they cannot diverge. + +If the vote itself fails (as opposed to being lost), `ucc_team_create_test` +returns the error and the handle is terminal. A handle the create allocated may +still be passed to `ucc_team_destroy`, which releases it. A handle that named a +cached team has already been handed back to the cache; `ucc_team_destroy` rejects +it and the caller must simply drop it. Either way, do not call +`ucc_team_create_test` on it again. + +### Team-cache settings must be identical on every rank + +> **The team-cache settings above are not per-rank tunables. A rank whose +> settings differ from its peers' can hang the job, not merely lose reuse.** + +When caching and agreement are both on, a cacheable multi-rank create posts a +member-scoped allreduce (the *agreement vote*) so that every member reaches the +same reuse-vs-rebuild decision. A rank only enters that vote if all of the +following hold on that rank: + +- `UCC_TEAM_CACHE_ENABLE=y` +- `UCC_TEAM_CACHE_AGREEMENT=y` +- the team is cacheable (no optional behavioral fields in `ucc_team_params_t`) +- the team has more than one member, and +- the caller passed `UCC_TEAM_PARAM_FIELD_EP_MAP`. + +A rank that fails any of these skips the vote entirely and proceeds to build its +team directly. Its peers, meanwhile, have posted an allreduce that now has no +matching contribution from that rank and will never complete: the create hangs. + +In practice this means: + +- Set the team-cache variables in the launcher environment so every rank + inherits the same values (`mpirun -x UCC_TEAM_CACHE_ENABLE=y ...` or + `ucc.conf`). Do not set them from a per-rank wrapper script or from a rank + conditional. +- If a middleware creates some teams with `EP_MAP` and others without, that is + safe only when the choice is the same on every rank for a given team, which + it is for MPI communicators. +- If you must disable caching for part of a job, disable it for the whole job. + +Setting `UCC_TEAM_CACHE_AGREEMENT=n` uniformly on every rank removes the vote +and with it this hazard, but it is only safe when communicator scopes never +overlap (see the table above). + +### Requirements + +- `UCC_TEAM_PARAM_FIELD_EP_MAP` must be set in `ucc_team_params_t` for a team + to be cacheable (it provides the membership the cache keys on). +- Teams with optional behavioral parameters (`ORDERING`, `OUTSTANDING_COLLS`, + `SYNC_TYPE`, `P2P_CONN`, `MEM_PARAMS`) are not cached because those + parameters are not part of the identity. + +### Usage example + +```bash +UCC_TEAM_CACHE_ENABLE=y UCC_TEAM_CACHE_MAX_SIZE=64 mpirun -np 8 ./my_app +``` + ## Known Issues - For the CUDA and NCCL TL CUDA device dependent data structures are created when UCC diff --git a/src/Makefile.am b/src/Makefile.am index aaadb968db2..eb059b48d1d 100644 --- a/src/Makefile.am +++ b/src/Makefile.am @@ -51,6 +51,7 @@ noinst_HEADERS = \ core/ucc_lib.h \ core/ucc_context.h \ core/ucc_team.h \ + core/ucc_team_cache.h \ core/ucc_ee.h \ core/ucc_progress_queue.h \ core/ucc_service_coll.h \ @@ -115,6 +116,7 @@ libucc_la_SOURCES = \ core/ucc_version.c \ core/ucc_context.c \ core/ucc_team.c \ + core/ucc_team_cache.c \ core/ucc_ee.c \ core/ucc_coll.c \ core/ucc_progress_queue.c \ diff --git a/src/components/cl/hier/allgatherv/allgatherv.c b/src/components/cl/hier/allgatherv/allgatherv.c index f6a6432b0a4..3f906b337ea 100755 --- a/src/components/cl/hier/allgatherv/allgatherv.c +++ b/src/components/cl/hier/allgatherv/allgatherv.c @@ -54,7 +54,7 @@ static inline ucc_status_t find_leader_rank(ucc_base_team_t *team, ucc_assert(team_rank < UCC_CL_TEAM_SIZE(cl_team)); ucc_assert(SBGP_EXISTS(cl_team, NODE_LEADERS)); - status = ucc_topo_get_node_leaders(core_team->topo, &node_leaders); + status = ucc_topo_get_node_leaders(UCC_TEAM_TOPO(core_team), &node_leaders); if (UCC_OK != status) { cl_error(team->context->lib, "Could not get node leaders"); return status; @@ -69,7 +69,8 @@ static inline ucc_status_t find_leader_rank(ucc_base_team_t *team, dst buffer is contiguous */ static inline ucc_status_t is_block_ordered(ucc_cl_hier_team_t *cl_team, int *ordered) { - ucc_topo_t *topo = cl_team->super.super.params.team->topo; + ucc_topo_t *topo = + UCC_TEAM_TOPO(cl_team->super.super.params.team); ucc_sbgp_t *all_nodes = NULL; int is_block_ordered = 1; int n_nodes; @@ -129,7 +130,7 @@ UCC_CL_HIER_PROFILE_FUNC(ucc_status_t, ucc_cl_hier_allgatherv_init, ucc_rank_t node_sbgp_size = SBGP_SIZE(cl_team, NODE); ucc_rank_t leader_sbgp_size = SBGP_SIZE(cl_team, NODE_LEADERS); ucc_rank_t team_size = UCC_CL_TEAM_SIZE(cl_team); - ucc_topo_t *topo = team->params.team->topo; + ucc_topo_t *topo = UCC_TEAM_TOPO(team->params.team); ucc_aint_t *node_disps = NULL; ucc_count_t *node_counts = NULL; ucc_aint_t *leader_disps = NULL; diff --git a/src/components/cl/hier/allgatherv/unpack.c b/src/components/cl/hier/allgatherv/unpack.c index 5a42d31de8e..ba246b5da16 100644 --- a/src/components/cl/hier/allgatherv/unpack.c +++ b/src/components/cl/hier/allgatherv/unpack.c @@ -72,7 +72,8 @@ ucc_status_t ucc_cl_hier_allgatherv_unpack_start(ucc_coll_task_t *task) ucc_rank_t *node_leaders = NULL; ucc_sbgp_t *all_nodes = NULL; ucc_sbgp_t *node_leaders_sbgp = NULL; - ucc_topo_t *topo = task->team->params.team->topo; + ucc_topo_t *topo = UCC_TEAM_TOPO( + task->team->params.team); ucc_ee_executor_t *exec; ucc_status_t status; ucc_rank_t i; diff --git a/src/components/cl/hier/allreduce/allreduce_split_rail.c b/src/components/cl/hier/allreduce/allreduce_split_rail.c index da0d088abc7..f3cb8c4c143 100644 --- a/src/components/cl/hier/allreduce/allreduce_split_rail.c +++ b/src/components/cl/hier/allreduce/allreduce_split_rail.c @@ -296,7 +296,7 @@ UCC_CL_HIER_PROFILE_FUNC(ucc_status_t, ucc_cl_hier_allreduce_split_rail_init, return UCC_ERR_NOT_SUPPORTED; } - if (!ucc_topo_isoppn(team->params.team->topo)) { + if (!ucc_topo_isoppn(UCC_TEAM_TOPO(team->params.team))) { cl_debug(team->context->lib, "split_rail algorithm does not support " "teams with non-uniform ppn across nodes"); return UCC_ERR_NOT_SUPPORTED; diff --git a/src/components/cl/hier/alltoallv/alltoallv.c b/src/components/cl/hier/alltoallv/alltoallv.c index 52686bf04a0..96c3e34e268 100644 --- a/src/components/cl/hier/alltoallv/alltoallv.c +++ b/src/components/cl/hier/alltoallv/alltoallv.c @@ -52,8 +52,8 @@ static ucc_status_t ucc_cl_hier_alltoallv_finalize(ucc_coll_task_t *task) for (_i = 0; _i < (_sbgp)->group_size; _i++) { \ _scount = ((_type *)(_coll_args)->args.src.info_v.counts)[_i]; \ _rcount = ((_type *)(_coll_args)->args.dst.info_v.counts)[_i]; \ - _is_local = \ - ucc_rank_on_local_node(_i, (_team)->params.team->topo); \ + _is_local = ucc_rank_on_local_node( \ + _i, UCC_TEAM_TOPO((_team)->params.team)); \ if ((_scount * _sdt_size > (_node_thresh)) && _is_local) { \ ((_type *)_sc_full)[_i] = 0; \ } else { \ diff --git a/src/components/cl/hier/bcast/bcast_2step.c b/src/components/cl/hier/bcast/bcast_2step.c index 6059d09183d..266a4675192 100644 --- a/src/components/cl/hier/bcast/bcast_2step.c +++ b/src/components/cl/hier/bcast/bcast_2step.c @@ -146,7 +146,7 @@ ucc_cl_hier_bcast_2step_init_schedule(ucc_base_coll_args_t *coll_args, if (SBGP_ENABLED(cl_team, NODE)) { args.args.root = root_on_local_node ? find_root_node_rank(root, cl_team) - : core_team->topo->node_leader_rank_id; + : UCC_TEAM_TOPO(core_team)->node_leader_rank_id; status = ucc_coll_init(SCORE_MAP(cl_team, NODE), &args, &tasks[n_tasks]); if (ucc_unlikely(UCC_OK != status)) { diff --git a/src/components/cl/hier/cl_hier_team.c b/src/components/cl/hier/cl_hier_team.c index 5879ec7fd74..1334b307f66 100644 --- a/src/components/cl/hier/cl_hier_team.c +++ b/src/components/cl/hier/cl_hier_team.c @@ -47,13 +47,13 @@ UCC_CLASS_INIT_FUNC(ucc_cl_hier_team_t, ucc_base_context_t *cl_context, ucc_tl_lib_t *tl_lib; ucc_base_lib_attr_t attr; - if (!params->team->topo) { + if (!UCC_TEAM_TOPO(params->team)) { cl_debug(cl_context->lib, "can't create hier team without topology data"); return UCC_ERR_INVALID_PARAM; } - if (ucc_topo_is_single_node(params->team->topo)) { + if (ucc_topo_is_single_node(UCC_TEAM_TOPO(params->team))) { cl_debug(cl_context->lib, "skipping single node team"); return UCC_ERR_INVALID_PARAM; } @@ -65,7 +65,8 @@ UCC_CLASS_INIT_FUNC(ucc_cl_hier_team_t, ucc_base_context_t *cl_context, for (i = 0; i < UCC_HIER_SBGP_LAST; i++) { hs = &self->sbgps[i]; if (hs->state == UCC_HIER_SBGP_ENABLED) { - hs->sbgp = ucc_topo_get_sbgp(params->team->topo, hs->sbgp_type); + hs->sbgp = ucc_topo_get_sbgp(UCC_TEAM_TOPO(params->team), + hs->sbgp_type); if (hs->sbgp->status != UCC_SBGP_ENABLED) { /* SBGP of that type either not exists or the calling process * is not part of subgroup diff --git a/src/components/tl/cuda/tl_cuda_team.c b/src/components/tl/cuda/tl_cuda_team.c index f2752a90f98..f0749be2c76 100644 --- a/src/components/tl/cuda/tl_cuda_team.c +++ b/src/components/tl/cuda/tl_cuda_team.c @@ -338,7 +338,7 @@ ucc_status_t ucc_tl_cuda_team_create_test(ucc_base_team_t *tl_team) team->scratch.rem[i] = NULL; } - if (!ucc_topo_has_device_info(UCC_TL_CORE_TEAM(team)->topo)) { + if (!ucc_topo_has_device_info(UCC_TEAM_TOPO(UCC_TL_CORE_TEAM(team)))) { tl_debug(tl_team->context->lib, "not all ranks have visible GPU device info; " "skipping TL/CUDA team creation"); diff --git a/src/components/tl/cuda/tl_cuda_team_topo.c b/src/components/tl/cuda/tl_cuda_team_topo.c index 07633ea7883..8000edb8ff4 100644 --- a/src/components/tl/cuda/tl_cuda_team_topo.c +++ b/src/components/tl/cuda/tl_cuda_team_topo.c @@ -342,7 +342,7 @@ static ucc_status_t ucc_tl_cuda_team_topo_init_matrix(const ucc_tl_cuda_team_t *team, ucc_rank_t *matrix) { - ucc_topo_t *topo = UCC_TL_CORE_TEAM(team)->topo; + ucc_topo_t *topo = UCC_TEAM_TOPO(UCC_TL_CORE_TEAM(team)); ucc_proc_info_t *procs = topo->topo->procs; ucc_device_id_t *dev_ids = topo->device_map.device_ids; int size = UCC_TL_TEAM_SIZE(team); @@ -406,7 +406,7 @@ ucc_status_t ucc_tl_cuda_team_topo_create(const ucc_tl_team_t *cuda_team, * connectivity. This handles NVSwitch, fabric clique, and direct NVLink * connections consistently and avoids rescanning the matrix for zeros. */ { - ucc_topo_t *utopo = UCC_TL_CORE_TEAM(team)->topo; + ucc_topo_t *utopo = UCC_TEAM_TOPO(UCC_TL_CORE_TEAM(team)); ucc_sbgp_t *node_sg = ucc_topo_get_sbgp(utopo, UCC_SBGP_NODE); topo->is_fully_connected = ucc_topo_is_nvlink_fully_connected(utopo, node_sg); diff --git a/src/components/tl/mlx5/tl_mlx5_team.c b/src/components/tl/mlx5/tl_mlx5_team.c index 20b6ca51c6d..ab5bc275dee 100644 --- a/src/components/tl/mlx5/tl_mlx5_team.c +++ b/src/components/tl/mlx5/tl_mlx5_team.c @@ -19,7 +19,7 @@ static ucc_status_t ucc_tl_mlx5_topo_init(ucc_tl_mlx5_team_t *team) ucc_subset_t subset; ucc_status_t status; - status = ucc_ep_map_create_nested(&UCC_TL_CORE_TEAM(team)->ctx_map, + status = ucc_ep_map_create_nested(&UCC_TEAM_CTX_MAP(UCC_TL_CORE_TEAM(team)), &UCC_TL_TEAM_MAP(team), &team->ctx_map); if (UCC_OK != status) { tl_debug(UCC_TL_TEAM_LIB(team), "failed to create ctx map"); diff --git a/src/components/tl/sharp/tl_sharp_team.c b/src/components/tl/sharp/tl_sharp_team.c index f86098a4467..b120663ab75 100644 --- a/src/components/tl/sharp/tl_sharp_team.c +++ b/src/components/tl/sharp/tl_sharp_team.c @@ -37,7 +37,8 @@ UCC_CLASS_INIT_FUNC(ucc_tl_sharp_team_t, ucc_base_context_t *tl_context, set.map = UCC_TL_TEAM_MAP(self); if (UCC_TL_SHARP_TEAM_LIB(self)->cfg.use_internal_oob) { - status = ucc_ep_map_create_nested(&UCC_TL_CORE_TEAM(self)->ctx_map, + status = + ucc_ep_map_create_nested(&UCC_TEAM_CTX_MAP(UCC_TL_CORE_TEAM(self)), &UCC_TL_TEAM_MAP(self), &self->oob_ctx.subset.map); if (status != UCC_OK) { diff --git a/src/components/tl/ucc_tl.c b/src/components/tl/ucc_tl.c index 3134c9fd144..e11bb609ad9 100644 --- a/src/components/tl/ucc_tl.c +++ b/src/components/tl/ucc_tl.c @@ -188,7 +188,7 @@ static ucc_status_t ucc_tl_is_reachable(const ucc_base_team_params_t *params, for (i = 0; i < params->size; i++) { rank = ucc_ep_map_eval(params->map, i); if (use_ctx) { - rank = ucc_ep_map_eval(core_team->ctx_map, rank); + rank = ucc_ep_map_eval(UCC_TEAM_CTX_MAP(core_team), rank); } addr_header = UCC_ADDR_STORAGE_RANK_HEADER(addr_storage, rank); for (j = 0; j < addr_header->n_components; j++) { diff --git a/src/components/tl/ucp/tl_ucp_tag.h b/src/components/tl/ucp/tl_ucp_tag.h index 06e3b623013..e3d71781d29 100644 --- a/src/components/tl/ucp/tl_ucp_tag.h +++ b/src/components/tl/ucp/tl_ucp_tag.h @@ -48,6 +48,7 @@ #define UCC_TL_UCP_MAX_COLL_TAG (UCC_TL_UCP_MAX_TAG - UCC_TL_UCP_RESERVED_TAGS) #define UCC_TL_UCP_SERVICE_TAG (UCC_TL_UCP_MAX_COLL_TAG + 1) #define UCC_TL_UCP_ACTIVE_SET_TAG (UCC_TL_UCP_MAX_COLL_TAG + 2) +/* Tags MAX_COLL_TAG+3..+7 are unused; team-cache voting reuses SERVICE_TAG */ #define UCC_TL_UCP_MAX_SENDER UCC_MASK(UCC_TL_UCP_SENDER_BITS) #define UCC_TL_UCP_MAX_ID UCC_MASK(UCC_TL_UCP_ID_BITS) diff --git a/src/components/tl/ucp/tl_ucp_team.c b/src/components/tl/ucp/tl_ucp_team.c index aa1a1dcd767..5656e9d6116 100644 --- a/src/components/tl/ucp/tl_ucp_team.c +++ b/src/components/tl/ucp/tl_ucp_team.c @@ -23,7 +23,8 @@ static inline ucc_status_t ucc_tl_ucp_get_topo(ucc_tl_ucp_team_t *team) return UCC_OK; } - status = ucc_ep_map_create_nested(&UCC_TL_CORE_TEAM(team)->ctx_map, + /* The nested map aliases ctx_map, kept at a fixed address by the holder */ + status = ucc_ep_map_create_nested(&UCC_TEAM_CTX_MAP(UCC_TL_CORE_TEAM(team)), &UCC_TL_TEAM_MAP(team), &team->ctx_map); if (UCC_OK != status) { diff --git a/src/components/topo/ucc_topo.c b/src/components/topo/ucc_topo.c index bb22f60cda9..a49d8b25cb6 100644 --- a/src/components/topo/ucc_topo.c +++ b/src/components/topo/ucc_topo.c @@ -311,6 +311,49 @@ void ucc_topo_cleanup(ucc_topo_t *topo) } } +#define UCC_TOPO_PREP_FATAL(_status) ((_status) == UCC_ERR_NO_MEMORY) + +ucc_status_t ucc_topo_prepare_shared(ucc_topo_t *topo) +{ + ucc_sbgp_t *sbgps; + ucc_rank_t *node_leaders; + ucc_status_t status; + int n_sbgps, i; + + if (!topo) { + return UCC_OK; + } + + /* Each sbgp records a terminal status, so a failure is never retried */ + for (i = 0; i < UCC_SBGP_LAST; i++) { + (void)ucc_topo_get_sbgp(topo, (ucc_sbgp_type_t)i); + } + + /* The all_* arrays stay NULL and retryable on failure, so report it */ + status = ucc_topo_get_all_sockets(topo, &sbgps, &n_sbgps); + if (UCC_TOPO_PREP_FATAL(status)) { + return status; + } + status = ucc_topo_get_all_numas(topo, &sbgps, &n_sbgps); + if (UCC_TOPO_PREP_FATAL(status)) { + return status; + } + status = ucc_topo_get_all_nodes(topo, &sbgps, &n_sbgps); + if (UCC_TOPO_PREP_FATAL(status)) { + return status; + } + + /* The node leaders map is only defined for multi-node teams */ + if (topo->topo->nnodes > 1) { + status = ucc_topo_get_node_leaders(topo, &node_leaders); + if (UCC_TOPO_PREP_FATAL(status)) { + return status; + } + } + + return UCC_OK; +} + ucc_sbgp_t *ucc_topo_get_sbgp(ucc_topo_t *topo, ucc_sbgp_type_t type) { if (topo->sbgps[type].status == UCC_SBGP_NOT_INIT) { diff --git a/src/components/topo/ucc_topo.h b/src/components/topo/ucc_topo.h index c1fd8a2b4f7..389818ff7ed 100644 --- a/src/components/topo/ucc_topo.h +++ b/src/components/topo/ucc_topo.h @@ -108,6 +108,9 @@ ucc_status_t ucc_topo_init( void ucc_topo_cleanup(ucc_topo_t *subset_topo); +/* Materializes the lazily built topo fields so it can be shared read-only */ +ucc_status_t ucc_topo_prepare_shared(ucc_topo_t *topo); + ucc_sbgp_t *ucc_topo_get_sbgp(ucc_topo_t *topo, ucc_sbgp_type_t type); int ucc_topo_is_single_node(ucc_topo_t *topo); diff --git a/src/core/ucc_context.c b/src/core/ucc_context.c index 8f2f43316ee..d5e178d0064 100644 --- a/src/core/ucc_context.c +++ b/src/core/ucc_context.c @@ -8,6 +8,8 @@ #include #include "ucc_context.h" #include "components/topo/ucc_topo.h" +#include "ucc_team.h" +#include "ucc_team_cache.h" #include "utils/ucc_proc_info.h" #include "components/cl/ucc_cl.h" #include "components/tl/ucc_tl.h" @@ -20,6 +22,9 @@ #include "utils/ucc_debug.h" #include "ucc_progress_queue.h" +/* Team ids kept free for in-flight creates when clamping the cache size */ +#define UCC_TEAM_CACHE_ID_HEADROOM 8U + static uint32_t ucc_context_seq_num = 0; static ucc_config_field_t ucc_context_config_table[] = { {"ESTIMATED_NUM_EPS", "0", @@ -46,6 +51,47 @@ static ucc_config_field_t ucc_context_config_table[] = { ucc_offsetof(ucc_context_config_t, team_ids_pool_size), UCC_CONFIG_TYPE_UINT}, + {"TEAM_CACHE_ENABLE", "n", + "Retain destroyed teams so that a later create with identical membership\n" + "can reuse the already built team. Experimental.", + ucc_offsetof(ucc_context_config_t, team_cache_enable), + UCC_CONFIG_TYPE_BOOL}, + + {"TEAM_CACHE_MAX_SIZE", "128", + "Maximum number of teams retained in the per-context team cache. Each\n" + "retained team holds a team id, so this is also clamped by\n" + "UCC_TEAM_IDS_POOL_SIZE.", + ucc_offsetof(ucc_context_config_t, team_cache_max_size), + UCC_CONFIG_TYPE_UINT}, + + {"TEAM_CACHE_EVICTION", "fifo", + "Eviction policy once the team cache is full. Only retained teams are\n" + "evictable. none - never evict, the new team stays uncached.\n" + "fifo - evict the oldest retained team.", + ucc_offsetof(ucc_context_config_t, team_cache_eviction), + UCC_CONFIG_TYPE_ENUM(ucc_team_cache_eviction_names)}, + + {"TEAM_CACHE_DISABLE_LINEAR_CHECK", "n", + "Trust the 64 bit identity hash alone on lookup and skip the exact\n" + "membership compare. Faster, but a hash collision would then reuse the\n" + "wrong team.", + ucc_offsetof(ucc_context_config_t, team_cache_disable_linear_check), + UCC_CONFIG_TYPE_BOOL}, + + {"TEAM_CACHE_DUMP_STATS", "n", + "Log team cache statistics (lookups, hits, misses, inserts, evictions)\n" + "when the context is destroyed.", + ucc_offsetof(ucc_context_config_t, team_cache_dump_stats), + UCC_CONFIG_TYPE_BOOL}, + + {"TEAM_CACHE_AGREEMENT", "y", + "Agree on the reuse decision across the members of every cacheable team\n" + "create, which makes reuse safe for overlapping subcommunicators at the\n" + "cost of one small allreduce per create. Disable only when team scopes\n" + "never overlap. Requires UCC_TEAM_CACHE_ENABLE.", + ucc_offsetof(ucc_context_config_t, team_cache_agreement), + UCC_CONFIG_TYPE_BOOL}, + {"INTERNAL_OOB", "1", "Use internal OOB transport for team creation. Available for ucc_context " "is configured with OOB (global mode). 0 - disable, 1 - try, 2 - force.", @@ -906,6 +952,52 @@ ucc_status_t ucc_context_create_proc_info( ctx->rank = UCC_RANK_MAX; ctx->lib = lib; ctx->ids.pool_size = config->team_ids_pool_size; + /* ctx is calloc'd, so team_cache is already NULL when caching is off */ + if (config->team_cache_enable) { + /* Retained teams hold team ids, so bound the cache by the id pool */ + uint64_t pool_capacity = (uint64_t)config->team_ids_pool_size * 64U; + /* Id 0 is reserved, hence the extra -1 */ + uint64_t reserved = (uint64_t)UCC_TEAM_CACHE_ID_HEADROOM + 1U; + uint32_t safe_max = (pool_capacity > reserved) + ? (uint32_t)(pool_capacity - reserved) + : 0U; + uint32_t cache_max = config->team_cache_max_size; + + if (cache_max > safe_max) { + ucc_info( + "UCC_TEAM_CACHE_MAX_SIZE=%u exceeds the safe bound " + "derived from UCC_TEAM_IDS_POOL_SIZE=%u " + "(pool_capacity=%llu, headroom=%u, safe_max=%u); " + "clamping cache to %u entries. " + "Raise UCC_TEAM_IDS_POOL_SIZE to raise the safe ceiling.", + cache_max, + config->team_ids_pool_size, + (unsigned long long)pool_capacity, + UCC_TEAM_CACHE_ID_HEADROOM, + safe_max, + safe_max); + cache_max = safe_max; + } + + status = ucc_team_cache_init( + &ctx->team_cache, + cache_max, + (ucc_team_cache_eviction_policy_t)config->team_cache_eviction, + config->team_cache_disable_linear_check); + if (UCC_OK != status) { + /* Fatal, not degraded: a cache-less rank would skip the vote its + peers wait on, so only configuration may decide caching */ + ucc_error("failed to init team cache: %s", + ucc_status_string(status)); + goto error_ctx; + } + ctx->team_cache->dump_stats = config->team_cache_dump_stats; + ctx->team_cache->agreement = config->team_cache_agreement; + /* cache_gen is seeded once ctx->id.seq_num is assigned below */ + ucc_debug("team cache enabled (max_size=%u)", cache_max); + } else { + ucc_debug("team cache disabled by configuration"); + } ucc_list_head_init(&ctx->progress_list); ucc_copy_context_params(&ctx->params, params); ucc_copy_context_params(&b_params.params, params); @@ -1083,7 +1175,9 @@ ucc_status_t ucc_context_create_proc_info( ucc_error("failed to init progress queue for context %p", ctx); goto error_ctx_create; } - + if (ctx->team_cache != NULL) { + ctx->team_cache->cache_gen = ((uint64_t)ctx->id.seq_num << 32); + } if (params->mask & UCC_CONTEXT_PARAM_FIELD_OOB && params->oob.n_oob_eps > 1) { do { @@ -1138,6 +1232,32 @@ ucc_status_t ucc_context_create_proc_info( } } + /* The agreement vote runs over the context service team, which is only + created above and whose absence is not fatal. Without it the vote cannot + run, so drop the cache entirely: reusing teams without agreement is safe + only when team scopes never overlap, and silently assuming that would + trade a performance feature for a correctness risk. The decision is + taken from configuration alone, so it is identical on every rank. */ + if (ctx->team_cache != NULL && ctx->team_cache->agreement && + ctx->service_team == NULL) { + ucc_warn("team cache disabled: the cross-rank agreement vote requires " + "a context service team, which is not available. Set " + "UCC_TEAM_CACHE_AGREEMENT=n to cache without it, which is " + "only safe when team scopes never overlap."); + ucc_team_cache_destroy(ctx->team_cache); + ctx->team_cache = NULL; + } + if (ctx->team_cache != NULL) { + /* Registered only once the cache is final for this context */ + status = ucc_context_progress_register( + ctx, ucc_team_cache_progress_cb, ctx->team_cache); + if (UCC_OK != status) { + ucc_error("failed to register team cache progress: %s", + ucc_status_string(status)); + goto error_ctx_create; + } + } + n_tl_ctx = ctx->n_tl_ctx; for (i = 0; i < n_tl_ctx; i++) { tl_ctx = ctx->tl_ctx[i]; @@ -1201,6 +1321,7 @@ ucc_status_t ucc_context_create_proc_info( ctx->topo = NULL; ucc_free(ctx->addr_storage.storage); ctx->addr_storage.storage = NULL; + ucc_team_cache_destroy(ctx->team_cache); /* NULL-safe */ ucc_free(ctx); error: return status; @@ -1235,6 +1356,18 @@ ucc_status_t ucc_context_destroy(ucc_context_t *context) if (UCC_OK != ucc_context_free_attr(&context->attr)) { ucc_error("failed to free context attributes"); } + + /* Retained teams hold CL/TL refs, so drain before those are destroyed */ + if (context->team_cache) { + if (context->team_cache->dump_stats) { + ucc_team_cache_dump_stats(context->team_cache); + } + ucc_context_progress_deregister( + context, ucc_team_cache_progress_cb, context->team_cache); + ucc_team_cache_drain(context); + ucc_team_cache_destroy(context->team_cache); + context->team_cache = NULL; + } for (i = 0; i < context->n_cl_ctx; i++) { cl_ctx = context->cl_ctx[i]; cl_lib = ucc_derived_of(cl_ctx->super.lib, ucc_cl_lib_t); diff --git a/src/core/ucc_context.h b/src/core/ucc_context.h index 5e2699719de..1fc98619060 100644 --- a/src/core/ucc_context.h +++ b/src/core/ucc_context.h @@ -13,6 +13,9 @@ #include "utils/ucc_proc_info.h" #include "components/topo/ucc_topo.h" +/* Forward declaration, so context callers need not see the cache */ +typedef struct ucc_team_cache ucc_team_cache_t; + #define UCC_MEM_MAP_TL_NAME_LEN 8 typedef struct ucc_lib_info ucc_lib_info_t; @@ -84,6 +87,7 @@ typedef struct ucc_context { uint64_t cl_flags; ucc_tl_team_t *service_team; int32_t throttle_progress; + ucc_team_cache_t *team_cache; } ucc_context_t; typedef struct ucc_context_config { @@ -93,6 +97,12 @@ typedef struct ucc_context_config { int n_cl_cfg; int n_tl_cfg; uint32_t team_ids_pool_size; + uint32_t team_cache_enable; + uint32_t team_cache_max_size; + uint32_t team_cache_eviction; + uint32_t team_cache_disable_linear_check; + uint32_t team_cache_dump_stats; + uint32_t team_cache_agreement; uint32_t estimated_num_eps; uint32_t estimated_num_ppn; uint32_t lock_free_progress_q; diff --git a/src/core/ucc_service_coll.c b/src/core/ucc_service_coll.c index 7ccd6e8d9ac..da24fc5cc64 100644 --- a/src/core/ucc_service_coll.c +++ b/src/core/ucc_service_coll.c @@ -18,7 +18,15 @@ uint64_t ucc_service_coll_map_cb(uint64_t ep, void *cb_ctx) ucc_rank_t team_rank; team_rank = ucc_ep_map_eval(req->subset.map, (ucc_rank_t)ep); - return ucc_ep_map_eval(team->ctx_map, team_rank); + return ucc_ep_map_eval(UCC_TEAM_CTX_MAP(team), team_rank); +} + +/* Maps a subset index straight to a context rank, bypassing team->ctx_map */ +static uint64_t ucc_service_coll_map_cb_direct(uint64_t ep, void *cb_ctx) +{ + ucc_service_coll_req_t *req = cb_ctx; + + return ucc_ep_map_eval(req->subset.map, (ucc_rank_t)ep); } static inline ucc_status_t @@ -36,8 +44,9 @@ ucc_service_coll_req_init(ucc_team_t *team, ucc_subset_t *subset, sizeof(*req)); return UCC_ERR_NO_MEMORY; } - req->team = team; - req->subset = *subset; + req->team = team; + req->subset = *subset; + req->embedded = 0; if (ctx->service_team) { *service_team = ctx->service_team; @@ -80,6 +89,51 @@ ucc_status_t ucc_service_allreduce(ucc_team_t *team, void *sbuf, void *rbuf, return UCC_OK; } +ucc_status_t ucc_service_allreduce_ctx(ucc_team_t *team, void *sbuf, + void *rbuf, ucc_datatype_t dt, + size_t count, ucc_reduction_op_t op, + ucc_subset_t subset, + ucc_service_coll_req_t *req) +{ + ucc_context_t *ctx = team->contexts[0]; + ucc_tl_team_t *steam; + ucc_tl_iface_t *tl_iface; + ucc_status_t status; + + /* A context service team exists only when the context was created with an + OOB and UCC_INTERNAL_OOB allowed it, and its creation is non-fatal, so + NULL is a supported state rather than a caller error. There is no + fallback: this collective addresses context ranks directly, which the + per-team service team cannot do. */ + if (ucc_unlikely(ctx->service_team == NULL)) { + ucc_debug("context %p has no service team, ctx-scoped service " + "allreduce is not available", + ctx); + return UCC_ERR_NOT_SUPPORTED; + } + + req->team = team; + req->subset = subset; + req->data = NULL; + req->embedded = 1; + + subset.map.type = UCC_EP_MAP_CB; + subset.map.cb.cb = ucc_service_coll_map_cb_direct; + subset.map.cb.cb_ctx = req; + + steam = ctx->service_team; + tl_iface = UCC_TL_TEAM_IFACE(steam); + status = tl_iface->scoll.allreduce(&steam->super, sbuf, rbuf, dt, count, + op, subset, &req->task); + if (status < 0) { + ucc_error("failed to start ctx service allreduce for team %p: %s", team, + ucc_status_string(status)); + return status; + } + + return UCC_OK; +} + ucc_status_t ucc_service_allgather(ucc_team_t *team, void *sbuf, void *rbuf, size_t msgsize, ucc_subset_t subset, ucc_service_coll_req_t **req) @@ -145,10 +199,13 @@ ucc_status_t ucc_service_coll_test(ucc_service_coll_req_t *req) ucc_status_t ucc_service_coll_finalize(ucc_service_coll_req_t *req) { + uint8_t embedded = req->embedded; ucc_status_t status; status = ucc_collective_finalize_internal(req->task); - ucc_free(req); + if (!embedded) { + ucc_free(req); + } return status; } diff --git a/src/core/ucc_service_coll.h b/src/core/ucc_service_coll.h index e1794de1c71..b7a0ac62f22 100644 --- a/src/core/ucc_service_coll.h +++ b/src/core/ucc_service_coll.h @@ -14,6 +14,7 @@ typedef struct ucc_service_coll_req { ucc_team_t *team; void * data; ucc_subset_t subset; + uint8_t embedded; /* req storage is caller-owned, never freed */ } ucc_service_coll_req_t; ucc_status_t ucc_service_allreduce(ucc_team_t *team, void *sbuf, void *rbuf, @@ -21,6 +22,15 @@ ucc_status_t ucc_service_allreduce(ucc_team_t *team, void *sbuf, void *rbuf, ucc_reduction_op_t op, ucc_subset_t subset, ucc_service_coll_req_t **req); +/* Service allreduce over ctx->service_team; @subset maps to ctx ranks. Unlike + the siblings above, @req is caller-provided (embedded in the team) rather + than allocated here. */ +ucc_status_t ucc_service_allreduce_ctx(ucc_team_t *team, void *sbuf, + void *rbuf, ucc_datatype_t dt, + size_t count, ucc_reduction_op_t op, + ucc_subset_t subset, + ucc_service_coll_req_t *req); + ucc_status_t ucc_service_allgather(ucc_team_t *team, void *sbuf, void *rbuf, size_t msgsize, ucc_subset_t subset, ucc_service_coll_req_t **req); diff --git a/src/core/ucc_team.c b/src/core/ucc_team.c index e3e4337ba75..b40cfc5dc92 100644 --- a/src/core/ucc_team.c +++ b/src/core/ucc_team.c @@ -11,9 +11,49 @@ #include "components/cl/ucc_cl.h" #include "components/tl/ucc_tl.h" #include "ucc_service_coll.h" +#include static ucc_status_t ucc_team_alloc_id(ucc_team_t *team); static void ucc_team_release_id(ucc_team_t *team); +static ucc_status_t ucc_team_destroy_single(ucc_team_h team); +static ucc_status_t ucc_team_destroy_single_ex( + ucc_team_h team, int for_rebuild); +static ucc_status_t ucc_team_teardown_for_rebuild(ucc_team_t *team); +static ucc_status_t ucc_team_reset_for_rebuild( + ucc_context_t *context, ucc_team_t *team); + +void ucc_team_artifacts_init_inline(ucc_team_artifacts_t *a) +{ + memset(a, 0, sizeof(*a)); + a->refcount = 1; + a->heap = 0; + ucc_spinlock_init(&a->lock, 0); +} + +void ucc_team_artifacts_put(ucc_team_artifacts_t *artifacts) +{ + int refcount; + + if (!artifacts) { + return; + } + ucc_spin_lock(&artifacts->lock); + ucc_assert(artifacts->refcount > 0); + refcount = --artifacts->refcount; + ucc_spin_unlock(&artifacts->lock); + + if (refcount > 0) { + return; + } + + /* ctx_ranks is NULL when ctx_map aliases the caller's ep_map */ + ucc_topo_cleanup(artifacts->topo); + ucc_free(artifacts->ctx_ranks); + ucc_spinlock_destroy(&artifacts->lock); + if (artifacts->heap) { + ucc_free(artifacts); + } +} void ucc_copy_team_params(ucc_team_params_t *dst, const ucc_team_params_t *src) { @@ -68,8 +108,9 @@ static ucc_status_t ucc_team_create_post_single(ucc_context_t *context, .map.type = UCC_EP_MAP_FULL}; status = ucc_internal_oob_init(team, subset, &team->bp.params.oob); if (UCC_OK != status) { - return status; + return status; /* bp.params.oob is still the caller's */ } + team->internal_oob = 1; team->bp.params.mask |= UCC_TEAM_PARAM_FIELD_OOB; } @@ -102,6 +143,193 @@ static ucc_status_t ucc_team_create_post_single(ucc_context_t *context, return UCC_OK; } +/* Allocate an unbuilt team; on @id_built the identity MOVES from @id to it. + The move happens only on success: when this returns NULL the caller still + owns @id and must free it. Both call sites rely on that split. */ +static ucc_team_t *ucc_team_alloc_shell( + ucc_context_h *contexts, uint32_t num_contexts, + const ucc_team_params_t *params, uint64_t team_size, uint64_t team_rank, + int id_built, ucc_team_cache_identity_t *id, ucc_status_t *status_out) +{ + ucc_team_t *team; + + team = ucc_calloc(1, sizeof(ucc_team_t), "ucc_team"); + if (!team) { + ucc_error( + "failed to allocate %zd bytes for ucc team", sizeof(ucc_team_t)); + *status_out = UCC_ERR_NO_MEMORY; + return NULL; + } + team->artifacts = &team->artifacts_inline; + ucc_team_artifacts_init_inline(team->artifacts); + team->runtime_oob = params->oob; + team->num_contexts = num_contexts; + team->size = (ucc_rank_t)team_size; + team->rank = (ucc_rank_t)team_rank; + team->seq_num = 0; + team->refcount = 1; + if (id_built) { + team->cache_identity = *id; /* the caller's id now owns nothing */ + team->cache_pending_insert = 1; + memset(id, 0, sizeof(*id)); + } else { + memset(&team->cache_identity, 0, sizeof(team->cache_identity)); + team->cache_pending_insert = 0; + } + ucc_list_head_init(&team->cache_link); + team->cache_state = UCC_TEAM_CACHE_STATE_NONE; + team->cache_local_action = UCC_TEAM_CACHE_ACTION_MISS; + team->contexts = ucc_malloc( + sizeof(ucc_context_t *) * num_contexts, "ucc_team_ctx"); + if (!team->contexts) { + ucc_error( + "failed to allocate %zd bytes for ucc team contexts array", + sizeof(ucc_context_t *) * num_contexts); + ucc_team_cache_identity_free(&team->cache_identity); + ucc_team_artifacts_put(team->artifacts); + ucc_free(team); + *status_out = UCC_ERR_NO_MEMORY; + return NULL; + } + memcpy(team->contexts, contexts, sizeof(ucc_context_t *) * num_contexts); + ucc_copy_team_params(&team->bp.params, params); + /* check if user provides team id and if it is not too large */ + if ((params->mask & UCC_TEAM_PARAM_FIELD_ID) && + (params->id <= UCC_TEAM_ID_MAX)) { + team->id = ((uint16_t)params->id) | UCC_TEAM_ID_EXTERNAL_BIT; + } + *status_out = UCC_OK; + return team; +} + +/* Classify this rank's cache action and post the vote that reconciles it */ +/* Return a RESERVED candidate to dormant; no refcount change happened yet */ +static void ucc_team_agreement_release_reserved( + ucc_team_cache_t *cache, ucc_team_t *handle, ucc_team_cache_action_t action) +{ + ucc_assert(action == UCC_TEAM_CACHE_ACTION_EXACT_REUSE); + ucc_spin_lock(&cache->lock); + handle->cache_state = UCC_TEAM_CACHE_STATE_DORMANT; + ucc_team_cache_registry_make_dormant(cache, handle); + ucc_spin_unlock(&cache->lock); +} + +/* Undo a posted vote setup after a fatal post failure */ +static void ucc_team_agreement_rollback( + ucc_team_cache_t *cache, ucc_team_t *handle, ucc_team_cache_action_t action) +{ + if (action == UCC_TEAM_CACHE_ACTION_EXACT_REUSE) { + ucc_team_agreement_release_reserved(cache, handle, action); + return; + } + ucc_team_destroy_single(handle); +} + +/* Vote error (not a lost vote): hand back or park, and never re-test the req */ +static void ucc_team_agreement_fail(ucc_context_t *context, ucc_team_t *team) +{ + ucc_team_cache_action_t action = team->cache_local_action; + + if (action == UCC_TEAM_CACHE_ACTION_EXACT_REUSE) { + ucc_team_agreement_release_reserved(context->team_cache, team, action); + } + team->state = UCC_TEAM_CREATE_FAILED; +} + +static ucc_status_t ucc_team_agreement_create_post( + ucc_context_h *contexts, uint32_t num_contexts, + const ucc_team_params_t *params, uint64_t team_size, uint64_t team_rank, + ucc_team_cache_t *cache, ucc_team_h *new_team) +{ + ucc_team_cache_identity_t id; + int id_built = 0; + ucc_team_t *cached = NULL; + ucc_team_t *handle; + ucc_team_cache_action_t action = UCC_TEAM_CACHE_ACTION_MISS; + uint64_t key = 0; + int is_rank0 = (team_rank == 0); + uint64_t proposed_cookie = 0; + ucc_subset_t subset; + ucc_status_t status; + + if (ucc_team_cache_identity_build(params, &id) == UCC_OK) { + id_built = 1; + ucc_spin_lock(&cache->lock); + /* Drawn unconditionally, so a MISS outcome has a cookie to adopt */ + if (is_rank0) { + proposed_cookie = ucc_team_cache_next_cookie(cache); + } + cached = ucc_team_cache_lookup(cache, &id); + if (cached != NULL) { + action = UCC_TEAM_CACHE_ACTION_EXACT_REUSE; + key = cached->id; + ucc_team_cache_registry_make_reserved(cache, cached); + cached->cache_state = UCC_TEAM_CACHE_STATE_RESERVED; + } + ucc_spin_unlock(&cache->lock); + } + + if (action == UCC_TEAM_CACHE_ACTION_EXACT_REUSE) { + handle = cached; + /* The candidate already carries its own identity */ + ucc_team_cache_identity_free(&id); + /* Refresh the map, since the candidate's may name a freed team */ + handle->bp.params.ep_map = params->ep_map; + handle->bp.params.mask |= UCC_TEAM_PARAM_FIELD_EP_MAP; + } else { + handle = ucc_team_alloc_shell( + contexts, + num_contexts, + params, + team_size, + team_rank, + id_built, + &id, + &status); + if (handle == NULL) { + if (id_built) { + ucc_team_cache_identity_free(&id); + } + return status; + } + status = ucc_team_create_post_single(contexts[0], handle); + if (status < 0) { + ucc_team_destroy_single(handle); + return status; + } + } + + handle->cache_local_action = action; + ucc_team_cache_vote_fill( + handle->cache_vote_in, + action != UCC_TEAM_CACHE_ACTION_MISS, + action, + key, + /*cookie=*/0, + /*parent_cookie=*/0, + is_rank0, + proposed_cookie); + /* ep_map is valid for this create and maps member index to ctx rank */ + subset.myrank = handle->rank; + subset.map = params->ep_map; + status = ucc_service_allreduce_ctx( + handle, + handle->cache_vote_in, + handle->cache_vote_out, + UCC_DT_UINT64, + UCC_TEAM_CACHE_VOTE_LANES, + UCC_OP_BAND, + subset, + &handle->cache_vote_req); + if (status < 0) { + ucc_team_agreement_rollback(cache, handle, action); + return status; + } + handle->state = UCC_TEAM_CACHE_AGREE; + *new_team = handle; + return UCC_OK; +} + ucc_status_t ucc_team_create_post(ucc_context_h *contexts, uint32_t num_contexts, const ucc_team_params_t *params, ucc_team_h *new_team) @@ -110,6 +338,9 @@ ucc_status_t ucc_team_create_post(ucc_context_h *contexts, uint32_t num_contexts uint64_t team_rank = UINT64_MAX; ucc_team_t *team; ucc_status_t status; + ucc_team_cache_t *cache = NULL; + ucc_team_cache_identity_t id; + int id_built = 0; if (num_contexts < 1) { return UCC_ERR_INVALID_PARAM; @@ -188,41 +419,79 @@ ucc_status_t ucc_team_create_post(ucc_context_h *contexts, uint32_t num_contexts return UCC_ERR_INVALID_PARAM; } - team = ucc_calloc(1, sizeof(ucc_team_t), "ucc_team"); - if (!team) { - ucc_error("failed to allocate %zd bytes for ucc team", - sizeof(ucc_team_t)); - return UCC_ERR_NO_MEMORY; + /* Agreement keeps every member on the same reuse decision. + + Every condition below must evaluate identically on every member. A rank + that skips the vote (caching or agreement off, a differing params->mask, + no EP_MAP) builds its team directly while its peers wait on a + member-scoped allreduce that never completes, so a non-uniform + configuration hangs the create rather than degrading it. This is + documented for users under "Team-cache settings must be identical on + every rank" in docs/user_guide.md. */ + cache = ((ucc_context_t *)contexts[0])->team_cache; + if (cache != NULL && cache->agreement && + ucc_team_cache_is_cacheable(params) && team_size > 1 && + (params->mask & UCC_TEAM_PARAM_FIELD_EP_MAP)) { + return ucc_team_agreement_create_post( + contexts, + num_contexts, + params, + team_size, + team_rank, + cache, + new_team); } - team->runtime_oob = params->oob; - team->num_contexts = num_contexts; - team->size = (ucc_rank_t)team_size; - team->rank = (ucc_rank_t)team_rank; - team->seq_num = 0; - team->contexts = ucc_malloc(sizeof(ucc_context_t *) * num_contexts, - "ucc_team_ctx"); - if (!team->contexts) { - ucc_error("failed to allocate %zd bytes for ucc team contexts array", - sizeof(ucc_context_t) * num_contexts); - status = UCC_ERR_NO_MEMORY; - goto err_ctx_alloc; + + /* Direct reuse, for a single-rank team or when agreement is disabled */ + if (cache != NULL && ucc_team_cache_is_cacheable(params)) { + status = ucc_team_cache_identity_build(params, &id); + if (status == UCC_OK) { + ucc_team_t *cached; + + id_built = 1; + + /* One lock spans lookup and adopt, so no team is adopted twice */ + ucc_spin_lock(&cache->lock); + cached = ucc_team_cache_lookup(cache, &id); + if (cached != NULL) { + cached->cache_local_action = UCC_TEAM_CACHE_ACTION_EXACT_REUSE; + ucc_team_cache_get(cached); + ucc_team_cache_registry_make_live(cache, cached); + } + ucc_spin_unlock(&cache->lock); + + if (cached != NULL) { + ucc_debug( + "team cache: dormant reuse / hit, team %p (hash=0x%" PRIx64 + ")", + (void *)cached, + id.hash); + ucc_team_cache_identity_free(&id); + *new_team = cached; + return UCC_OK; + } + /* On a miss @id moves onto the new team, to insert once ACTIVE */ + } } - memcpy(team->contexts, contexts, sizeof(ucc_context_t *) * num_contexts); - ucc_copy_team_params(&team->bp.params, params); - /* check if user provides team id and if it is not too large */ - if ((params->mask & UCC_TEAM_PARAM_FIELD_ID) && - (params->id <= UCC_TEAM_ID_MAX)) { - team->id = ((uint16_t)params->id) | UCC_TEAM_ID_EXTERNAL_BIT; + team = ucc_team_alloc_shell( + contexts, + num_contexts, + params, + team_size, + team_rank, + id_built, + &id, + &status); + if (team == NULL) { + if (id_built) { + ucc_team_cache_identity_free(&id); + } + return status; } status = ucc_team_create_post_single(contexts[0], team); *new_team = team; return status; - -err_ctx_alloc: - *new_team = NULL; - ucc_free(team); - return status; } static ucc_status_t ucc_team_create_service_team(ucc_context_t *context, @@ -277,14 +546,24 @@ static ucc_status_t ucc_team_create_cls(ucc_context_t *context, ucc_subset_t subset; int i; - if (context->topo && !team->topo && team->size > 1) { + if (context->topo && !UCC_TEAM_TOPO(team) && team->size > 1) { /* Context->topo is not NULL if any of the enabled CLs reported topo_required through the lib_attr */ - subset.map = team->ctx_map; + subset.map = UCC_TEAM_CTX_MAP(team); subset.myrank = team->rank; - status = ucc_topo_init(subset, context->topo, &team->topo); + status = ucc_topo_init(subset, context->topo, &UCC_TEAM_TOPO(team)); if (UCC_OK != status) { ucc_warn("failed to init team topo"); + } else if (team->cache_pending_insert) { + /* A cacheable topo may be shared, so materialize it up front */ + status = ucc_topo_prepare_shared(UCC_TEAM_TOPO(team)); + if (UCC_OK != status) { + /* Keep the team, but leave it un-shared and lazily filled */ + ucc_warn("failed to prepare shared topo (%s); team %p will " + "not be cached", + ucc_status_string(status), (void *)team); + team->cache_pending_insert = 0; + } } } @@ -346,22 +625,47 @@ static inline ucc_status_t ucc_team_exchange(ucc_context_t *context, /* We only need to exchange ctx_ranks and build map to ctx array */ ucc_assert(context->addr_storage.storage != NULL); if (team->bp.params.mask & UCC_TEAM_PARAM_FIELD_EP_MAP) { - team->ctx_map = team->bp.params.ep_map; + if (team->cache_pending_insert) { + /* A cached team outlives the caller's map, so copy it now */ + ucc_rank_t i; + + if (!UCC_TEAM_CTX_RANKS(team)) { + UCC_TEAM_CTX_RANKS(team) = ucc_malloc( + team->size * sizeof(ucc_rank_t), "ctx_ranks"); + if (!UCC_TEAM_CTX_RANKS(team)) { + ucc_error( + "failed to allocate %zd bytes for ctx ranks array", + team->size * sizeof(ucc_rank_t)); + return UCC_ERR_NO_MEMORY; + } + for (i = 0; i < team->size; i++) { + UCC_TEAM_CTX_RANKS(team)[i] = + (ucc_rank_t)ucc_ep_map_eval(team->bp.params.ep_map, i); + } + } + UCC_TEAM_CTX_MAP(team) = ucc_ep_map_from_array( + &UCC_TEAM_CTX_RANKS(team), team->size, + context->addr_storage.size, 1); + } else { + /* The caller's ep_map outlives the team, so aliasing it is safe */ + UCC_TEAM_CTX_MAP(team) = team->bp.params.ep_map; + } } else { - if (!team->ctx_ranks) { - team->ctx_ranks = + if (!UCC_TEAM_CTX_RANKS(team)) { + UCC_TEAM_CTX_RANKS(team) = ucc_malloc(team->size * sizeof(ucc_rank_t), "ctx_ranks"); - if (!team->ctx_ranks) { + if (!UCC_TEAM_CTX_RANKS(team)) { ucc_error("failed to allocate %zd bytes for ctx ranks array", team->size * sizeof(ucc_rank_t)); return UCC_ERR_NO_MEMORY; } - status = oob.allgather(&context->rank, team->ctx_ranks, + status = oob.allgather(&context->rank, UCC_TEAM_CTX_RANKS(team), sizeof(ucc_rank_t), oob.coll_info, &team->oob_req); if (UCC_OK != status) { ucc_error("failed to start oob allgather for proc info exchange"); - ucc_free(team->ctx_ranks); + ucc_free(UCC_TEAM_CTX_RANKS(team)); + UCC_TEAM_CTX_RANKS(team) = NULL; return status; } } @@ -375,11 +679,12 @@ static inline ucc_status_t ucc_team_exchange(ucc_context_t *context, } oob.req_free(team->oob_req); ucc_assert(team->size >= 2); - team->ctx_map = ucc_ep_map_from_array(&team->ctx_ranks, team->size, - context->addr_storage.size, 1); + UCC_TEAM_CTX_MAP(team) = + ucc_ep_map_from_array(&UCC_TEAM_CTX_RANKS(team), team->size, + context->addr_storage.size, 1); } ucc_debug("team %p rank %d, ctx_rank %d, map_type %d", team, team->rank, - context->rank, team->ctx_map.type); + context->rank, UCC_TEAM_CTX_MAP(team).type); return UCC_OK; } @@ -422,12 +727,155 @@ static ucc_status_t ucc_team_build_score_map(ucc_team_t *team) return status; } +/* Detach @team from the cache table and from whichever list it is on, and mark + it uncached. Callers hold @cache->lock, so the state change happens under the + lock as ucc_team_cache_state_t requires. The identity is left intact: the + teardown paths still log from it, and the callers that must release it do so + themselves. */ +static void ucc_team_cache_detach(ucc_team_cache_t *cache, ucc_team_t *team) +{ + ucc_team_cache_table_erase(cache, team); + ucc_team_cache_registry_remove(team); + team->cache_state = UCC_TEAM_CACHE_STATE_NONE; +} + +/* Add a freshly built, now ACTIVE team to its context cache */ +static void ucc_team_cache_admit(ucc_team_t *team) +{ + ucc_context_t *ctx = team->contexts[0]; + ucc_team_cache_t *cache = ctx->team_cache; + uint64_t new_cookie; + + team->cache_pending_insert = 0; + + /* The direct path leaves the vote zeroed, which reuse never consults */ + new_cookie = ucc_team_cache_vote_new_cookie(team->cache_vote_out); + if (new_cookie != 0 && new_cookie != ~(uint64_t)0) { + team->cache_identity.instance_cookie = new_cookie; + } + + if (cache == NULL) { + return; + } + + /* Both take cache->lock, so they must run before the insert below */ + ucc_team_cache_progress_pending(cache); + + if (cache->eviction != UCC_TEAM_CACHE_EVICTION_NONE && + cache->size >= cache->max_size) { + if (UCC_ERR_NO_RESOURCE == ucc_team_cache_evict_one(cache)) { + ucc_debug( + "team cache at pool-safe capacity (size=%u/%u), all entries " + "live; admitting team %p (hash=0x%" PRIx64 ") un-cached", + cache->size, + cache->max_size, + (void *)team, + team->cache_identity.hash); + } + } + + ucc_spin_lock(&cache->lock); + if (UCC_OK == ucc_team_cache_insert(cache, team) && + team->cache_state == UCC_TEAM_CACHE_STATE_DORMANT) { + ucc_list_del(&team->cache_link); + team->cache_state = UCC_TEAM_CACHE_STATE_LIVE; + ucc_team_cache_registry_add_live(cache, team); + ucc_debug( + "team cache: insert (hash=0x%" PRIx64 ") team %p -> LIVE " + "refcount=%d", + team->cache_identity.hash, + (void *)team, + team->refcount); + } else if (team->cache_state == UCC_TEAM_CACHE_STATE_NONE) { + ucc_debug( + "team cache: insert skipped (hash=0x%" PRIx64 ") team %p stays " + "uncached", + team->cache_identity.hash, + (void *)team); + } + ucc_spin_unlock(&cache->lock); +} + ucc_status_t ucc_team_create_test_single(ucc_context_t *context, ucc_team_t *team) { - ucc_status_t status = UCC_OK; + ucc_status_t status = UCC_OK; + ucc_team_cache_action_t agreed; switch (team->state) { + case UCC_TEAM_CACHE_AGREE: + status = ucc_service_coll_test(&team->cache_vote_req); + if (status == UCC_INPROGRESS) { + return UCC_INPROGRESS; + } + if (status < 0) { + ucc_service_coll_finalize(&team->cache_vote_req); + ucc_error( + "team cache: agreement vote failed: %s", + ucc_status_string(status)); + ucc_team_agreement_fail(context, team); + return status; + } + ucc_service_coll_finalize(&team->cache_vote_req); + + agreed = ucc_team_cache_vote_result(team->cache_vote_out); + if (agreed == UCC_TEAM_CACHE_ACTION_EXACT_REUSE && + team->cache_local_action == UCC_TEAM_CACHE_ACTION_EXACT_REUSE) { + ucc_team_cache_t *vote_cache = context->team_cache; + + ucc_spin_lock(&vote_cache->lock); + ucc_team_cache_get(team); /* RESERVED -> LIVE */ + ucc_team_cache_registry_make_live(vote_cache, team); + ucc_spin_unlock(&vote_cache->lock); + ucc_debug("team cache: agreed EXACT reuse, team %p", (void *)team); + team->state = UCC_TEAM_ACTIVE; + return UCC_OK; + } + /* Not unanimous, so every member builds a fresh team */ + if (team->cache_local_action == UCC_TEAM_CACHE_ACTION_EXACT_REUSE) { + /* Detach the rejected candidate and rebuild the same handle */ + ucc_team_cache_t *vote_cache = context->team_cache; + + ucc_spin_lock(&vote_cache->lock); + ucc_team_cache_table_erase(vote_cache, team); + ucc_team_cache_registry_remove(team); + ucc_spin_unlock(&vote_cache->lock); + team->cache_pending_insert = 1; + team->state = UCC_TEAM_CACHE_MISS_TEARDOWN; + ucc_debug( + "team cache: agreement lost, rebuilding team %p in place", + (void *)team); + /* fall through to CACHE_MISS_TEARDOWN */ + } else { + team->state = (team->size > 1) ? UCC_TEAM_ADDR_EXCHANGE + : UCC_TEAM_CL_CREATE; + return UCC_INPROGRESS; + } + /* fall through */ + case UCC_TEAM_CACHE_MISS_TEARDOWN: + /* A candidate that lost the vote is torn down rather than just marked + uncached: peers are about to build a fresh team reusing this team's + id, and its CL/TL teams still hold the matching wire tags. Destroying + them is the alias barrier that keeps a late in-flight message from + the retired team out of the rebuilt one's tag space. + + This state is re-entered on every progress call until the teardown + completes, so ucc_team_teardown_for_rebuild must be restartable: it + advances per component and returns UCC_INPROGRESS without repeating + the destroys it already finished. */ + status = ucc_team_teardown_for_rebuild(team); + if (status == UCC_INPROGRESS) { + ucc_context_progress(context); + return UCC_INPROGRESS; + } + if (status < 0) { + goto out; + } + status = ucc_team_reset_for_rebuild(context, team); + if (status < 0) { + goto out; + } + return UCC_INPROGRESS; /* re-enter at the reset start state */ case UCC_TEAM_ADDR_EXCHANGE: status = ucc_team_exchange(context, team); if (UCC_OK != status) { @@ -470,6 +918,9 @@ ucc_status_t ucc_team_create_test_single(ucc_context_t *context, break; case UCC_TEAM_ACTIVE: return UCC_OK; + case UCC_TEAM_CREATE_FAILED: + ucc_error("team %p: create already failed, handle is terminal", team); + return UCC_ERR_INVALID_PARAM; } out: if (UCC_OK == status) { @@ -488,6 +939,9 @@ ucc_status_t ucc_team_create_test_single(ucc_context_t *context, } /* TODO: add team/coll selection and check if some teams are never used after selection and clean them up */ + if (UCC_OK == status && team->cache_pending_insert) { + ucc_team_cache_admit(team); + } return status; } @@ -505,7 +959,8 @@ ucc_status_t ucc_team_create_test(ucc_team_h team) return ucc_team_create_test_single(team->contexts[0], team); } -static ucc_status_t ucc_team_destroy_single(ucc_team_h team) +/* Tear down a team; @for_rebuild keeps what reset_for_rebuild reuses */ +static ucc_status_t ucc_team_destroy_single_ex(ucc_team_h team, int for_rebuild) { ucc_cl_iface_t *cl_iface; int i; @@ -530,10 +985,14 @@ static ucc_status_t ucc_team_destroy_single(ucc_team_h team) team->cl_teams[i] = NULL; } - ucc_topo_cleanup(team->topo); + /* Safe here: the TL nested maps aliasing ctx_map are destroyed above */ + ucc_team_artifacts_put(team->artifacts); + team->artifacts = NULL; - if (team->contexts[0]->service_team && team->size > 1) { + /* Ownership, not service-team presence: a failed init leaves the caller's */ + if (team->internal_oob) { ucc_internal_oob_finalize(&team->bp.params.oob); + team->internal_oob = 0; } if ((ucc_global_config.log_component.log_level >= UCC_LOG_LEVEL_INFO) && @@ -541,16 +1000,64 @@ static ucc_status_t ucc_team_destroy_single(ucc_team_h team) ucc_info("team destroyed, team_id %d", team->id); } - ucc_coll_score_free_map(team->score_map); + if (team->score_map) { /* NULL for a shell that never finished building */ + ucc_coll_score_free_map(team->score_map); + team->score_map = NULL; + } ucc_free(team->addr_storage.storage); - ucc_free(team->ctx_ranks); ucc_team_release_id(team); + + if (for_rebuild) { + /* create_post_single reallocates cl_teams */ + ucc_free(team->cl_teams); + team->cl_teams = NULL; + memset(&team->addr_storage, 0, sizeof(team->addr_storage)); + return UCC_OK; + } ucc_free(team->cl_teams); ucc_free(team->contexts); + ucc_team_cache_identity_free(&team->cache_identity); ucc_free(team); return UCC_OK; } +static ucc_status_t ucc_team_destroy_single(ucc_team_h team) +{ + return ucc_team_destroy_single_ex(team, 0); +} + +/* Polls, returning UCC_INPROGRESS until every component has been destroyed */ +static ucc_status_t ucc_team_teardown_for_rebuild(ucc_team_t *team) +{ + return ucc_team_destroy_single_ex(team, 1); +} + +/* Return a torn-down team to its pre-build state, keeping its membership */ +static ucc_status_t ucc_team_reset_for_rebuild( + ucc_context_t *context, ucc_team_t *team) +{ + ucc_assert(team->service_team == NULL); + ucc_assert(team->sreq == NULL); + + team->n_cl_teams = 0; + team->seq_num = 0; + /* A pool id was released by the teardown and has to be redrawn */ + if (!UCC_TEAM_ID_IS_EXTERNAL(team)) { + team->id = 0; + } + team->bp.id = 0; + team->oob_req = NULL; + team->refcount = 1; + team->cache_state = UCC_TEAM_CACHE_STATE_NONE; + team->cache_pending_insert = 1; + ucc_list_head_init(&team->cache_link); + /* Teardown released the old holder, so re-init the inline one */ + team->artifacts = &team->artifacts_inline; + ucc_team_artifacts_init_inline(team->artifacts); + + return ucc_team_create_post_single(context, team); +} + ucc_status_t ucc_team_destroy(ucc_team_h team) { if (NULL == team) { @@ -558,16 +1065,192 @@ ucc_status_t ucc_team_destroy(ucc_team_h team) return UCC_ERR_INVALID_PARAM; } + if (team->state == UCC_TEAM_CREATE_FAILED) { + if (team->cache_state != UCC_TEAM_CACHE_STATE_NONE) { + /* Failed adoption: the cache already reclaimed this team */ + ucc_error("team %p failed to adopt a cached team; the handle is " + "not the caller's to destroy", + team); + return UCC_ERR_INVALID_PARAM; + } + return ucc_team_destroy_single(team); /* shell owned by this create */ + } + if (team->state != UCC_TEAM_ACTIVE) { ucc_error("team %p is used before team_create is completed", team); return UCC_ERR_INVALID_PARAM; } + /* A dormant team is cache-owned; only a live user may release it */ + if (team->cache_state == UCC_TEAM_CACHE_STATE_DORMANT || + team->cache_state == UCC_TEAM_CACHE_STATE_RESERVED) { + ucc_error("team %p is retained by the team cache and has no live " + "user; refusing to destroy it", + team); + return UCC_ERR_INVALID_PARAM; + } + /* we don't support multiple contexts per team yet */ ucc_assert(team->num_contexts == 1); + + /* A cached team is retained, keeping its id and its CL/TL teams */ + if (team->cache_state == UCC_TEAM_CACHE_STATE_LIVE) { + ucc_context_t *ctx = team->contexts[0]; + ucc_team_cache_t *cache = ctx->team_cache; + int n; + + ucc_assert(cache != NULL); + ucc_spin_lock(&cache->lock); + n = ucc_team_cache_put(team); /* LIVE -> DORMANT */ + ucc_team_cache_registry_make_dormant(cache, team); + ucc_spin_unlock(&cache->lock); + + ucc_debug( + "team cache: team %p now dormant, retained for reuse " + "(hash=0x%" PRIx64 ", live_users=%d)", + (void *)team, + team->cache_identity.hash, + n); + return UCC_OK; /* callers spin on UCC_INPROGRESS, so never return it */ + } + return ucc_team_destroy_single(team); } +/* Terminal teardown failure: leak the id too, surviving TL teams still use it */ +static void ucc_team_cache_abandon_failed(ucc_team_t *team, ucc_status_t status) +{ + ucc_error( + "cached team %p teardown failed terminally (%s); abandoning " + "team-id %u together with the partially destroyed component state", + (void *)team, + ucc_status_string(status), + (unsigned)team->id); +} + +void ucc_team_cache_drain(ucc_context_t *context) +{ + ucc_team_cache_t *cache = context->team_cache; + ucc_team_t *team, *tmp; + ucc_status_t status; + + if (cache == NULL) { + return; + } + + /* Context teardown is single threaded, so the walk needs no lock */ + ucc_list_for_each_safe (team, tmp, &cache->dormant, cache_link) { + ucc_assert(team->cache_state == UCC_TEAM_CACHE_STATE_DORMANT); + + ucc_team_cache_detach(cache, team); + + while (UCC_INPROGRESS == (status = ucc_team_destroy_single(team))) { + ucc_context_progress(context); + } + if (status < 0) { + ucc_team_cache_abandon_failed(team, status); + } + } + + ucc_assert(ucc_list_is_empty(&cache->dormant)); + + while (!ucc_list_is_empty(&cache->pending_destroy)) { + ucc_team_cache_progress_pending(cache); + if (!ucc_list_is_empty(&cache->pending_destroy)) { + ucc_context_progress(context); + } + } +} + +/* Teams are popped under the lock, since a destroy must not run holding it */ +void ucc_team_cache_progress_pending(ucc_team_cache_t *cache) +{ + ucc_team_t *team; + ucc_status_t status; + unsigned pending, i; + uint16_t team_id; + + if (cache == NULL) { + return; + } + + ucc_spin_lock(&cache->lock); + pending = (unsigned)ucc_list_length(&cache->pending_destroy); + ucc_spin_unlock(&cache->lock); + + /* Bounded by the initial count so a re-queued team is not retried here */ + for (i = 0; i < pending; i++) { + ucc_spin_lock(&cache->lock); + if (ucc_list_is_empty(&cache->pending_destroy)) { + ucc_spin_unlock(&cache->lock); + break; + } + team = ucc_list_extract_head(&cache->pending_destroy, ucc_team_t, + cache_link); + ucc_spin_unlock(&cache->lock); + + team_id = team->id; + status = ucc_team_destroy_single(team); + if (status == UCC_INPROGRESS) { + ucc_spin_lock(&cache->lock); + ucc_list_add_tail(&cache->pending_destroy, &team->cache_link); + ucc_spin_unlock(&cache->lock); + } else if (status < 0) { + ucc_team_cache_abandon_failed(team, status); + } else { + ucc_debug("team cache: evicted team destroy complete " + "(UCC_OK, id=%u); team-id pool headroom restored", + (unsigned)team_id); + } + } +} + +unsigned ucc_team_cache_progress_cb(void *arg) +{ + ucc_team_cache_t *cache = arg; + + /* Unlocked peek; a stale read only delays the retry to the next call */ + if (ucc_list_is_empty(&cache->pending_destroy)) { + return 0; + } + ucc_team_cache_progress_pending(cache); + return 0; +} + +ucc_status_t ucc_team_cache_evict_one(ucc_team_cache_t *cache) +{ + ucc_team_t *victim; + uint64_t hash; + + ucc_spin_lock(&cache->lock); + + victim = ucc_team_cache_pick_victim(cache); + if (victim == NULL) { + ucc_spin_unlock(&cache->lock); + return UCC_ERR_NO_RESOURCE; + } + ucc_assert(victim->cache_state == UCC_TEAM_CACHE_STATE_DORMANT); + hash = victim->cache_identity.hash; + + ucc_team_cache_detach(cache, victim); + ucc_list_add_tail(&cache->pending_destroy, &victim->cache_link); + cache->stats.evictions++; + + ucc_debug( + "team cache %p: evicting dormant team %p (hash=0x%" PRIx64 + ", size now %u, evictions=%" PRIu64 ")", + (void *)cache, + (void *)victim, + hash, + cache->size, + cache->stats.evictions); + + ucc_spin_unlock(&cache->lock); + + ucc_team_cache_progress_pending(cache); + return UCC_OK; +} + int ucc_team_id_pool_ffs_clear(uint64_t *value) { int i; @@ -621,6 +1304,8 @@ static ucc_status_t ucc_team_alloc_id(ucc_team_t *team) ucc_subset_t subset = {.map.type = UCC_EP_MAP_FULL, .map.ep_num = team->size, .myrank = team->rank}; + /* Let finished evictions return their ids before voting on the pool */ + ucc_team_cache_progress_pending(ctx->team_cache); status = ucc_service_allreduce(team, local, global, UCC_DT_UINT64, ctx->ids.pool_size, UCC_OP_BAND, subset, diff --git a/src/core/ucc_team.h b/src/core/ucc_team.h index 3318ad26d97..a01525f06b1 100644 --- a/src/core/ucc_team.h +++ b/src/core/ucc_team.h @@ -10,22 +10,48 @@ #include "ucc/api/ucc.h" #include "utils/ucc_datastruct.h" #include "utils/ucc_coll_utils.h" +#include "utils/ucc_list.h" #include "ucc_context.h" +#include "ucc_team_cache.h" #include "utils/ucc_math.h" #include "components/base/ucc_base_iface.h" #include "components/cl/ucc_cl.h" #include "components/tl/ucc_tl.h" #include "coll_score/ucc_coll_score.h" +#include "ucc_service_coll.h" /* ucc_service_coll_req_t is embedded below */ -typedef struct ucc_service_coll_req ucc_service_coll_req_t; typedef enum { - UCC_TEAM_ADDR_EXCHANGE, + UCC_TEAM_ADDR_EXCHANGE, /* zero, so it is the calloc default */ UCC_TEAM_SERVICE_TEAM, UCC_TEAM_ALLOC_ID, UCC_TEAM_CL_CREATE, UCC_TEAM_ACTIVE, + UCC_TEAM_CACHE_AGREE, /* cache-action vote in flight */ + UCC_TEAM_CACHE_MISS_TEARDOWN, /* vote lost, draining before a rebuild */ + /* Terminal, vote failed: destroy frees a shell, rejects a cache-owned one */ + UCC_TEAM_CREATE_FAILED, } ucc_team_state_t; +/* Refcounted holder of the per-team state that derived teams may share */ +typedef struct ucc_team_artifacts { + ucc_ep_map_t ctx_map; /*< map to the ctx ranks, set if CTX is global */ + ucc_rank_t *ctx_ranks; /*< UCC-owned backing array of ctx_map, or NULL */ + ucc_topo_t *topo; /*< subset topology */ + int refcount; /*< number of teams pointing at this holder */ + ucc_spinlock_t lock; /*< guards refcount */ + uint8_t heap; /*< 1: heap-allocated, 0: embedded in a team */ +} ucc_team_artifacts_t; + +/* Init an embedded holder in place: refcount 1, heap 0 */ +void ucc_team_artifacts_init_inline(ucc_team_artifacts_t *a); + +/* Drop a reference; at zero release contents, free the struct if heap */ +void ucc_team_artifacts_put(ucc_team_artifacts_t *a); + +#define UCC_TEAM_CTX_MAP(_team) ((_team)->artifacts->ctx_map) +#define UCC_TEAM_CTX_RANKS(_team) ((_team)->artifacts->ctx_ranks) +#define UCC_TEAM_TOPO(_team) ((_team)->artifacts->topo) + typedef struct ucc_team { ucc_team_state_t state; ucc_context_t ** contexts; @@ -41,13 +67,21 @@ typedef struct ucc_team { ucc_tl_team_t * service_team; ucc_service_coll_req_t *sreq; ucc_addr_storage_t addr_storage; /*< addresses of team endpoints */ - ucc_rank_t * ctx_ranks; void * oob_req; - ucc_ep_map_t ctx_map; /*< map to the ctx ranks, defined if CTX - type is global (oob provided) */ - ucc_topo_t *topo; + ucc_team_artifacts_t *artifacts; /*< ctx_map/ctx_ranks/topo holder */ + ucc_team_artifacts_t artifacts_inline; /*< holder used when not shared */ ucc_score_map_t *score_map; /*< score map of CLs */ uint32_t seq_num; + int refcount; /* live teams backing a cache entry */ + ucc_team_cache_identity_t cache_identity; + ucc_list_link_t cache_link; /* live, dormant or reserved list */ + ucc_team_cache_state_t cache_state; + int cache_pending_insert; /* cacheable, not yet in */ + ucc_team_cache_action_t cache_local_action; /* this rank's vote */ + ucc_service_coll_req_t cache_vote_req; /* embedded, never freed */ + uint64_t cache_vote_in[UCC_TEAM_CACHE_VOTE_LANES]; + uint64_t cache_vote_out[UCC_TEAM_CACHE_VOTE_LANES]; + uint8_t internal_oob; /* bp.params.oob is UCC-owned */ } ucc_team_t; /* If the bit is set then team_id is provided by the user */ @@ -61,6 +95,18 @@ void ucc_copy_team_params(ucc_team_params_t *dst, const ucc_team_params_t *src); int ucc_team_id_pool_ffs_clear(uint64_t *value); void ucc_team_id_pool_set_bit(uint64_t *local, int id); +/* Destroy every dormant team; call before the CL/TL contexts are destroyed */ +void ucc_team_cache_drain(ucc_context_t *context); + +/* Drive one teardown attempt for each team on the pending-destroy list */ +void ucc_team_cache_progress_pending(ucc_team_cache_t *cache); + +/* Context progress callback for the above; @arg is the cache */ +unsigned ucc_team_cache_progress_cb(void *arg); + +/* Move the eviction victim to the pending-destroy list and start its teardown */ +ucc_status_t ucc_team_cache_evict_one(ucc_team_cache_t *cache); + /* Returns addressing information for "rank" in a team. If ucc context was created with OOB then addr storage is located on context. In that case we need to map rank to ctx_rank first. Otherwise, addr @@ -77,7 +123,8 @@ ucc_get_team_ep_header(ucc_context_t *context, ucc_team_t *team, : &team->addr_storage; ucc_rank_t storage_rank = context->addr_storage.storage - ? (team ? ucc_ep_map_eval(team->ctx_map, rank) : rank) + ? (team ? ucc_ep_map_eval(UCC_TEAM_CTX_MAP(team), rank) + : rank) : rank; return UCC_ADDR_STORAGE_RANK_HEADER(storage, storage_rank); @@ -107,12 +154,14 @@ static inline void *ucc_get_team_ep_addr(ucc_context_t *context, static inline ucc_rank_t ucc_get_ctx_rank(ucc_team_t *team, ucc_rank_t team_rank) { - return ucc_ep_map_eval(team->ctx_map, team_rank); + return ucc_ep_map_eval(UCC_TEAM_CTX_MAP(team), team_rank); } static inline ucc_host_id_t ucc_team_rank_host_id(ucc_rank_t rank, ucc_team_t *team) { - return team->topo->topo->procs[ucc_get_ctx_rank(team, rank)].host_id; + ucc_topo_t *topo = UCC_TEAM_TOPO(team); + + return topo->topo->procs[ucc_get_ctx_rank(team, rank)].host_id; } static inline int ucc_team_ranks_on_same_node(ucc_rank_t rank1, ucc_rank_t rank2, diff --git a/src/core/ucc_team_cache.c b/src/core/ucc_team_cache.c new file mode 100644 index 00000000000..e84d4c42923 --- /dev/null +++ b/src/core/ucc_team_cache.c @@ -0,0 +1,477 @@ +/** + * Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * See file LICENSE for terms. + */ + +#include "config.h" +#include "ucc_team_cache.h" +#include "ucc_team.h" +#include "utils/ucc_malloc.h" +#include "utils/ucc_log.h" +#include "utils/ucc_coll_utils.h" +#include "utils/ucc_compiler_def.h" +#include "utils/khash.h" +#include +#include + +/* One team per key; the cache holds the table as an opaque void * */ +KHASH_MAP_INIT_INT64(ucc_team_cache_map, ucc_team_t *) +typedef khash_t(ucc_team_cache_map) ucc_team_cache_map_t; + +/* Order must match ucc_team_cache_eviction_policy_t */ +const char *ucc_team_cache_eviction_names[] = { + [UCC_TEAM_CACHE_EVICTION_NONE] = "none", + [UCC_TEAM_CACHE_EVICTION_FIFO] = "fifo", + NULL}; + +ucc_status_t ucc_team_cache_init( + ucc_team_cache_t **cache, uint32_t max_size, + ucc_team_cache_eviction_policy_t eviction, uint32_t disable_linear_check) +{ + ucc_team_cache_t *c; + + c = ucc_calloc(1, sizeof(*c), "ucc_team_cache"); + if (ucc_unlikely(!c)) { + ucc_error("failed to allocate ucc_team_cache_t"); + return UCC_ERR_NO_MEMORY; + } + + c->table = kh_init(ucc_team_cache_map); + if (ucc_unlikely(!c->table)) { + ucc_error("failed to allocate ucc_team_cache hash table"); + goto err; + } + + ucc_list_head_init(&c->live); + ucc_list_head_init(&c->dormant); + ucc_list_head_init(&c->reserved); + ucc_list_head_init(&c->pending_destroy); + ucc_spinlock_init(&c->lock, 0); + + c->max_size = max_size; + c->eviction = eviction; + c->disable_linear_check = disable_linear_check; + + ucc_debug( + "ucc_team_cache created: %p, max_size=%u, eviction=%s, " + "disable_linear_check=%u", + (void *)c, + max_size, + ucc_team_cache_eviction_names[eviction], + disable_linear_check); + *cache = c; + return UCC_OK; + +err: + ucc_free(c); + return UCC_ERR_NO_MEMORY; +} + +void ucc_team_cache_destroy(ucc_team_cache_t *cache) +{ + if (!cache) { + return; + } + + if (cache->size != 0) { + ucc_warn( + "ucc_team_cache_destroy called with %u entries still present", + cache->size); + } + + if (!ucc_list_is_empty(&cache->pending_destroy)) { + ucc_warn( + "ucc_team_cache_destroy called with pending-destroy entries " + "still present (eviction teardown not flushed)"); + } + + /* All teams are quiesced at context destroy, so no vote is in flight */ + ucc_assert(ucc_list_is_empty(&cache->reserved)); + + 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", (void *)cache); + ucc_free(cache); +} + +void ucc_team_cache_dump_stats(ucc_team_cache_t *cache) +{ + double hit_rate = 0.0; + + if (!cache) { + return; + } + + if (cache->stats.lookups > 0) { + hit_rate = (cache->stats.hits * 100.0) / cache->stats.lookups; + } + + ucc_info( + "team_cache stats: lookups=%" PRIu64 " hits=%" PRIu64 " (%.1f%%) " + "misses=%" PRIu64 " inserts=%" PRIu64 " evictions=%" PRIu64, + cache->stats.lookups, + cache->stats.hits, + hit_rate, + cache->stats.misses, + cache->stats.inserts, + cache->stats.evictions); +} + +/* FNV-1a over {size, self_ep, members}, used only as a bucket key */ +#define UCC_TEAM_CACHE_FNV1A_OFFSET 0xcbf29ce484222325ULL +#define UCC_TEAM_CACHE_FNV1A_PRIME 0x00000100000001b3ULL + +static void ucc_team_cache_fnv1a_accumulate( + uint64_t *h, const void *data, size_t len) +{ + const uint8_t *p = (const uint8_t *)data; + size_t b; + + for (b = 0; b < len; b++) { + *h ^= (uint64_t)p[b]; + *h *= UCC_TEAM_CACHE_FNV1A_PRIME; + } +} + +ucc_status_t ucc_team_cache_identity_build( + const ucc_team_params_t *params, ucc_team_cache_identity_t *identity) +{ + ucc_rank_t size; + ucc_rank_t self_ep; + ucc_rank_t *members; + ucc_rank_t i; + uint64_t h; + + if (!(params->mask & UCC_TEAM_PARAM_FIELD_EP_MAP)) { + ucc_debug("team cache identity: no EP_MAP in params, not cacheable"); + return UCC_ERR_INVALID_PARAM; + } + + size = (ucc_rank_t)params->ep_map.ep_num; + self_ep = (params->mask & UCC_TEAM_PARAM_FIELD_EP) ? (ucc_rank_t)params->ep + : UCC_RANK_INVALID; + + /* An external id is its own id/tag domain, so it is part of the identity */ + identity->ext_id = ((params->mask & UCC_TEAM_PARAM_FIELD_ID) && + (params->id <= UCC_TEAM_ID_MAX)) + ? (uint16_t)(((uint16_t)params->id) | + UCC_TEAM_ID_EXTERNAL_BIT) + : 0; + + if (size < 1) { + ucc_debug("team cache identity: empty ep_map"); + return UCC_ERR_INVALID_PARAM; + } + + members = ucc_malloc(size * sizeof(*members), "team_cache_members"); + if (ucc_unlikely(!members)) { + ucc_error( + "failed to allocate %zu bytes for team cache members", + size * sizeof(*members)); + return UCC_ERR_NO_MEMORY; + } + + /* Copy the map so the identity never aliases caller-owned storage */ + for (i = 0; i < size; i++) { + members[i] = ucc_ep_map_eval(params->ep_map, i); + } + + identity->size = size; + identity->self_ep = self_ep; + identity->members = members; + identity->instance_cookie = 0; /* stamped later by the agreement vote */ + + /* ext_id is deliberately not hashed, so the bucket is membership only */ + h = UCC_TEAM_CACHE_FNV1A_OFFSET; + ucc_team_cache_fnv1a_accumulate(&h, &size, sizeof(size)); + ucc_team_cache_fnv1a_accumulate(&h, &self_ep, sizeof(self_ep)); + ucc_team_cache_fnv1a_accumulate( + &h, members, (size_t)size * sizeof(*members)); + identity->hash = h; + + return UCC_OK; +} + +int ucc_team_cache_identity_equal( + const ucc_team_cache_identity_t *a, const ucc_team_cache_identity_t *b) +{ + return ucc_team_cache_identity_equal_membership(a, b) && + a->ext_id == b->ext_id; +} + +int ucc_team_cache_identity_equal_membership( + const ucc_team_cache_identity_t *a, const ucc_team_cache_identity_t *b) +{ + if (a->hash != b->hash || a->size != b->size || a->self_ep != b->self_ep) { + return 0; + } + return memcmp( + a->members, b->members, (size_t)a->size * sizeof(*a->members)) == + 0; +} + +void ucc_team_cache_identity_free(ucc_team_cache_identity_t *identity) +{ + if (!identity) { + return; + } + ucc_free(identity->members); + memset(identity, 0, sizeof(*identity)); +} + +uint64_t ucc_team_cache_next_cookie(ucc_team_cache_t *cache) +{ + uint64_t c = ++cache->cache_gen; /* 0 is the unstamped sentinel */ + if (ucc_unlikely(c == 0)) { + c = ++cache->cache_gen; + } + return c; +} + +void ucc_team_cache_vote_fill( + uint64_t *v, int prepared, ucc_team_cache_action_t action, uint64_t key, + uint64_t cookie, uint64_t parent_cookie, int is_rank0, + uint64_t proposed_cookie) +{ + if (!prepared) { + /* All-ones equality lanes are a BAND no-op */ + v[0] = 0; + v[1] = ~(uint64_t)0; + v[2] = ~(uint64_t)0; + v[3] = ~(uint64_t)0; + v[4] = ~(uint64_t)0; + v[5] = ~(uint64_t)0; + v[6] = ~(uint64_t)0; + v[7] = ~(uint64_t)0; + v[8] = ~(uint64_t)0; + } else { + v[0] = 1; + v[1] = (uint64_t)action; + v[2] = ~(uint64_t)action; + v[3] = key; + v[4] = ~key; + v[5] = cookie; + v[6] = ~cookie; + v[7] = parent_cookie; + v[8] = ~parent_cookie; + } + /* Only rank 0 contributes here, so every member reads its value */ + v[9] = is_rank0 ? proposed_cookie : ~(uint64_t)0; +} + +ucc_team_cache_action_t ucc_team_cache_vote_result(const uint64_t *v) +{ + int all_prepared = (v[0] == 1); + int action_agree = (v[1] == ~v[2]); + int key_agree = (v[3] == ~v[4]); + int cookie_agree = (v[5] == ~v[6]); + int pcookie_agree = (v[7] == ~v[8]); + + if (all_prepared && action_agree && key_agree && cookie_agree && + pcookie_agree) { + return (ucc_team_cache_action_t)v[1]; + } + return UCC_TEAM_CACHE_ACTION_MISS; +} + +uint64_t ucc_team_cache_vote_new_cookie(const uint64_t *v) +{ + return v[9]; +} + +int ucc_team_cache_is_cacheable(const ucc_team_params_t *params) +{ + /* These are not part of the identity, so a reuse could change semantics */ + uint64_t optional = UCC_TEAM_PARAM_FIELD_ORDERING | + UCC_TEAM_PARAM_FIELD_OUTSTANDING_COLLS | + UCC_TEAM_PARAM_FIELD_SYNC_TYPE | + UCC_TEAM_PARAM_FIELD_P2P_CONN | + UCC_TEAM_PARAM_FIELD_MEM_PARAMS; + + return (params->mask & optional) == 0; +} + +/* A cached team is on exactly one of live/dormant/reserved, or on none */ + +void ucc_team_cache_registry_add_live(ucc_team_cache_t *cache, ucc_team_t *team) +{ + ucc_list_add_tail(&cache->live, &team->cache_link); +} + +void ucc_team_cache_registry_make_dormant( + ucc_team_cache_t *cache, ucc_team_t *team) +{ + ucc_list_del(&team->cache_link); + ucc_list_add_tail(&cache->dormant, &team->cache_link); +} + +void ucc_team_cache_registry_make_live( + ucc_team_cache_t *cache, ucc_team_t *team) +{ + ucc_list_del(&team->cache_link); + ucc_list_add_tail(&cache->live, &team->cache_link); +} + +void ucc_team_cache_registry_make_reserved( + ucc_team_cache_t *cache, ucc_team_t *team) +{ + /* Stays in the bucket, but no lookup, eviction or drain can reach it */ + ucc_list_del(&team->cache_link); + ucc_list_add_tail(&cache->reserved, &team->cache_link); +} + +void ucc_team_cache_registry_remove(ucc_team_t *team) +{ + ucc_list_del(&team->cache_link); +} + +void ucc_team_cache_table_erase(ucc_team_cache_t *cache, ucc_team_t *team) +{ + ucc_team_cache_map_t *h = (ucc_team_cache_map_t *)cache->table; + uint64_t hash = team->cache_identity.hash; + khiter_t k; + + k = kh_get(ucc_team_cache_map, h, hash); + if (k == kh_end(h) || kh_value(h, k) != team) { + /* Absent, or a colliding team owns the bucket */ + return; + } + kh_del(ucc_team_cache_map, h, k); + ucc_assert(cache->size > 0); + cache->size--; +} + +/* NONE --insert--> DORMANT --get--> LIVE --put-to-0--> DORMANT */ + +ucc_team_t *ucc_team_cache_lookup( + ucc_team_cache_t *cache, const ucc_team_cache_identity_t *id) +{ + ucc_team_cache_map_t *h = (ucc_team_cache_map_t *)cache->table; + khiter_t k; + ucc_team_t *team; + + cache->stats.lookups++; + + k = kh_get(ucc_team_cache_map, h, id->hash); + if (k == kh_end(h)) { + cache->stats.misses++; + ucc_debug( + "team_cache %p: lookup miss (hash=0x%" PRIx64 ")", + (void *)cache, + id->hash); + return NULL; + } + + team = kh_value(h, k); + + /* A different ext_id is a different tag domain, so not a valid reuse */ + if (team->cache_identity.ext_id != id->ext_id) { + cache->stats.misses++; + return NULL; + } + + if (!cache->disable_linear_check && + !ucc_team_cache_identity_equal_membership(&team->cache_identity, id)) { + cache->stats.misses++; + return NULL; + } + + /* Only a DORMANT team is free to re-adopt */ + if (team->cache_state != UCC_TEAM_CACHE_STATE_DORMANT) { + cache->stats.misses++; + return NULL; + } + + cache->stats.hits++; + ucc_debug( + "team_cache %p: lookup HIT team %p (hash=0x%" PRIx64 ")", + (void *)cache, + (void *)team, + id->hash); + return team; +} + +ucc_status_t ucc_team_cache_insert(ucc_team_cache_t *cache, ucc_team_t *team) +{ + ucc_team_cache_map_t *h = (ucc_team_cache_map_t *)cache->table; + uint64_t hash = team->cache_identity.hash; + khiter_t k; + int ret; + + if (cache->size >= cache->max_size) { + ucc_info( + "team_cache %p: full (size=%u, max=%u) - team %p not cached", + (void *)cache, + cache->size, + cache->max_size, + (void *)team); + return UCC_OK; /* cache_state stays NONE, so destroy tears it down */ + } + + /* An occupied bucket leaves the new team uncached but still functional */ + k = kh_get(ucc_team_cache_map, h, hash); + if (k != kh_end(h)) { + ucc_info( + "team_cache %p: bucket occupied (hash=0x%" PRIx64 + ") - team %p not cached (duplicate or hash collision)", + (void *)cache, + hash, + (void *)team); + return UCC_OK; + } + + k = kh_put(ucc_team_cache_map, h, hash, &ret); + if (ucc_unlikely(ret < 0)) { + ucc_error( + "team_cache %p: kh_put failed for hash=0x%" PRIx64, + (void *)cache, + hash); + return UCC_ERR_NO_MEMORY; + } + kh_value(h, k) = team; + + /* The caller immediately promotes this to LIVE via registry_make_live */ + ucc_list_add_tail(&cache->dormant, &team->cache_link); + + team->cache_state = UCC_TEAM_CACHE_STATE_DORMANT; + cache->size++; + cache->stats.inserts++; + + ucc_debug( + "team_cache %p: inserted team %p (hash=0x%" PRIx64 ", size=%u)", + (void *)cache, + (void *)team, + hash, + cache->size); + return UCC_OK; +} + +void ucc_team_cache_get(ucc_team_t *team) +{ + team->refcount++; + team->cache_state = UCC_TEAM_CACHE_STATE_LIVE; +} + +int ucc_team_cache_put(ucc_team_t *team) +{ + int rc; + + ucc_assert(team->refcount > 0); /* catches a double put in debug builds */ + rc = --team->refcount; + if (rc == 0) { + team->cache_state = UCC_TEAM_CACHE_STATE_DORMANT; + } + return rc; +} + +ucc_team_t *ucc_team_cache_pick_victim(ucc_team_cache_t *cache) +{ + if (ucc_list_is_empty(&cache->dormant)) { + ucc_debug( + "team_cache %p: pick_victim - dormant list empty", (void *)cache); + return NULL; + } + + /* For FIFO and NONE alike, the list head is the oldest insert */ + return ucc_list_head(&cache->dormant, ucc_team_t, cache_link); +} diff --git a/src/core/ucc_team_cache.h b/src/core/ucc_team_cache.h new file mode 100644 index 00000000000..08cc9a4815e --- /dev/null +++ b/src/core/ucc_team_cache.h @@ -0,0 +1,167 @@ +/** + * Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * See file LICENSE for terms. + */ + +#ifndef UCC_TEAM_CACHE_H_ +#define UCC_TEAM_CACHE_H_ + +#include "config.h" +#include "ucc/api/ucc.h" +#include "ucc/api/ucc_status.h" +#include "utils/ucc_datastruct.h" +#include "utils/ucc_list.h" +#include "utils/ucc_spinlock.h" +#include + +typedef struct ucc_team ucc_team_t; +typedef struct ucc_context ucc_context_t; + +/* NONE -> DORMANT -> RESERVED -> LIVE, all changed under cache->lock */ +typedef enum ucc_team_cache_state { + UCC_TEAM_CACHE_STATE_NONE = 0, /* never cached */ + UCC_TEAM_CACHE_STATE_DORMANT = 1, /* cached, no live backing team */ + UCC_TEAM_CACHE_STATE_RESERVED = 2, /* pinned for an in-flight vote */ + UCC_TEAM_CACHE_STATE_LIVE = 3, /* cached, backing one live team */ +} ucc_team_cache_state_t; + +/* Order must match ucc_team_cache_eviction_names[] */ +typedef enum ucc_team_cache_eviction_policy { + UCC_TEAM_CACHE_EVICTION_NONE = 0, + UCC_TEAM_CACHE_EVICTION_FIFO = 1, +} ucc_team_cache_eviction_policy_t; + +/* UCC_TEAM_CACHE_EVICTION choices, indexed by the enum, NULL terminated */ +extern const char *ucc_team_cache_eviction_names[]; + +/* Normalized team identity; @hash covers membership only, not @ext_id */ +typedef struct ucc_team_cache_identity { + uint64_t hash; + ucc_rank_t size; + ucc_rank_t self_ep; + uint16_t ext_id; /* 0 for pool-id teams */ + /* Per-adoption stamp that identifies WHICH cached team was picked. + Normally the vote's key lane carries @ext_id, which is enough to prove + every rank selected the same entry. RESEAT breaks that: its lookup + deliberately ignores @ext_id so a drifted id can still match on + membership, and two ranks holding different dormant teams of identical + membership would then agree on a lane that no longer distinguishes them. + The cookie is stamped when a team is adopted and voted on instead, so the + agreement still proves a common choice. 0 until the vote stamps it. */ + uint64_t instance_cookie; + ucc_rank_t *members; /* heap owned, length == size */ +} ucc_team_cache_identity_t; + +/* Materialize membership from @params into @identity, which owns its members */ +ucc_status_t ucc_team_cache_identity_build( + const ucc_team_params_t *params, ucc_team_cache_identity_t *identity); + +/* Compare membership and ext_id */ +int ucc_team_cache_identity_equal( + const ucc_team_cache_identity_t *a, const ucc_team_cache_identity_t *b); + +/* Compare membership only */ +int ucc_team_cache_identity_equal_membership( + const ucc_team_cache_identity_t *a, const ucc_team_cache_identity_t *b); + +/* Free @identity->members and zero @identity; idempotent */ +void ucc_team_cache_identity_free(ucc_team_cache_identity_t *identity); + +/* Non-zero if no optional behavioral team param is set in params->mask */ +int ucc_team_cache_is_cacheable(const ucc_team_params_t *params); + +typedef enum ucc_team_cache_action { + UCC_TEAM_CACHE_ACTION_MISS = 0, /* fresh full build */ + UCC_TEAM_CACHE_ACTION_EXACT_REUSE = 1, /* re-adopt a DORMANT team */ +} ucc_team_cache_action_t; + +/* Vote lanes: [0] prepared, [1..8] (value, ~value) pairs, [9] new cookie */ +#define UCC_TEAM_CACHE_VOTE_LANES 10 + +/* Fill one rank's vote buffer @v ahead of the UCC_OP_BAND allreduce */ +void ucc_team_cache_vote_fill( + uint64_t *v, int prepared, ucc_team_cache_action_t action, uint64_t key, + uint64_t cookie, uint64_t parent_cookie, int is_rank0, + uint64_t proposed_cookie); + +/* Agreed action from a BAND reduced vote buffer, or MISS if not unanimous */ +ucc_team_cache_action_t ucc_team_cache_vote_result(const uint64_t *v); + +/* Team rank 0's proposed instance cookie, from a BAND reduced vote buffer */ +uint64_t ucc_team_cache_vote_new_cookie(const uint64_t *v); + +typedef struct ucc_team_cache_stats { + uint64_t lookups; + uint64_t hits; + uint64_t misses; + uint64_t evictions; + uint64_t inserts; +} ucc_team_cache_stats_t; + +/* All fields are protected by @lock */ +typedef struct ucc_team_cache { + void *table; /* khash key64 -> ucc_team_t* */ + ucc_list_link_t live; + ucc_list_link_t dormant; /* head = oldest victim */ + ucc_list_link_t reserved; + ucc_list_link_t pending_destroy; + ucc_spinlock_t lock; + uint32_t max_size; + uint32_t size; + ucc_team_cache_eviction_policy_t eviction; + uint32_t disable_linear_check; + uint32_t dump_stats; + uint32_t agreement; + uint64_t cache_gen; /* source of instance cookies */ + ucc_team_cache_stats_t stats; +} ucc_team_cache_t; + +/* Reserve the next instance cookie; team rank 0 only, under @cache->lock */ +uint64_t ucc_team_cache_next_cookie(ucc_team_cache_t *cache); + +ucc_status_t ucc_team_cache_init( + ucc_team_cache_t **cache, uint32_t max_size, + ucc_team_cache_eviction_policy_t eviction, uint32_t disable_linear_check); + +/* Free @cache, which the caller has already drained; may be NULL */ +void ucc_team_cache_destroy(ucc_team_cache_t *cache); + +/* Find a DORMANT team by exact identity, or NULL; under @cache->lock */ +ucc_team_t *ucc_team_cache_lookup( + ucc_team_cache_t *cache, const ucc_team_cache_identity_t *id); + +/* Insert @team as DORMANT; a full cache or hash collision skips the insert */ +ucc_status_t ucc_team_cache_insert(ucc_team_cache_t *cache, ucc_team_t *team); + +/* Adopt a DORMANT team: refcount++, DORMANT -> LIVE; under @cache->lock */ +void ucc_team_cache_get(ucc_team_t *team); + +/* Release a LIVE team: refcount--, LIVE -> DORMANT at zero; returns refcount */ +int ucc_team_cache_put(ucc_team_t *team); + +/* Eviction victim per @cache->eviction, or NULL if no team is dormant */ +ucc_team_t *ucc_team_cache_pick_victim(ucc_team_cache_t *cache); + +/* Registry helpers below move @team between lists, all under @cache->lock */ +void ucc_team_cache_registry_add_live( + ucc_team_cache_t *cache, ucc_team_t *team); + +void ucc_team_cache_registry_make_dormant( + ucc_team_cache_t *cache, ucc_team_t *team); + +void ucc_team_cache_registry_make_live( + ucc_team_cache_t *cache, ucc_team_t *team); + +void ucc_team_cache_registry_make_reserved( + ucc_team_cache_t *cache, ucc_team_t *team); + +/* Remove @team from whichever list it is on */ +void ucc_team_cache_registry_remove(ucc_team_t *team); + +/* Erase @team from the bucket table and decrement @cache->size */ +void ucc_team_cache_table_erase(ucc_team_cache_t *cache, ucc_team_t *team); + +/* Log hit/miss/eviction counters, for UCC_TEAM_CACHE_DUMP_STATS */ +void ucc_team_cache_dump_stats(ucc_team_cache_t *cache); + +#endif /* UCC_TEAM_CACHE_H_ */ diff --git a/src/utils/ucc_coll_utils.c b/src/utils/ucc_coll_utils.c index f5c185fc282..3707549df22 100644 --- a/src/utils/ucc_coll_utils.c +++ b/src/utils/ucc_coll_utils.c @@ -768,7 +768,7 @@ void ucc_coll_str(const ucc_coll_task_t *task, char *str, size_t len, ucc_snprintf_safe(task_info, sizeof(task_info), " rank %u, ctx_rank %u, seq_num %d, req %p", team->rank, - ucc_ep_map_eval(team->ctx_map, team->rank), + ucc_ep_map_eval(UCC_TEAM_CTX_MAP(team), team->rank), task->seq_num, task); strncat(str, task_info, len - strlen(str)); } diff --git a/test/gtest/Makefile.am b/test/gtest/Makefile.am index ef2eb28b9a8..8b29a1588b2 100644 --- a/test/gtest/Makefile.am +++ b/test/gtest/Makefile.am @@ -78,6 +78,7 @@ gtest_SOURCES = \ core/test_mc_reduce.cc \ core/test_team.cc \ core/test_team_id_pool.cc \ + core/test_team_cache.cc \ core/test_schedule.cc \ core/test_topo.cc \ core/test_service_coll.cc \ diff --git a/test/gtest/core/test_team_cache.cc b/test/gtest/core/test_team_cache.cc new file mode 100644 index 00000000000..6ee6009ea05 --- /dev/null +++ b/test/gtest/core/test_team_cache.cc @@ -0,0 +1,1560 @@ +/** + * Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * See file LICENSE for terms. + */ +extern "C" { +#include "core/ucc_team_cache.h" +#include "core/ucc_team.h" +#include "core/ucc_context.h" +#include "utils/ucc_spinlock.h" +} +#include +#include +#include +#include +#include +#include +#include +#include + +/* Unit tests for the team-cache identity (build/hash/equal/free), the + cacheability policy, the locked cache API and eviction, with no MPI job. */ + +/* CB closure returning member[ep] from a heap vector (mimics OMPI coll/ucc's + rank_map_cb); freeable after build to prove identity does not retain it. */ +struct cb_ctx { + std::vector members; +}; + +static uint64_t member_cb(uint64_t ep, void *ctx) +{ + cb_ctx *c = static_cast(ctx); + return (uint64_t)c->members[ep]; +} + +static ucc_team_params_t make_cb_params(cb_ctx *ctx, ucc_rank_t self_ep) +{ + ucc_team_params_t p; + memset(&p, 0, sizeof(p)); + p.mask = UCC_TEAM_PARAM_FIELD_EP_MAP | UCC_TEAM_PARAM_FIELD_EP; + p.ep = self_ep; + p.ep_map.type = UCC_EP_MAP_CB; + p.ep_map.ep_num = ctx->members.size(); + p.ep_map.cb.cb = member_cb; + p.ep_map.cb.cb_ctx = ctx; + return p; +} + +/* ARRAY+OOB style, as OpenSHMEM scoll/ucc passes it: a user-owned array. */ +static ucc_team_params_t make_array_params( + ucc_rank_t *arr, ucc_rank_t size, ucc_rank_t self_ep) +{ + ucc_team_params_t p; + memset(&p, 0, sizeof(p)); + p.mask = UCC_TEAM_PARAM_FIELD_EP_MAP | + UCC_TEAM_PARAM_FIELD_EP; + p.ep = self_ep; + p.ep_map.type = UCC_EP_MAP_ARRAY; + p.ep_map.ep_num = size; + p.ep_map.array.map = arr; + p.ep_map.array.elem_size = sizeof(ucc_rank_t); + return p; +} + +static ucc_team_params_t make_strided_params( + uint64_t start, int64_t stride, ucc_rank_t size, ucc_rank_t self_ep) +{ + ucc_team_params_t p; + memset(&p, 0, sizeof(p)); + p.mask = UCC_TEAM_PARAM_FIELD_EP_MAP | + UCC_TEAM_PARAM_FIELD_EP; + p.ep = self_ep; + p.ep_map.type = UCC_EP_MAP_STRIDED; + p.ep_map.ep_num = size; + p.ep_map.strided.start = start; + p.ep_map.strided.stride = stride; + return p; +} + +/* ARRAY membership plus a caller-supplied external team id (FIELD_ID), as an + MPI communicator passes its context id. */ +static ucc_team_params_t make_array_id_params( + ucc_rank_t *arr, ucc_rank_t size, ucc_rank_t self_ep, uint64_t id) +{ + ucc_team_params_t p = make_array_params(arr, size, self_ep); + p.mask |= UCC_TEAM_PARAM_FIELD_ID; + p.id = id; + return p; +} + +/* Zero-init an identity and build it from @p, asserting success. Callers must + ucc_team_cache_identity_free the result. */ +static void build_identity( + const ucc_team_params_t &p, ucc_team_cache_identity_t &id) +{ + memset(&id, 0, sizeof(id)); + ASSERT_EQ(UCC_OK, ucc_team_cache_identity_build(&p, &id)); +} + +/* Assert two params materialize to an equal identity (same hash AND + exact-compare equal). Frees both identities. */ +static void expect_identities_equal( + const ucc_team_params_t &pa, const ucc_team_params_t &pb) +{ + ucc_team_cache_identity_t a, b; + build_identity(pa, a); + build_identity(pb, b); + EXPECT_EQ(a.hash, b.hash); + EXPECT_NE(0, ucc_team_cache_identity_equal(&a, &b)); + ucc_team_cache_identity_free(&a); + ucc_team_cache_identity_free(&b); +} + +class test_team_cache : public ucc::test {}; + +/* Identical membership + DIFFERENT external ids must NOT be full-equal (no + dormant reuse across id/tag domains), but must stay membership-equal so + coexistence/derived detection finds the live parent. Same id -> fully equal. */ +UCC_TEST_F(test_team_cache, external_id_isolates_dormant_reuse) +{ + ucc_rank_t arr[4] = {0, 1, 2, 3}; + ucc_team_params_t p3 = make_array_id_params(arr, 4, 1, 3); + ucc_team_params_t p3b = make_array_id_params(arr, 4, 1, 3); + ucc_team_params_t p5 = make_array_id_params(arr, 4, 1, 5); + + ucc_team_cache_identity_t a, b, c; + build_identity(p3, a); + build_identity(p3b, b); + build_identity(p5, c); + + /* Membership-only hash: all three share a hash bucket. */ + EXPECT_EQ(a.hash, b.hash); + EXPECT_EQ(a.hash, c.hash); + + /* Same members + same id -> full match; different id -> no full match. */ + EXPECT_NE(0, ucc_team_cache_identity_equal(&a, &b)); + EXPECT_EQ(0, ucc_team_cache_identity_equal(&a, &c)); + /* Membership matches regardless of id (derived/coexistence detection). */ + EXPECT_NE(0, ucc_team_cache_identity_equal_membership(&a, &c)); + + ucc_team_cache_identity_free(&a); + ucc_team_cache_identity_free(&b); + ucc_team_cache_identity_free(&c); +} + +/* Identity ignores ep_map style / closure pointers: two CB closures, CB vs + ARRAY, and FULL/STRIDED(0,1) vs ARRAY [0..size) all produce an equal identity. */ +UCC_TEST_F(test_team_cache, cross_style_cb_vs_array_equal) +{ + ucc_rank_t same[4] = {3, 5, 7, 9}; + expect_identities_equal( + make_array_params(same, 4, 1), make_array_params(same, 4, 1)); + + /* Two distinct CB closures materializing the same members. */ + cb_ctx c1, c2; + c1.members = {2, 4, 6, 8}; + c2.members = {2, 4, 6, 8}; + expect_identities_equal(make_cb_params(&c1, 0), make_cb_params(&c2, 0)); + + /* Cross-style: EP_MAP CB (coll/ucc) vs EP_MAP_ARRAY (scoll/ucc). */ + cb_ctx c; + c.members = {10, 20, 30}; + ucc_rank_t arr3[3] = {10, 20, 30}; + expect_identities_equal( + make_cb_params(&c, 2), make_array_params(arr3, 3, 2)); + + /* CB / ARRAY / FULL / STRIDED(0,1) over [0..5) must all be equal. */ + ucc_rank_t arr[5] = {0, 1, 2, 3, 4}; + cb_ctx c5; + c5.members = {0, 1, 2, 3, 4}; + ucc_team_params_t pcb = make_cb_params(&c5, 0); + ucc_team_params_t parr = make_array_params(arr, 5, 0); + ucc_team_params_t pstr = make_strided_params(0, 1, 5, 0); + + ucc_team_params_t pfull; + memset(&pfull, 0, sizeof(pfull)); + pfull.mask = UCC_TEAM_PARAM_FIELD_EP_MAP | UCC_TEAM_PARAM_FIELD_EP; + pfull.ep = 0; + pfull.ep_map.type = UCC_EP_MAP_FULL; + pfull.ep_map.ep_num = 5; + + ucc_team_cache_identity_t cb, a, b, full; + build_identity(pcb, cb); + build_identity(parr, a); + build_identity(pstr, b); + build_identity(pfull, full); + + EXPECT_NE(0, ucc_team_cache_identity_equal(&a, &cb)); + EXPECT_NE(0, ucc_team_cache_identity_equal(&a, &b)); + EXPECT_NE(0, ucc_team_cache_identity_equal(&a, &full)); + + ucc_team_cache_identity_free(&cb); + ucc_team_cache_identity_free(&a); + ucc_team_cache_identity_free(&b); + ucc_team_cache_identity_free(&full); +} + +/* Freeing/mutating the caller's closure and user array after build does not + change the identity (params are materialized, not retained). */ +UCC_TEST_F(test_team_cache, identity_owns_materialized_members) +{ + cb_ctx *c = new cb_ctx(); + c->members = {11, 13, 17, 19}; + + ucc_rank_t *arr = (ucc_rank_t *)malloc(4 * sizeof(ucc_rank_t)); + ASSERT_NE(nullptr, arr); + arr[0] = 11; + arr[1] = 13; + arr[2] = 17; + arr[3] = 19; + + ucc_team_params_t pcb = make_cb_params(c, 3); + ucc_team_params_t parr = make_array_params(arr, 4, 3); + + ucc_team_cache_identity_t from_cb, from_arr, ref; + build_identity(pcb, from_cb); + build_identity(parr, from_arr); + + ucc_rank_t refarr[4] = {11, 13, 17, 19}; + ucc_team_params_t pref = make_array_params(refarr, 4, 3); + build_identity(pref, ref); + + /* Destroy/mutate the caller-owned inputs. */ + delete c; + arr[0] = 999; + arr[2] = 42; + free(arr); + + EXPECT_EQ(ref.hash, from_cb.hash); + EXPECT_EQ(ref.hash, from_arr.hash); + EXPECT_NE(0, ucc_team_cache_identity_equal(&ref, &from_cb)); + EXPECT_NE(0, ucc_team_cache_identity_equal(&ref, &from_arr)); + + ucc_team_cache_identity_free(&from_cb); + ucc_team_cache_identity_free(&from_arr); + ucc_team_cache_identity_free(&ref); +} + +/* Differing membership -> not equal (size, self_ep, array contents, stride). */ +UCC_TEST_F(test_team_cache, differing_membership_not_equal) +{ + ucc_rank_t base[4] = {1, 2, 3, 4}; + ucc_rank_t diff_val[4] = {1, 2, 3, 5}; + ucc_rank_t diff_len[3] = {1, 2, 3}; + + ucc_team_params_t pbase = make_array_params(base, 4, 1); + ucc_team_params_t pval = make_array_params(diff_val, 4, 1); + ucc_team_params_t plen = make_array_params(diff_len, 3, 1); + ucc_team_params_t pep = make_array_params(base, 4, 2); + ucc_team_params_t pstr1 = make_strided_params(0, 1, 4, 0); + ucc_team_params_t pstr2 = make_strided_params(0, 2, 4, 0); + + ucc_team_cache_identity_t base_id, val_id, len_id, ep_id, s1, s2; + build_identity(pbase, base_id); + build_identity(pval, val_id); + build_identity(plen, len_id); + build_identity(pep, ep_id); + build_identity(pstr1, s1); + build_identity(pstr2, s2); + + EXPECT_EQ(0, ucc_team_cache_identity_equal(&base_id, &val_id)); + EXPECT_EQ(0, ucc_team_cache_identity_equal(&base_id, &len_id)); + EXPECT_EQ(0, ucc_team_cache_identity_equal(&base_id, &ep_id)); + EXPECT_EQ(0, ucc_team_cache_identity_equal(&s1, &s2)); + + ucc_team_cache_identity_free(&base_id); + ucc_team_cache_identity_free(&val_id); + ucc_team_cache_identity_free(&len_id); + ucc_team_cache_identity_free(&ep_id); + ucc_team_cache_identity_free(&s1); + ucc_team_cache_identity_free(&s2); +} + +/* Agreement vote: a UCC_OP_BAND allreduce over the members' fill buffers. */ +static void vote_band_reduce( + const std::vector> &in, uint64_t *out) +{ + for (int l = 0; l < UCC_TEAM_CACHE_VOTE_LANES; l++) { + out[l] = ~(uint64_t)0; + } + for (const auto &v : in) { + for (int l = 0; l < UCC_TEAM_CACHE_VOTE_LANES; l++) { + out[l] &= v[l]; + } + } +} + +/* All-hit agreement: EXACT_REUSE with a stable key agrees; rank-0's proposed + cookie is distributed via lane [9]. */ +UCC_TEST_F(test_team_cache, vote_agreement_and_cookie) +{ + uint64_t out[UCC_TEAM_CACHE_VOTE_LANES]; + + /* EXACT_REUSE with cookie=0 everywhere still agrees (ext_id pins the + instance). */ + { + std::vector> in( + 3, std::vector(UCC_TEAM_CACHE_VOTE_LANES)); + for (int r = 0; r < 3; r++) { + ucc_team_cache_vote_fill( + in[r].data(), + 1, + UCC_TEAM_CACHE_ACTION_EXACT_REUSE, + 0x55, + /*cookie=*/0, + /*parent_cookie=*/0, + /*is_rank0=*/(r == 0), + 0x1); + } + vote_band_reduce(in, out); + EXPECT_EQ( + UCC_TEAM_CACHE_ACTION_EXACT_REUSE, ucc_team_cache_vote_result(out)); + /* rank-0's proposed cookie must survive in lane [9]. */ + EXPECT_EQ((uint64_t)0x1, ucc_team_cache_vote_new_cookie(out)); + } + + /* MISS from all ranks (all prepared=0) -> result stays MISS. */ + { + std::vector> in( + 2, std::vector(UCC_TEAM_CACHE_VOTE_LANES)); + for (int r = 0; r < 2; r++) { + ucc_team_cache_vote_fill( + in[r].data(), + /*prepared=*/0, + UCC_TEAM_CACHE_ACTION_MISS, + 0, + 0, + 0, + /*is_rank0=*/(r == 0), + /*proposed_cookie=*/0xBEEF); + } + vote_band_reduce(in, out); + EXPECT_EQ(UCC_TEAM_CACHE_ACTION_MISS, ucc_team_cache_vote_result(out)); + EXPECT_EQ((uint64_t)0xBEEF, ucc_team_cache_vote_new_cookie(out)); + } +} + +/* A single miss rank (prepared=0) forces a global MISS even if every other rank + agrees, and rank-0's new-cookie proposal still survives for the fresh build. */ +UCC_TEST_F(test_team_cache, vote_one_miss_rank_forces_miss_keeps_cookie) +{ + std::vector> in( + 3, std::vector(UCC_TEAM_CACHE_VOTE_LANES)); + uint64_t out[UCC_TEAM_CACHE_VOTE_LANES]; + + ucc_team_cache_vote_fill( + in[0].data(), + 1, + UCC_TEAM_CACHE_ACTION_EXACT_REUSE, + 0x1234, + 0, + 0, + /*is_rank0=*/1, + /*proposed_cookie=*/0x900D); + ucc_team_cache_vote_fill( + in[1].data(), + 1, + UCC_TEAM_CACHE_ACTION_EXACT_REUSE, + 0x1234, + 0, + 0, + /*is_rank0=*/0, + 0); + ucc_team_cache_vote_fill( + in[2].data(), + /*prepared=*/0, + UCC_TEAM_CACHE_ACTION_MISS, + 0, + 0, + 0, + /*is_rank0=*/0, + 0); + vote_band_reduce(in, out); + EXPECT_EQ(UCC_TEAM_CACHE_ACTION_MISS, ucc_team_cache_vote_result(out)); + EXPECT_EQ((uint64_t)0x900D, ucc_team_cache_vote_new_cookie(out)); +} + +/* Ranks that agree on the action and key but disagree on the instance cookie + must not reuse: the cookie equality pair breaks, so the vote degrades to + MISS. This is what keeps a re-seated instance from being silently adopted. */ +UCC_TEST_F(test_team_cache, vote_reseat_different_cookie_misses) +{ + std::vector> in( + 2, std::vector(UCC_TEAM_CACHE_VOTE_LANES)); + uint64_t out[UCC_TEAM_CACHE_VOTE_LANES]; + + ucc_team_cache_vote_fill( + in[0].data(), + 1, + UCC_TEAM_CACHE_ACTION_EXACT_REUSE, + 0xABCD, + /*cookie=*/0x11, + /*parent_cookie=*/0, + /*is_rank0=*/1, + /*proposed_cookie=*/0x77); + ucc_team_cache_vote_fill( + in[1].data(), + 1, + UCC_TEAM_CACHE_ACTION_EXACT_REUSE, + 0xABCD, + /*cookie=*/0x22, /* re-seated instance */ + /*parent_cookie=*/0, + /*is_rank0=*/0, + 0); + vote_band_reduce(in, out); + EXPECT_EQ(UCC_TEAM_CACHE_ACTION_MISS, ucc_team_cache_vote_result(out)); + EXPECT_EQ((uint64_t)0x77, ucc_team_cache_vote_new_cookie(out)); + + /* A parent-cookie disagreement alone is equally disqualifying. */ + ucc_team_cache_vote_fill( + in[1].data(), + 1, + UCC_TEAM_CACHE_ACTION_EXACT_REUSE, + 0xABCD, + /*cookie=*/0x11, + /*parent_cookie=*/0x99, + /*is_rank0=*/0, + 0); + vote_band_reduce(in, out); + EXPECT_EQ(UCC_TEAM_CACHE_ACTION_MISS, ucc_team_cache_vote_result(out)); +} + +/* next_cookie is strictly monotonic and never returns 0, since 0 is the + "unstamped" sentinel that makes a vote accept any instance. */ +UCC_TEST_F(test_team_cache, next_cookie_monotonic_nonzero) +{ + ucc_team_cache_t c; + memset(&c, 0, sizeof(c)); + c.cache_gen = 0; + uint64_t a = ucc_team_cache_next_cookie(&c); + uint64_t b = ucc_team_cache_next_cookie(&c); + EXPECT_NE((uint64_t)0, a); + EXPECT_NE((uint64_t)0, b); + EXPECT_LT(a, b); + + /* The wrap is only reachable after 2^64 adoptions, so seed it: the next + increment rolls cache_gen to 0, which must be skipped. */ + c.cache_gen = UINT64_MAX; + EXPECT_NE((uint64_t)0, ucc_team_cache_next_cookie(&c)) + << "a wrapped cookie must skip the unstamped sentinel"; +} + +/* is_cacheable: true when only EP_MAP/EP/OOB/etc are set; false when any + optional behavioral field is set. FLAGS is ignored. */ +UCC_TEST_F(test_team_cache, is_cacheable_policy) +{ + ucc_rank_t arr[2] = {0, 1}; + ucc_team_params_t p = make_array_params(arr, 2, 0); + + EXPECT_NE(0, ucc_team_cache_is_cacheable(&p)); + + ucc_team_params_t pflags = p; + pflags.mask |= UCC_TEAM_PARAM_FIELD_FLAGS; + pflags.flags = 0x1; + EXPECT_NE(0, ucc_team_cache_is_cacheable(&pflags)); + + const uint64_t optional[] = { + UCC_TEAM_PARAM_FIELD_ORDERING, + UCC_TEAM_PARAM_FIELD_OUTSTANDING_COLLS, + UCC_TEAM_PARAM_FIELD_SYNC_TYPE, + UCC_TEAM_PARAM_FIELD_P2P_CONN, + UCC_TEAM_PARAM_FIELD_MEM_PARAMS, + }; + for (size_t i = 0; i < sizeof(optional) / sizeof(optional[0]); i++) { + ucc_team_params_t po = p; + po.mask |= optional[i]; + EXPECT_EQ(0, ucc_team_cache_is_cacheable(&po)) + << "optional field index " << i << " should block caching"; + } +} + +/* Build a lookup key from ARRAY membership + external id. */ +static void build_id_key( + ucc_rank_t *arr, ucc_rank_t size, ucc_rank_t self_ep, uint64_t id, + ucc_team_cache_identity_t &key) +{ + ucc_team_params_t p = make_array_id_params(arr, size, self_ep, id); + build_identity(p, key); +} + +/* RAII wrapper for ucc_team_cache_init/destroy; args match ucc_team_cache_init. + Implicitly usable as a ucc_team_cache_t*. */ +struct ScopedCache { + ucc_team_cache_t *cache = nullptr; + + ScopedCache( + uint32_t max_size, ucc_team_cache_eviction_policy_t evict, + int disable_linear_check) + { + EXPECT_EQ( + UCC_OK, + ucc_team_cache_init(&cache, max_size, evict, disable_linear_check)); + } + ~ScopedCache() + { + ucc_team_cache_destroy(cache); + } + + operator ucc_team_cache_t *() const + { + return cache; + } + ucc_team_cache_t *operator->() const + { + return cache; + } +}; + +static bool team_on_list(ucc_list_link_t *head, ucc_team_t *team) +{ + ucc_team_t *t; + ucc_list_for_each (t, head, cache_link) { + if (t == team) { + return true; + } + } + return false; +} + +/* Registry list-surgery matrix: add-live -> make-dormant -> make-live -> remove. + Asserts the team is on exactly one list (or none after remove) at each step. */ +UCC_TEST_F(test_team_cache, registry_add_dormant_live_remove) +{ + ScopedCache cache(16, UCC_TEAM_CACHE_EVICTION_FIFO, 0); + + ucc_team_t team; + memset(&team, 0, sizeof(team)); + ucc_list_head_init(&team.cache_link); + + EXPECT_FALSE(team_on_list(&cache->live, &team)); + EXPECT_FALSE(team_on_list(&cache->dormant, &team)); + + ucc_team_cache_registry_add_live(cache, &team); + EXPECT_TRUE(team_on_list(&cache->live, &team)); + EXPECT_FALSE(team_on_list(&cache->dormant, &team)); + EXPECT_FALSE(ucc_list_is_empty(&cache->live)); + + ucc_team_cache_registry_make_dormant(cache, &team); + EXPECT_FALSE(team_on_list(&cache->live, &team)); + EXPECT_TRUE(team_on_list(&cache->dormant, &team)); + EXPECT_TRUE(ucc_list_is_empty(&cache->live)); + EXPECT_FALSE(ucc_list_is_empty(&cache->dormant)); + + ucc_team_cache_registry_make_live(cache, &team); + EXPECT_TRUE(team_on_list(&cache->live, &team)); + EXPECT_FALSE(team_on_list(&cache->dormant, &team)); + + ucc_team_cache_registry_remove(&team); + EXPECT_FALSE(team_on_list(&cache->live, &team)); + EXPECT_FALSE(team_on_list(&cache->dormant, &team)); + EXPECT_TRUE(ucc_list_is_empty(&cache->live)); + EXPECT_TRUE(ucc_list_is_empty(&cache->dormant)); +} + +/* Cache API tests (lookup / insert / get / put) use bare stub teams with only + the cache-related fields initialized; they never call create_post. */ + +/* Allocate and minimally initialize a stub team. refcount starts at 0 to match + the production convention for a cached DORMANT team. */ +static ucc_team_t *alloc_stub_team(void) +{ + ucc_team_t *t = (ucc_team_t *)calloc(1, sizeof(*t)); + if (!t) { + return nullptr; + } + t->refcount = 0; + t->cache_state = UCC_TEAM_CACHE_STATE_NONE; + ucc_list_head_init(&t->cache_link); + memset(&t->cache_identity, 0, sizeof(t->cache_identity)); + return t; +} + +static void free_stub_team(ucc_team_t *t) +{ + ucc_team_cache_identity_free(&t->cache_identity); + free(t); +} + +/* Detach a stub team from the table + registry. Caller must hold cache->lock. */ +static void erase_stub(ucc_team_cache_t *cache, ucc_team_t *t) +{ + ucc_team_cache_table_erase(cache, t); + ucc_team_cache_registry_remove(t); +} + +/* Lock, erase, unlock, and free each team - the common stub teardown. */ +static void erase_and_free( + ucc_team_cache_t *cache, std::initializer_list teams) +{ + ucc_spin_lock(&cache->lock); + for (ucc_team_t *t : teams) { + erase_stub(cache, t); + } + ucc_spin_unlock(&cache->lock); + for (ucc_team_t *t : teams) { + free_stub_team(t); + } +} + +/* lookup does NOT return a LIVE team: insert -> get (DORMANT->LIVE) -> miss. */ +UCC_TEST_F(test_team_cache, lookup_does_not_return_live) +{ + ScopedCache cache(16, UCC_TEAM_CACHE_EVICTION_FIFO, 0); + + ucc_rank_t arr[3] = {2, 4, 6}; + ucc_team_params_t p = make_array_params(arr, 3, 0); + + ucc_team_t *team = alloc_stub_team(); + ASSERT_NE(nullptr, team); + ASSERT_EQ(UCC_OK, ucc_team_cache_identity_build(&p, &team->cache_identity)); + + ucc_team_cache_identity_t key; + build_identity(p, key); + + ucc_spin_lock(&cache->lock); + ASSERT_EQ(UCC_OK, ucc_team_cache_insert(cache, team)); + ucc_team_cache_get(team); + EXPECT_EQ(UCC_TEAM_CACHE_STATE_LIVE, team->cache_state); + ucc_team_t *found = ucc_team_cache_lookup(cache, &key); + ucc_spin_unlock(&cache->lock); + + EXPECT_EQ(nullptr, found); + EXPECT_EQ(0u, cache->stats.hits); + + ucc_team_cache_identity_free(&key); + /* Release team back to dormant so erase_and_free can clean it up. */ + ucc_spin_lock(&cache->lock); + ucc_team_cache_put(team); + ucc_team_cache_registry_make_dormant(cache, team); + ucc_spin_unlock(&cache->lock); + erase_and_free(cache, {team}); +} + +/* Full-cache skip: second insert into a full max_size==1 cache stays NONE. */ +UCC_TEST_F(test_team_cache, full_cache_skip) +{ + ScopedCache cache(1, UCC_TEAM_CACHE_EVICTION_NONE, 0); + + ucc_rank_t arr1[2] = {0, 1}; + ucc_rank_t arr2[2] = {2, 3}; + ucc_team_params_t p1 = make_array_params(arr1, 2, 0); + ucc_team_params_t p2 = make_array_params(arr2, 2, 0); + + ucc_team_t *t1 = alloc_stub_team(); + ucc_team_t *t2 = alloc_stub_team(); + ASSERT_NE(nullptr, t1); + ASSERT_NE(nullptr, t2); + + ASSERT_EQ(UCC_OK, ucc_team_cache_identity_build(&p1, &t1->cache_identity)); + ASSERT_EQ(UCC_OK, ucc_team_cache_identity_build(&p2, &t2->cache_identity)); + + ucc_spin_lock(&cache->lock); + EXPECT_EQ(UCC_OK, ucc_team_cache_insert(cache, t1)); + EXPECT_EQ(UCC_TEAM_CACHE_STATE_DORMANT, t1->cache_state); + EXPECT_EQ(1u, cache->size); + + EXPECT_EQ(UCC_OK, ucc_team_cache_insert(cache, t2)); + EXPECT_EQ(UCC_TEAM_CACHE_STATE_NONE, t2->cache_state); + EXPECT_EQ(1u, cache->size); + ucc_spin_unlock(&cache->lock); + + erase_and_free( + cache, {t1}); /* t1 was inserted (DORMANT); remove before free */ + free_stub_team(t2); /* t2 was not inserted (NONE); free directly */ +} + +/* A duplicate identity finds the bucket occupied and is left uncached. */ +UCC_TEST_F(test_team_cache, duplicate_identity_skips_insert) +{ + ScopedCache cache(16, UCC_TEAM_CACHE_EVICTION_FIFO, 0); + + ucc_rank_t arr[3] = {5, 6, 7}; + ucc_team_params_t p = make_array_params(arr, 3, 0); + + ucc_team_t *t1 = alloc_stub_team(); + ucc_team_t *t2 = alloc_stub_team(); + ASSERT_NE(nullptr, t1); + ASSERT_NE(nullptr, t2); + + ASSERT_EQ(UCC_OK, ucc_team_cache_identity_build(&p, &t1->cache_identity)); + ASSERT_EQ(UCC_OK, ucc_team_cache_identity_build(&p, &t2->cache_identity)); + + ucc_spin_lock(&cache->lock); + EXPECT_EQ(UCC_OK, ucc_team_cache_insert(cache, t1)); + EXPECT_EQ(UCC_OK, ucc_team_cache_insert(cache, t2)); + ucc_spin_unlock(&cache->lock); + + EXPECT_EQ(UCC_TEAM_CACHE_STATE_DORMANT, t1->cache_state); + EXPECT_EQ(UCC_TEAM_CACHE_STATE_NONE, t2->cache_state); + EXPECT_EQ(1u, cache->size); + EXPECT_EQ(1u, cache->stats.inserts); + + erase_and_free(cache, {t1}); + free_stub_team(t2); +} + +/* get/put refcount arithmetic and LIVE<->DORMANT transitions. */ +UCC_TEST_F(test_team_cache, get_put_refcount_and_state) +{ + ScopedCache cache(16, UCC_TEAM_CACHE_EVICTION_FIFO, 0); + + ucc_rank_t arr[2] = {0, 1}; + ucc_team_params_t p = make_array_params(arr, 2, 0); + + ucc_team_t *team = alloc_stub_team(); /* refcount 0 (DORMANT) */ + ASSERT_NE(nullptr, team); + ASSERT_EQ(UCC_OK, ucc_team_cache_identity_build(&p, &team->cache_identity)); + + ucc_team_cache_identity_t key; + build_identity(p, key); + + ucc_spin_lock(&cache->lock); + ASSERT_EQ(UCC_OK, ucc_team_cache_insert(cache, team)); + EXPECT_EQ(0, team->refcount); + EXPECT_EQ(UCC_TEAM_CACHE_STATE_DORMANT, team->cache_state); + + /* Adopt the dormant team: 0 -> 1, LIVE. */ + ucc_team_cache_get(team); + ucc_team_cache_registry_make_live(cache, team); + EXPECT_EQ(1, team->refcount); + EXPECT_EQ(UCC_TEAM_CACHE_STATE_LIVE, team->cache_state); + + /* Last user drops: 1 -> 0, DORMANT. */ + EXPECT_EQ(0, ucc_team_cache_put(team)); + EXPECT_EQ(UCC_TEAM_CACHE_STATE_DORMANT, team->cache_state); + ucc_team_cache_registry_make_dormant(cache, team); + + /* A team that went LIVE then back to DORMANT is look-up-able again. */ + EXPECT_EQ(team, ucc_team_cache_lookup(cache, &key)); + + /* Two live users need two puts to return to DORMANT. */ + ucc_team_cache_get(team); + ucc_team_cache_registry_make_live(cache, team); + ucc_team_cache_get(team); + EXPECT_EQ(2, team->refcount); + EXPECT_EQ(UCC_TEAM_CACHE_STATE_LIVE, team->cache_state); + + EXPECT_EQ(1, ucc_team_cache_put(team)); + EXPECT_EQ(UCC_TEAM_CACHE_STATE_LIVE, team->cache_state); + EXPECT_EQ(0, ucc_team_cache_put(team)); + EXPECT_EQ(UCC_TEAM_CACHE_STATE_DORMANT, team->cache_state); + ucc_team_cache_registry_make_dormant(cache, team); + ucc_spin_unlock(&cache->lock); + + ucc_team_cache_identity_free(&key); + erase_and_free(cache, {team}); +} + +/* RESERVED state (the agreement-vote pin): a DORMANT candidate moved to RESERVED + is off the dormant/live lists (lookup can't return it) but stays in the bucket + with refcount unchanged, so a vote-FAIL rolls it back to DORMANT and a vote-PASS + promotes it to LIVE via get (0 -> 1). */ +UCC_TEST_F(test_team_cache, reserved_state_pin_and_rollback) +{ + ScopedCache cache(16, UCC_TEAM_CACHE_EVICTION_FIFO, 0); + ucc_rank_t arr[3] = {10, 20, 30}; + ucc_team_params_t p = make_array_params(arr, 3, 0); + + ucc_team_t *t = alloc_stub_team(); + ASSERT_NE(nullptr, t); + ASSERT_EQ(UCC_OK, ucc_team_cache_identity_build(&p, &t->cache_identity)); + + ucc_spin_lock(&cache->lock); + ASSERT_EQ(UCC_OK, ucc_team_cache_insert(cache, t)); + ASSERT_EQ(UCC_TEAM_CACHE_STATE_DORMANT, t->cache_state); + + ucc_team_cache_identity_t key; + build_identity(p, key); + + ASSERT_EQ(t, ucc_team_cache_lookup(cache, &key)); + + /* Pin for an in-flight vote: DORMANT -> RESERVED, refcount untouched. */ + ucc_team_cache_registry_make_reserved(cache, t); + t->cache_state = UCC_TEAM_CACHE_STATE_RESERVED; + EXPECT_EQ(0, t->refcount); + /* lookup is DORMANT-only: a RESERVED team must not be returned. */ + EXPECT_EQ(nullptr, ucc_team_cache_lookup(cache, &key)); + + /* Vote FAIL: roll back RESERVED -> DORMANT, re-adoptable. */ + ucc_team_cache_registry_make_dormant(cache, t); + t->cache_state = UCC_TEAM_CACHE_STATE_DORMANT; + EXPECT_EQ(t, ucc_team_cache_lookup(cache, &key)); + + /* Vote PASS: RESERVED -> LIVE via get (0 -> 1). */ + ucc_team_cache_registry_make_reserved(cache, t); + t->cache_state = UCC_TEAM_CACHE_STATE_RESERVED; + ucc_team_cache_get(t); + ucc_team_cache_registry_make_live(cache, t); + EXPECT_EQ(1, t->refcount); + EXPECT_EQ(UCC_TEAM_CACHE_STATE_LIVE, t->cache_state); + EXPECT_EQ( + nullptr, ucc_team_cache_lookup(cache, &key)); /* LIVE not returned */ + ucc_spin_unlock(&cache->lock); + + ucc_team_cache_identity_free(&key); + erase_and_free(cache, {t}); +} + +/* ucc_team_destroy must never tear down a handle the cache owns: a dormant or + reserved team, or a terminal handle from a failed adoption. */ +UCC_TEST_F(test_team_cache, destroy_rejects_cache_owned_handle) +{ + ucc_context_t *ctxs[1] = {nullptr}; /* never dereferenced on these paths */ + ucc_team_t *t = alloc_stub_team(); + ASSERT_NE(nullptr, t); + + t->contexts = ctxs; + t->num_contexts = 1; + t->state = UCC_TEAM_ACTIVE; + t->cache_state = UCC_TEAM_CACHE_STATE_DORMANT; + EXPECT_EQ(UCC_ERR_INVALID_PARAM, ucc_team_destroy(t)); + + t->cache_state = UCC_TEAM_CACHE_STATE_RESERVED; + EXPECT_EQ(UCC_ERR_INVALID_PARAM, ucc_team_destroy(t)); + + /* Failed adoption: terminal state on a team the cache reclaimed */ + t->state = UCC_TEAM_CREATE_FAILED; + t->cache_state = UCC_TEAM_CACHE_STATE_DORMANT; + EXPECT_EQ(UCC_ERR_INVALID_PARAM, ucc_team_destroy(t)); + + /* A terminal handle is not re-testable either */ + EXPECT_EQ(UCC_ERR_INVALID_PARAM, ucc_team_create_test(t)); + + /* The stub was never touched, so it is still ours to free */ + EXPECT_EQ(UCC_TEAM_CREATE_FAILED, t->state); + free_stub_team(t); +} + +/* Build a stub team with @arr membership and insert it into @cache as DORMANT. */ +static ucc_team_t *insert_stub_dormant( + ucc_team_cache_t *cache, ucc_rank_t *arr, ucc_rank_t n, ucc_rank_t self_ep) +{ + ucc_team_params_t p = make_array_params(arr, n, self_ep); + + ucc_team_t *t = alloc_stub_team(); + if (!t) { + return nullptr; + } + if (UCC_OK != ucc_team_cache_identity_build(&p, &t->cache_identity)) { + free_stub_team(t); + return nullptr; + } + + ucc_spin_lock(&cache->lock); + ucc_status_t st = ucc_team_cache_insert(cache, t); + ucc_spin_unlock(&cache->lock); + + if (st != UCC_OK || t->cache_state != UCC_TEAM_CACHE_STATE_DORMANT) { + free_stub_team(t); + return nullptr; + } + return t; +} + +/* Insert three distinct-membership DORMANT stub teams (list order A,B,C). */ +static void insert_three_dormant( + ucc_team_cache_t *cache, ucc_team_t **tA, ucc_team_t **tB, ucc_team_t **tC) +{ + ucc_rank_t mA[3] = {10, 20, 30}; + ucc_rank_t mB[3] = {11, 21, 31}; + ucc_rank_t mC[3] = {12, 22, 32}; + *tA = insert_stub_dormant(cache, mA, 3, 0); + *tB = insert_stub_dormant(cache, mB, 3, 0); + *tC = insert_stub_dormant(cache, mC, 3, 0); + ASSERT_NE(nullptr, *tA); + ASSERT_NE(nullptr, *tB); + ASSERT_NE(nullptr, *tC); +} + +/* Victim selection: FIFO returns the insertion head (oldest). */ +UCC_TEST_F(test_team_cache, evict_victim_selection) +{ + SCOPED_TRACE("fifo"); + ScopedCache cache(8, UCC_TEAM_CACHE_EVICTION_FIFO, 0); + + ucc_team_t *tA, *tB, *tC; + insert_three_dormant(cache, &tA, &tB, &tC); + ASSERT_EQ(3u, cache->size); + + tA->seq_num = 10; + tB->seq_num = 3; + tC->seq_num = 20; + + ucc_spin_lock(&cache->lock); + ucc_team_t *victim = ucc_team_cache_pick_victim(cache); + ucc_spin_unlock(&cache->lock); + + /* FIFO: oldest inserted (tA) regardless of seq_num. */ + EXPECT_EQ(tA, victim); + + erase_and_free(cache, {tA, tB, tC}); +} + +/* pick_victim skips LIVE teams: adopt both -> dormant empty -> NULL; + release both -> a victim is returned. */ +UCC_TEST_F(test_team_cache, evict_skips_live_returns_no_resource) +{ + ScopedCache cache(8, UCC_TEAM_CACHE_EVICTION_FIFO, 0); + + ucc_rank_t mA[2] = {100, 200}; + ucc_rank_t mB[2] = {101, 201}; + + ucc_team_t *tA = insert_stub_dormant(cache, mA, 2, 0); + ucc_team_t *tB = insert_stub_dormant(cache, mB, 2, 0); + ASSERT_NE(nullptr, tA); + ASSERT_NE(nullptr, tB); + + ucc_spin_lock(&cache->lock); + ucc_team_cache_get(tA); + ucc_team_cache_registry_make_live(cache, tA); + ucc_team_cache_get(tB); + ucc_team_cache_registry_make_live(cache, tB); + + EXPECT_EQ(nullptr, ucc_team_cache_pick_victim(cache)); + + ucc_team_cache_put(tA); + ucc_team_cache_registry_make_dormant(cache, tA); + ucc_team_cache_put(tB); + ucc_team_cache_registry_make_dormant(cache, tB); + + EXPECT_NE(nullptr, ucc_team_cache_pick_victim(cache)); + ucc_spin_unlock(&cache->lock); + + erase_and_free(cache, {tA, tB}); +} + +/* Linear-check knob: disable_linear_check controls whether lookup runs the exact + rank-array compare after a hash match. Collisions are injected by overwriting + the lookup key's hash. */ + +/* disable=0 (safe): the exact compare rejects a collision as MISS. + disable=1 (trust-hash): the compare is skipped and the hash match returns. */ +UCC_TEST_F(test_team_cache, linear_check_on_rejects_collision) +{ + auto run_collision = [](int disable_linear_check, bool expect_hit) { + SCOPED_TRACE( + disable_linear_check ? "disable_linear_check=1 (trust-hash)" + : "disable_linear_check=0 (safe)"); + ScopedCache cache( + 16, UCC_TEAM_CACHE_EVICTION_FIFO, disable_linear_check); + + ucc_rank_t arrA[3] = {1, 2, 3}; + ucc_team_t *teamA = alloc_stub_team(); + ASSERT_NE(nullptr, teamA); + ucc_team_params_t pA = make_array_params(arrA, 3, 0); + ASSERT_EQ( + UCC_OK, ucc_team_cache_identity_build(&pA, &teamA->cache_identity)); + + ucc_spin_lock(&cache->lock); + ASSERT_EQ(UCC_OK, ucc_team_cache_insert(cache, teamA)); + ucc_spin_unlock(&cache->lock); + EXPECT_EQ(UCC_TEAM_CACHE_STATE_DORMANT, teamA->cache_state); + + /* Lookup key {7,8,9}; hash overridden to teamA's to force a collision. */ + ucc_rank_t arrB[3] = {7, 8, 9}; + ucc_team_params_t pB = make_array_params(arrB, 3, 0); + ucc_team_cache_identity_t keyB; + build_identity(pB, keyB); + + ASSERT_NE(teamA->cache_identity.hash, keyB.hash); + keyB.hash = teamA->cache_identity.hash; + + ucc_spin_lock(&cache->lock); + ucc_team_t *found = ucc_team_cache_lookup(cache, &keyB); + ucc_spin_unlock(&cache->lock); + + if (expect_hit) { + EXPECT_EQ(teamA, found); + EXPECT_EQ(1u, cache->stats.hits); + EXPECT_EQ(0u, cache->stats.misses); + } else { + EXPECT_EQ(nullptr, found); + EXPECT_EQ(1u, cache->stats.misses); + EXPECT_EQ(0u, cache->stats.hits); + } + + ucc_team_cache_identity_free(&keyB); + erase_and_free(cache, {teamA}); + }; + + run_collision(/*disable_linear_check=*/0, /*expect_hit=*/false); + run_collision(/*disable_linear_check=*/1, /*expect_hit=*/true); +} + +/* Trust-hash mode skips the membership compare but NOT the ext_id compare, so a + dormant team under one ext_id is not re-adopted for a different ext_id. */ +UCC_TEST_F(test_team_cache, linear_check_off_still_honors_ext_id) +{ + ScopedCache cache(16, UCC_TEAM_CACHE_EVICTION_FIFO, 1); + + ucc_rank_t arr[3] = {1, 2, 3}; + ucc_team_t *teamA = alloc_stub_team(); + ASSERT_NE(nullptr, teamA); + ucc_team_params_t pA = make_array_id_params(arr, 3, 0, 7); + ASSERT_EQ( + UCC_OK, ucc_team_cache_identity_build(&pA, &teamA->cache_identity)); + + ucc_spin_lock(&cache->lock); + ASSERT_EQ(UCC_OK, ucc_team_cache_insert(cache, teamA)); + ucc_spin_unlock(&cache->lock); + EXPECT_EQ(UCC_TEAM_CACHE_STATE_DORMANT, teamA->cache_state); + + /* Same membership, DIFFERENT external id -> same hash, different ext_id. */ + ucc_team_cache_identity_t keyB; + build_id_key(arr, 3, 0, 8, keyB); + + ASSERT_EQ(teamA->cache_identity.hash, keyB.hash); + ASSERT_NE(teamA->cache_identity.ext_id, keyB.ext_id); + + ucc_spin_lock(&cache->lock); + ucc_team_t *found = ucc_team_cache_lookup(cache, &keyB); + ucc_spin_unlock(&cache->lock); + + EXPECT_EQ(nullptr, found) + << "trust-hash mode must still reject a differing ext_id"; + EXPECT_EQ(1u, cache->stats.misses); + EXPECT_EQ(0u, cache->stats.hits); + + ucc_team_cache_identity_free(&keyB); + erase_and_free(cache, {teamA}); +} + +/* All counters start zeroed; insert/hit/miss/hit-after-put accumulate them as + expected. Also folds the insert->DORMANT + size==1 postcondition. */ +UCC_TEST_F(test_team_cache, stats_accumulate_correctly) +{ + ScopedCache cache(16, UCC_TEAM_CACHE_EVICTION_FIFO, 0); + + EXPECT_EQ(0u, cache->stats.lookups); + EXPECT_EQ(0u, cache->stats.hits); + EXPECT_EQ(0u, cache->stats.misses); + EXPECT_EQ(0u, cache->stats.inserts); + EXPECT_EQ(0u, cache->stats.evictions); + + ucc_rank_t arr1[2] = {0, 1}; + ucc_rank_t arr2[3] = {0, 1, 2}; + ucc_team_params_t p1 = make_array_params(arr1, 2, 0); + ucc_team_params_t p2 = make_array_params(arr2, 3, 0); + + ucc_team_t *t1 = alloc_stub_team(); + ucc_team_t *t2 = alloc_stub_team(); + ASSERT_NE(nullptr, t1); + ASSERT_NE(nullptr, t2); + + ASSERT_EQ(UCC_OK, ucc_team_cache_identity_build(&p1, &t1->cache_identity)); + ASSERT_EQ(UCC_OK, ucc_team_cache_identity_build(&p2, &t2->cache_identity)); + + ucc_team_cache_identity_t k1, k2; + build_identity(p1, k1); + build_identity(p2, k2); + + ucc_spin_lock(&cache->lock); + + /* Insert t1: DORMANT, size 1, one insert, no lookup counted. */ + EXPECT_EQ(UCC_OK, ucc_team_cache_insert(cache, t1)); + EXPECT_EQ(UCC_TEAM_CACHE_STATE_DORMANT, t1->cache_state); + EXPECT_EQ(1u, cache->size); + EXPECT_EQ(1u, cache->stats.inserts); + EXPECT_EQ(0u, cache->stats.lookups); + + /* Lookup t1 (hit). */ + EXPECT_EQ(t1, ucc_team_cache_lookup(cache, &k1)); + EXPECT_EQ(1u, cache->stats.lookups); + EXPECT_EQ(1u, cache->stats.hits); + EXPECT_EQ(0u, cache->stats.misses); + + /* Lookup t2 identity (miss, not inserted). */ + EXPECT_EQ(nullptr, ucc_team_cache_lookup(cache, &k2)); + EXPECT_EQ(2u, cache->stats.lookups); + EXPECT_EQ(1u, cache->stats.hits); + EXPECT_EQ(1u, cache->stats.misses); + + /* Insert t2. */ + EXPECT_EQ(UCC_OK, ucc_team_cache_insert(cache, t2)); + EXPECT_EQ(2u, cache->stats.inserts); + + /* Lookup t1 again (another hit). */ + EXPECT_EQ(t1, ucc_team_cache_lookup(cache, &k1)); + EXPECT_EQ(3u, cache->stats.lookups); + EXPECT_EQ(2u, cache->stats.hits); + EXPECT_EQ(1u, cache->stats.misses); + + ucc_spin_unlock(&cache->lock); + + ucc_team_cache_identity_free(&k1); + ucc_team_cache_identity_free(&k2); + erase_and_free(cache, {t1, t2}); +} + +/* dump_stats executes without fault (indirectly validates the format string), + including the zero-lookups divide-by-zero guard and the NULL no-op. */ +UCC_TEST_F(test_team_cache, dump_stats_no_crash) +{ + ScopedCache cache(16, UCC_TEAM_CACHE_EVICTION_FIFO, 0); + + ucc_spin_lock(&cache->lock); + cache->stats.lookups = 100; + cache->stats.hits = 80; + cache->stats.misses = 20; + cache->stats.inserts = 50; + cache->stats.evictions = 10; + ucc_spin_unlock(&cache->lock); + ucc_team_cache_dump_stats(cache); + + ucc_spin_lock(&cache->lock); + cache->stats.lookups = 0; + ucc_spin_unlock(&cache->lock); + ucc_team_cache_dump_stats(cache); + + ucc_team_cache_dump_stats(nullptr); +} + +/* Cache-concurrency stress: overlapping create_post on the same context is + invalid; cache->lock guards concurrent DESTROY/LOOKUP/CREATE on different + contexts. These tests drive the locked cache API directly from std::threads. */ +static bool cache_concurrency_runnable(void) +{ + unsigned hw = std::thread::hardware_concurrency(); + return (hw == 0) || (hw >= 2); +} + +/* Build @n stub teams with pairwise-distinct membership and refcount 0. */ +static void build_distinct_stub_teams( + std::vector &teams, + std::vector> &members, int n, int size_mod) +{ + for (int i = 0; i < n; i++) { + int sz = 2 + (i % size_mod); + members[i].resize(sz); + for (int j = 0; j < sz; j++) { + members[i][j] = (ucc_rank_t)j; + } + ucc_team_params_t p = make_array_params( + members[i].data(), (ucc_rank_t)sz, 0); + + teams[i] = alloc_stub_team(); + ASSERT_NE(nullptr, teams[i]); + teams[i]->refcount = 0; + ASSERT_EQ( + UCC_OK, + ucc_team_cache_identity_build(&p, &teams[i]->cache_identity)); + } +} + +/* Erase each resident stub under the lock and free it - concurrency teardown. */ +static void drain_stub_teams( + ucc_team_cache_t *cache, std::vector &teams) +{ + for (auto *t : teams) { + ucc_spin_lock(&cache->lock); + erase_stub(cache, t); + ucc_spin_unlock(&cache->lock); + free_stub_team(t); + } +} + +/* Contended DESTROY + LOOKUP path. A pool of DORMANT stub teams; N threads + race, each iteration under cache->lock: lookup, adopt on a DORMANT hit, then + release. */ +UCC_TEST_F(test_team_cache, concurrent_lookup_adopt_release_stress) +{ + if (!cache_concurrency_runnable()) { + GTEST_SKIP() << "host lacks >= 2 concurrent threads for cache stress"; + } + + const char *iters_env = std::getenv("UCC_GTEST_CACHE_STRESS_ITERS"); + const char *threads_env = std::getenv("UCC_GTEST_CACHE_STRESS_THREADS"); + const int n_iters = iters_env ? std::atoi(iters_env) : 500; + const int n_threads = threads_env ? std::atoi(threads_env) : 8; + const int n_teams = 16; + + ScopedCache cache(64, UCC_TEAM_CACHE_EVICTION_FIFO, 0); + + std::vector teams(n_teams); + std::vector keys(n_teams); + std::vector> members(n_teams); + std::vector> owned(n_teams); + + build_distinct_stub_teams(teams, members, n_teams, /*size_mod=*/n_teams); + for (int i = 0; i < n_teams; i++) { + ucc_spin_lock(&cache->lock); + ASSERT_EQ(UCC_OK, ucc_team_cache_insert(cache, teams[i])); + ucc_spin_unlock(&cache->lock); + ASSERT_EQ(UCC_TEAM_CACHE_STATE_DORMANT, teams[i]->cache_state); + + ucc_team_params_t p = make_array_params( + members[i].data(), (ucc_rank_t)members[i].size(), 0); + build_identity(p, keys[i]); + owned[i].store(false); + } + + const uint32_t size_before = cache->size; + + std::atomic double_adopt{false}; + std::atomic bad_refcount{false}; + + auto worker = [&](int seed) { + std::mt19937 rng((unsigned)(seed * 2654435761u + 1)); + for (int it = 0; it < n_iters; it++) { + int idx = rng() % n_teams; + + /* Adopt phase: lookup -> get -> make_live, all under the lock. */ + ucc_spin_lock(&cache->lock); + ucc_team_t *t = ucc_team_cache_lookup(cache, &keys[idx]); + if (t != nullptr) { + if (t->cache_state != UCC_TEAM_CACHE_STATE_DORMANT || + t->refcount != 0) { + bad_refcount.store(true); + } + ucc_team_cache_get(t); + ucc_team_cache_registry_make_live(cache, t); + } + ucc_spin_unlock(&cache->lock); + + if (t == nullptr) { + continue; /* another thread holds it live: legal miss */ + } + + /* Only one thread may hold this team live at a time. */ + if (owned[idx].exchange(true)) { + double_adopt.store(true); + } + std::this_thread::yield(); + owned[idx].store(false); + + /* Release phase: put -> make_dormant on last drop. */ + ucc_spin_lock(&cache->lock); + int rc = ucc_team_cache_put(t); + if (rc < 0) { + bad_refcount.store(true); + } + if (rc == 0) { + ucc_team_cache_registry_make_dormant(cache, t); + } + ucc_spin_unlock(&cache->lock); + } + }; + + std::vector pool; + for (int i = 0; i < n_threads; i++) { + pool.emplace_back(worker, i); + } + for (auto &th : pool) { + th.join(); + } + + EXPECT_FALSE(double_adopt.load()) + << "a LIVE team was adopted by two threads concurrently"; + EXPECT_FALSE(bad_refcount.load()) + << "cache refcount/state invariant violated under concurrency"; + + EXPECT_EQ(size_before, cache->size); + for (int i = 0; i < n_teams; i++) { + EXPECT_EQ(UCC_TEAM_CACHE_STATE_DORMANT, teams[i]->cache_state) + << "team " << i << " must settle back to DORMANT"; + EXPECT_EQ(0, teams[i]->refcount) + << "team " << i << " refcount must settle to 0"; + } + + for (int i = 0; i < n_teams; i++) { + ucc_team_cache_identity_free(&keys[i]); + } + drain_stub_teams(cache, teams); +} + +/* Integration tests: full create->use->destroy->recreate cycle through + UccJob/UccTeam (real create_post / create_test / destroy). White-box access + to ucc_team_t.cache_state and ucc_context_t.team_cache asserts dormant-reuse + invariants. Caching is enabled per-test via the UccJob env-var mechanism. */ +class test_team_cache_integration : public ucc::test {}; + +/* Return the underlying ucc_team_t* from a per-process team handle. */ +static ucc_team_t *team_ptr(UccTeam_h &team, int proc_idx = 0) +{ + return (ucc_team_t *)team->procs[proc_idx].team; +} + +/* Return the ucc_context_t* for the given process in a team. */ +static ucc_context_t *ctx_ptr(UccTeam_h &team, int proc_idx = 0) +{ + return (ucc_context_t *)team->procs[proc_idx].p.get()->ctx_h; +} + +/* True if the team that lived at @handle is on @ctx's dormant list. @handle is + compared by value only: when a destroy does NOT admit the team to the cache + the object is freed, so reading handle->cache_state would be a use-after-free + in exactly the case a dormancy assertion exists to catch. */ +static bool is_dormant(ucc_context_t *ctx, const ucc_team_t *handle) +{ + ucc_team_t *dt; + + ucc_list_for_each (dt, &ctx->team_cache->dormant, cache_link) { + if (dt == handle) { + return true; + } + } + return false; +} + +/* Run a single barrier collective on a team and assert it completes. */ +static void run_barrier(UccTeam_h &team) +{ + ucc_coll_args_t coll; + coll.mask = 0; + coll.coll_type = UCC_COLL_TYPE_BARRIER; + UccReq req(team, &coll); + req.start(); + ASSERT_EQ(UCC_OK, req.wait()); +} + +/* create->barrier->destroy->recreate-identical must re-adopt the SAME + ucc_team_t (pointer + team-id preserved) and record a hit on the second + create, then drive a second lifetime cycle. */ +UCC_TEST_F(test_team_cache_integration, dormant_reuse) +{ + UccJob job( + 4, + UccJob::UCC_JOB_CTX_GLOBAL, + {ucc_env_var_t("UCC_TEAM_CACHE_ENABLE", "y")}); + + UccTeam_h t1 = job.create_team(2, /*use_team_ep_map=*/true); + + ucc_team_t *tp0_before = team_ptr(t1, 0); + ucc_team_t *tp1_before = team_ptr(t1, 1); + uint16_t id0_before = tp0_before->id; + + EXPECT_EQ(UCC_TEAM_CACHE_STATE_LIVE, tp0_before->cache_state); + + /* Both contexts are captured before the reset below drops the team. */ + ucc_context_t *ctx0 = ctx_ptr(t1, 0); + ucc_context_t *ctx1 = ctx_ptr(t1, 1); + ASSERT_NE(nullptr, ctx0->team_cache); + + /* First create: a miss, no hit, one insert. */ + EXPECT_GE(ctx0->team_cache->stats.lookups, 1u); + EXPECT_EQ(0u, ctx0->team_cache->stats.hits); + EXPECT_EQ(1u, ctx0->team_cache->stats.inserts); + + run_barrier(t1); + + t1.reset(); + EXPECT_TRUE(is_dormant(ctx0, tp0_before)); + EXPECT_TRUE(is_dormant(ctx1, tp1_before)); + + /* Second team, IDENTICAL membership -> re-adopt the dormant team. */ + UccTeam_h t2 = job.create_team(2, /*use_team_ep_map=*/true); + + ucc_team_t *tp0_after = team_ptr(t2, 0); + ucc_team_t *tp1_after = team_ptr(t2, 1); + + EXPECT_EQ(tp0_before, tp0_after) << "team not reused"; + EXPECT_EQ(tp1_before, tp1_after) << "team not reused"; + EXPECT_EQ(UCC_TEAM_CACHE_STATE_LIVE, tp0_after->cache_state); + EXPECT_EQ(UCC_TEAM_CACHE_STATE_LIVE, tp1_after->cache_state); + EXPECT_EQ(id0_before, tp0_after->id) + << "re-adopted team must retain its original team ID"; + + /* Second create: one more lookup, exactly one hit, no extra insert. */ + EXPECT_GE(ctx0->team_cache->stats.lookups, 2u); + EXPECT_EQ(1u, ctx0->team_cache->stats.hits); + EXPECT_EQ(1u, ctx0->team_cache->stats.inserts); + + run_barrier(t2); + + t2.reset(); + EXPECT_EQ(UCC_TEAM_CACHE_STATE_DORMANT, tp0_after->cache_state); +} + +/* With UCC_TEAM_CACHE_ENABLE=n, create->destroy->recreate produces DISTINCT + ucc_team_t pointers; a team without EP_MAP is never cacheable. */ +UCC_TEST_F(test_team_cache_integration, knob_off) +{ + UccJob job( + 4, + UccJob::UCC_JOB_CTX_GLOBAL, + {ucc_env_var_t("UCC_TEAM_CACHE_ENABLE", "n")}); + + UccTeam_h t1 = job.create_team(2, /*use_team_ep_map=*/true); + ucc_team_t *tp_before = team_ptr(t1, 0); + + ucc_context_t *ctx0 = ctx_ptr(t1, 0); + EXPECT_EQ(nullptr, ctx0->team_cache) + << "team_cache must be NULL when caching is disabled"; + + run_barrier(t1); + + /* Create t2 while t1 is still live so pointer-distinctness is meaningful. */ + UccTeam_h t2 = job.create_team(2, /*use_team_ep_map=*/true); + ucc_team_t *tp_after = team_ptr(t2, 0); + + EXPECT_NE(tp_before, tp_after); + EXPECT_EQ(UCC_TEAM_CACHE_STATE_NONE, tp_after->cache_state); + + run_barrier(t2); + t1.reset(); + + /* A team created WITHOUT EP_MAP is never cacheable (identity_build requires + FIELD_EP_MAP): cache_state==NONE and distinct pointers. */ + UccTeam_h n1 = job.create_team(2, /*use_team_ep_map=*/false); + ucc_team_t *np1 = team_ptr(n1, 0); + EXPECT_EQ(UCC_TEAM_CACHE_STATE_NONE, np1->cache_state); + run_barrier(n1); + + UccTeam_h n2 = job.create_team(2, /*use_team_ep_map=*/false); + ucc_team_t *np2 = team_ptr(n2, 0); + EXPECT_NE(np1, np2); + EXPECT_EQ(UCC_TEAM_CACHE_STATE_NONE, np2->cache_state); + run_barrier(n2); + n1.reset(); +} + +/* Real UccJob, 2-slot FIFO cache. Create/destroy T1{2p}, T2{3p} (both DORMANT, + cache full), then create T3{4p}: admission evicts T1 and drains its destroy + synchronously; T3 lands in the freed slot. Asserts evictions >= 1 and + size <= max_size. */ +UCC_TEST_F(test_team_cache_integration, evict_id_release_on_evict) +{ + UccJob job( + 4, + UccJob::UCC_JOB_CTX_GLOBAL, + {ucc_env_var_t("UCC_TEAM_CACHE_ENABLE", "y"), + ucc_env_var_t("UCC_TEAM_CACHE_MAX_SIZE", "2"), + ucc_env_var_t("UCC_TEAM_CACHE_EVICTION", "fifo"), + ucc_env_var_t("UCC_TEAM_IDS_POOL_SIZE", "1")}); + + UccTeam_h t1 = job.create_team(2, /*use_team_ep_map=*/true); + ASSERT_NE(nullptr, t1); + EXPECT_EQ(UCC_TEAM_CACHE_STATE_LIVE, team_ptr(t1)->cache_state); + run_barrier(t1); + t1.reset(); /* T1 -> DORMANT (slot 1) */ + + UccTeam_h t2 = job.create_team(3, /*use_team_ep_map=*/true); + ASSERT_NE(nullptr, t2); + EXPECT_EQ(UCC_TEAM_CACHE_STATE_LIVE, team_ptr(t2)->cache_state); + run_barrier(t2); + + ucc_context_t *ctx0 = ctx_ptr(t2, 0); + ucc_team_cache_t *cache = ctx0->team_cache; + ASSERT_NE(nullptr, cache); + uint64_t evictions_before = cache->stats.evictions; + + t2.reset(); /* T2 -> DORMANT (slot 2, cache full) */ + + UccTeam_h t3 = job.create_team(4, /*use_team_ep_map=*/true); + ASSERT_NE(nullptr, t3); + run_barrier(t3); + + EXPECT_GT(cache->stats.evictions, evictions_before) + << "inserting T3 into a full cache must evict"; + EXPECT_LE(cache->size, cache->max_size); + EXPECT_TRUE(ucc_list_is_empty(&cache->pending_destroy)) + << "pending destroys must complete synchronously in gtest context"; +} + +/* Pool-pressure headroom: six DISTINCT-membership teams churned through a 2-slot + cache + size-1 ID pool. */ +UCC_TEST_F(test_team_cache_integration, id_pool_headroom_with_dormant_teams) +{ + UccJob job( + 16, + UccJob::UCC_JOB_CTX_GLOBAL, + {ucc_env_var_t("UCC_TEAM_CACHE_ENABLE", "y"), + ucc_env_var_t("UCC_TEAM_CACHE_MAX_SIZE", "2"), + ucc_env_var_t("UCC_TEAM_CACHE_EVICTION", "fifo"), + ucc_env_var_t("UCC_TEAM_IDS_POOL_SIZE", "1")}); + + ucc_context_t *ctx0 = nullptr; + ucc_team_cache_t *cache = nullptr; + + const int kTeams = 6; + for (int i = 0; i < kTeams; i++) { + int sz = 2 + i; /* sizes 2..7, all distinct, all fit in a 16-proc job */ + UccTeam_h t = job.create_team(sz, /*use_team_ep_map=*/true); + ASSERT_NE(nullptr, t) + << "create_team must succeed (no UCC_ERR_NO_RESOURCE) at team " + << i; + if (i == 0) { + ctx0 = ctx_ptr(t, 0); + cache = ctx0->team_cache; + } + run_barrier(t); + t.reset(); + } + + ASSERT_NE(nullptr, cache); + EXPECT_LE(cache->size, cache->max_size); + EXPECT_GT(cache->stats.evictions, 0u) + << "eviction must fire across " << kTeams << " distinct teams"; +} + +/* With UCC_TEAM_CACHE_DUMP_STATS=y the knob is wired onto the cache struct and + the destroy path dumps stats; the knob-off leg reads back zero. */ +UCC_TEST_F(test_team_cache_integration, dump_stats_integration_knob_on) +{ + UccJob job( + 4, + UccJob::UCC_JOB_CTX_GLOBAL, + {ucc_env_var_t("UCC_TEAM_CACHE_ENABLE", "y"), + ucc_env_var_t("UCC_TEAM_CACHE_DUMP_STATS", "y")}); + + UccTeam_h t = job.create_team(2, /*use_team_ep_map=*/true); + ASSERT_NE(nullptr, t); + + ucc_context_t *ctx = ctx_ptr(t, 0); + ASSERT_NE(nullptr, ctx->team_cache); + EXPECT_NE(0u, ctx->team_cache->dump_stats); + + run_barrier(t); + t.reset(); + + UccJob job_off( + 4, + UccJob::UCC_JOB_CTX_GLOBAL, + {ucc_env_var_t("UCC_TEAM_CACHE_ENABLE", "y"), + ucc_env_var_t("UCC_TEAM_CACHE_DUMP_STATS", "n")}); + + UccTeam_h t_off = job_off.create_team(2, /*use_team_ep_map=*/true); + ASSERT_NE(nullptr, t_off); + + ucc_context_t *ctx_off = ctx_ptr(t_off, 0); + ASSERT_NE(nullptr, ctx_off->team_cache); + EXPECT_EQ(0u, ctx_off->team_cache->dump_stats); + + run_barrier(t_off); + t_off.reset(); +} + +/* With DISABLE_LINEAR_CHECK=y the hash-trust path is used; for non-colliding + membership it produces the same dormant reuse as the safe path. */ +UCC_TEST_F(test_team_cache_integration, disable_linear_check_accepted) +{ + UccJob job( + 4, + UccJob::UCC_JOB_CTX_GLOBAL, + {ucc_env_var_t("UCC_TEAM_CACHE_ENABLE", "y"), + ucc_env_var_t("UCC_TEAM_CACHE_DISABLE_LINEAR_CHECK", "y")}); + + UccTeam_h t1 = job.create_team(2, /*use_team_ep_map=*/true); + ASSERT_NE(nullptr, t1); + + ucc_team_t *tp0_before = team_ptr(t1, 0); + ucc_context_t *ctx0 = ctx_ptr(t1, 0); + + ASSERT_NE(nullptr, ctx0->team_cache); + EXPECT_NE(0u, ctx0->team_cache->disable_linear_check); + + run_barrier(t1); + t1.reset(); + EXPECT_TRUE(is_dormant(ctx0, tp0_before)); + + UccTeam_h t2 = job.create_team(2, /*use_team_ep_map=*/true); + ASSERT_NE(nullptr, t2); + + ucc_team_t *tp0_after = team_ptr(t2, 0); + EXPECT_EQ(tp0_before, tp0_after) + << "hash-trust re-adopt must return the same ucc_team_t pointer"; + EXPECT_EQ(UCC_TEAM_CACHE_STATE_LIVE, tp0_after->cache_state); + EXPECT_GE(ctx0->team_cache->stats.hits, 1u); + + run_barrier(t2); +} diff --git a/test/gtest/core/test_topo.cc b/test/gtest/core/test_topo.cc index 8b58a013593..69628ff4612 100644 --- a/test/gtest/core/test_topo.cc +++ b/test/gtest/core/test_topo.cc @@ -703,3 +703,105 @@ UCC_TEST_F(test_topo, node_leaders) EXPECT_EQ(3, node_leaders[2]); EXPECT_EQ(3, node_leaders[3]); } + +/* ucc_topo_prepare_shared materializes every lazily built field so that a topo + can be shared read-only by a derived team. After it returns, no sbgp may + still be NOT_INIT and the all_* arrays must be populated, because a later + lazy build on a shared topo would be an unsynchronized write from whichever + team touched it first. */ +UCC_TEST_F(test_topo, prepare_shared_materializes_all_fields) +{ + const ucc_rank_t ctx_size = 8; + addr_storage s(ctx_size); + ucc_subset_t set; + int i; + + /* 2 nodes x 2 sockets, so every sbgp kind and the node-leaders map exist */ + SET_PI(s, 0, 0xaaa, 0, 0); + SET_PI(s, 1, 0xaaa, 0, 1); + SET_PI(s, 2, 0xaaa, 1, 2); + SET_PI(s, 3, 0xaaa, 1, 3); + SET_PI(s, 4, 0xbbb, 0, 4); + SET_PI(s, 5, 0xbbb, 0, 5); + SET_PI(s, 6, 0xbbb, 1, 6); + SET_PI(s, 7, 0xbbb, 1, 7); + + set.map.ep_num = ctx_size; + set.map.type = UCC_EP_MAP_FULL; + set.myrank = 0; + + EXPECT_EQ(UCC_OK, ucc_context_topo_init(&s.storage, &ctx_topo)); + EXPECT_EQ(UCC_OK, ucc_topo_init(set, ctx_topo, &topo)); + + /* Lazily built, so nothing is materialized yet. */ + EXPECT_EQ(nullptr, topo->all_sockets); + EXPECT_EQ(nullptr, topo->all_nodes); + EXPECT_EQ(nullptr, topo->node_leaders); + + EXPECT_EQ(UCC_OK, ucc_topo_prepare_shared(topo)); + + for (i = 0; i < UCC_SBGP_LAST; i++) { + EXPECT_NE(UCC_SBGP_NOT_INIT, topo->sbgps[i].status) + << "sbgp " << i << " left uninitialized for sharing"; + } + EXPECT_NE(nullptr, topo->all_sockets); + EXPECT_GT(topo->n_sockets, 0); + EXPECT_NE(nullptr, topo->all_nodes); + EXPECT_GT(topo->n_nodes, 0); + EXPECT_EQ(2, topo->topo->nnodes); + /* nnodes > 1, so the node-leaders map is defined and must be built. */ + EXPECT_NE(nullptr, topo->node_leaders); +} + +/* Preparing twice must be a no-op: a derived team may re-share an already + prepared parent topo, and re-running must not rebuild or reallocate. */ +UCC_TEST_F(test_topo, prepare_shared_is_idempotent) +{ + const ucc_rank_t ctx_size = 4; + addr_storage s(ctx_size); + ucc_subset_t set; + ucc_sbgp_t *sockets_first; + ucc_sbgp_t *nodes_first; + + SET_PI(s, 0, 0xaaa, 0, 0); + SET_PI(s, 1, 0xaaa, 1, 1); + SET_PI(s, 2, 0xbbb, 0, 2); + SET_PI(s, 3, 0xbbb, 1, 3); + + set.map.ep_num = ctx_size; + set.map.type = UCC_EP_MAP_FULL; + set.myrank = 0; + + EXPECT_EQ(UCC_OK, ucc_context_topo_init(&s.storage, &ctx_topo)); + EXPECT_EQ(UCC_OK, ucc_topo_init(set, ctx_topo, &topo)); + + EXPECT_EQ(UCC_OK, ucc_topo_prepare_shared(topo)); + sockets_first = topo->all_sockets; + nodes_first = topo->all_nodes; + + EXPECT_EQ(UCC_OK, ucc_topo_prepare_shared(topo)); + EXPECT_EQ(sockets_first, topo->all_sockets) + << "second prepare reallocated the socket sbgps"; + EXPECT_EQ(nodes_first, topo->all_nodes) + << "second prepare reallocated the node sbgps"; +} + +/* A NULL topo is the single-rank case (no topo is ever built) and must be + accepted, since the caller shares artifacts without checking. */ +UCC_TEST_F(test_topo, prepare_shared_accepts_null) +{ + EXPECT_EQ(UCC_OK, ucc_topo_prepare_shared(NULL)); + + /* The fixture destructor cleans these up, so give it something valid. */ + const ucc_rank_t ctx_size = 2; + addr_storage s(ctx_size); + ucc_subset_t set; + + SET_PI(s, 0, 0xaaa, 0, 0); + SET_PI(s, 1, 0xaaa, 0, 1); + set.map.ep_num = ctx_size; + set.map.type = UCC_EP_MAP_FULL; + set.myrank = 0; + EXPECT_EQ(UCC_OK, ucc_context_topo_init(&s.storage, &ctx_topo)); + EXPECT_EQ(UCC_OK, ucc_topo_init(set, ctx_topo, &topo)); +} diff --git a/test/mpi/Makefile.am b/test/mpi/Makefile.am index 8f00081c16e..5ad24dd1ff9 100644 --- a/test/mpi/Makefile.am +++ b/test/mpi/Makefile.am @@ -29,7 +29,10 @@ ucc_test_mpi_SOURCES = \ test_gatherv.cc \ test_scatter.cc \ test_scatterv.cc \ - test_mem_map.cc + test_mem_map.cc \ + test_team_cache.cc + +EXTRA_DIST = run_cache_equivalence.sh CXX=$(MPICXX) LD=$(MPICXX) diff --git a/test/mpi/main.cc b/test/mpi/main.cc index cf3435076ed..d8b21f02afb 100644 --- a/test/mpi/main.cc +++ b/test/mpi/main.cc @@ -5,6 +5,7 @@ * See file LICENSE for terms. */ +#include #include #include #include @@ -598,10 +599,16 @@ void ProcessArgs(int argc, char** argv) } } +/* The report table has one row per collective type; the team-cache suite is not + a collective type, so it gets the trailing row. */ +#define UCC_TEST_TEAM_CACHE_ROW (ucc_ilog2(UCC_COLL_TYPE_LAST) + 1) + int main(int argc, char *argv[]) { int failed = 0; - int total_done_skipped_failed[ucc_ilog2(UCC_COLL_TYPE_LAST) + 1][4]; + int total_done_skipped_failed[UCC_TEST_TEAM_CACHE_ROW + 1][4]; + ucc_test_suite_result_t tc_result = {0, 0, 0}; + ucc_context_h cache_ctx = NULL; std::chrono::steady_clock::time_point begin; int size, required, provided, completed, rank; UccTestMpi *test; @@ -712,6 +719,28 @@ int main(int argc, char *argv[]) } std::cout << std::flush; + /* Team-cache correctness suite. Requested with + UCC_TEAM_CACHE_CORRECTNESS_TESTS=y (and UCC_TEAM_CACHE_ENABLE=y); its + results are tallied below so that a failure fails the job and a wholly + skipped run is reported as skipped rather than as a pass. + + The harness teams are destroyed first. They are live, cacheable and of + world membership, which would otherwise divert the suite's creates onto + the derived-from-live path and occupy the cache entries that the eviction + tests reason about. */ + if (!test->teams.empty()) { + const char *cache_tests_env = + getenv("UCC_TEAM_CACHE_CORRECTNESS_TESTS"); + + if (cache_tests_env && + (tolower((unsigned char)cache_tests_env[0]) == 'y' || + cache_tests_env[0] == '1')) { + cache_ctx = test->teams[0].ctx; + test->destroy_teams(); + tc_result = run_team_cache_tests(cache_ctx, rank, size); + } + } + for (auto s : test->results) { int coll_num = ucc_ilog2(std::get<0>(s)); switch(std::get<1>(s)) { @@ -727,6 +756,11 @@ int main(int argc, char *argv[]) } total_done_skipped_failed[coll_num][0]++; } + total_done_skipped_failed[UCC_TEST_TEAM_CACHE_ROW][0] = + tc_result.passed + tc_result.failed + tc_result.skipped; + total_done_skipped_failed[UCC_TEST_TEAM_CACHE_ROW][1] = tc_result.passed; + total_done_skipped_failed[UCC_TEST_TEAM_CACHE_ROW][2] = tc_result.skipped; + total_done_skipped_failed[UCC_TEST_TEAM_CACHE_ROW][3] = tc_result.failed; MPI_Iallreduce(MPI_IN_PLACE, total_done_skipped_failed, sizeof(total_done_skipped_failed)/sizeof(int), MPI_INT, MPI_MAX, MPI_COMM_WORLD, &req); @@ -771,6 +805,21 @@ int main(int argc, char *argv[]) std::endl; } + if (total_done_skipped_failed[UCC_TEST_TEAM_CACHE_ROW][0] != 0) { + int *row = total_done_skipped_failed[UCC_TEST_TEAM_CACHE_ROW]; + + num_all += row[0]; + num_done += row[1]; + num_skipped += row[2]; + num_failed += row[3]; + std::cout << + std::setw(22) << std::left << "team_cache" << + std::setw(10) << std::right << row[0] << + std::setw(10) << std::right << row[1] << + std::setw(10) << std::right << row[3] << + std::setw(10) << std::right << row[2] << + std::endl; + } std::cout << " \n===== UCC MPI TEST SUMMARY =====\n" << "total tests: " << num_all << "\n" << diff --git a/test/mpi/run_cache_equivalence.sh b/test/mpi/run_cache_equivalence.sh new file mode 100755 index 00000000000..567fdb99a01 --- /dev/null +++ b/test/mpi/run_cache_equivalence.sh @@ -0,0 +1,154 @@ +#!/bin/bash +# run_cache_equivalence.sh - cache enabled-vs-disabled equivalence test. +# +# Runs ucc_test_mpi over the same team/collective set with the cache on (pass A) +# and off (pass B), using its built-in per-collective correctness checks as the +# equivalence oracle (no cross-run diffing). Exits non-zero if any run fails. +# +# Usage: bash test/mpi/run_cache_equivalence.sh [NP] +# +# NP defaults to 8; --oversubscribe keeps it usable on a single-host container. +# Pass a larger NP when a bigger cluster is available. + +set -eE + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd -P)" + +MPIRUN="${MPIRUN:-$(command -v mpirun || true)}" +# EXE defaults to the ucc_test_mpi built next to this script, but may be +# overridden (e.g. by CI, where the build tree differs from the source tree). +EXE="${EXE:-${SCRIPT_DIR}/ucc_test_mpi}" +NP="${1:-8}" + +if [ ! -x "${MPIRUN}" ]; then + echo "CACHE_EQUIV: ERROR - mpirun not found; set MPIRUN" >&2 + exit 1 +fi + +TEAMS="world,half,odd_even,reverse" +COLLS="barrier,allreduce,bcast,alltoall,allgather" +MTYPES="host" + +# ----------------------------------------------------------------------------- +# Helpers +# ----------------------------------------------------------------------------- + +PASS_COUNT=0 +FAIL_COUNT=0 + +emit_result() { + local state="$1" # CACHE_ON | CACHE_OFF + local team="$2" # team label or "all" + local result="$3" # PASS | FAIL + local detail="${4:-}" + if [ -n "${detail}" ]; then + echo "CACHE_EQUIV_RESULT state=${state} team=${team} result=${result} (${detail})" + else + echo "CACHE_EQUIV_RESULT state=${state} team=${team} result=${result}" + fi + if [ "${result}" = "PASS" ]; then + PASS_COUNT=$((PASS_COUNT + 1)) + else + FAIL_COUNT=$((FAIL_COUNT + 1)) + fi +} + +# run_mpi_pass +# Runs ucc_test_mpi over the collective suite AND the team-cache correctness suite +# (UCC_TEAM_CACHE_CORRECTNESS_TESTS=1). The collective suite proves on-vs-off +# equivalence; the correctness suite exercises actual reuse/derivation. For the +# cache-on pass we also assert the correctness suite did NOT skip itself (which it +# does when the cache is disabled), so a silently-inert cache cannot masquerade as +# "equivalent" by never touching the cache at all. The correctness suite is +# tallied into the ucc_test_mpi report and exit code, so a failure is caught +# below. +# +# UCC_TEAM_CACHE_MAX_SIZE=2 keeps the cache small enough that the overlapping +# subcommunicator tests actually evict and diverge; without it those tests skip. +run_mpi_pass() { + local label="$1" enable_val="$2" teams_arg="$3" + + local rc=0 + local log + log="$(mktemp)" + trap 'rm -f "${log}"' RETURN + "${MPIRUN}" \ + --allow-run-as-root \ + -np "${NP}" \ + --oversubscribe \ + -x "UCC_TEAM_CACHE_ENABLE=${enable_val}" \ + -x "UCC_TEAM_CACHE_MAX_SIZE=2" \ + -x "UCC_TEAM_CACHE_CORRECTNESS_TESTS=y" \ + "${EXE}" \ + -t "${teams_arg}" \ + -c "${COLLS}" \ + --mtypes "${MTYPES}" \ + > "${log}" 2>&1 || rc=$? + cat "${log}" + + local t + local result="PASS" + [ "${rc}" -eq 0 ] || result="FAIL" + + if [ "${label}" = "CACHE_ON" ] && [ "${result}" = "PASS" ]; then + if grep -qa "SKIP all team-cache tests" "${log}"; then + echo "CACHE_EQUIV: ERROR - cache-on pass did not exercise the cache" \ + "(correctness suite skipped itself); the cache may be inert" + result="FAIL" + fi + fi + IFS=',' read -ra team_list <<< "${teams_arg}" + for t in "${team_list[@]}"; do + emit_result "${label}" "${t}" "${result}" + done + + return "${rc}" +} + +# ----------------------------------------------------------------------------- +# Pre-flight +# ----------------------------------------------------------------------------- + +if [ ! -x "${EXE}" ]; then + echo "ERROR: ucc_test_mpi not found or not executable at ${EXE}" + exit 1 +fi + +echo "========================================================================" +echo " cache equivalence test NP=${NP} teams=${TEAMS} colls=${COLLS}" +echo "========================================================================" + +# ----------------------------------------------------------------------------- +# Pass A - cache ENABLED +# ----------------------------------------------------------------------------- + +echo "" +echo "--- Pass A: UCC_TEAM_CACHE_ENABLE=y ---" +pass_a_rc=0 +run_mpi_pass "CACHE_ON" "y" "${TEAMS}" || pass_a_rc=$? + +# ----------------------------------------------------------------------------- +# Pass B - cache DISABLED +# ----------------------------------------------------------------------------- + +echo "" +echo "--- Pass B: UCC_TEAM_CACHE_ENABLE=n ---" +pass_b_rc=0 +run_mpi_pass "CACHE_OFF" "n" "${TEAMS}" || pass_b_rc=$? + +# ----------------------------------------------------------------------------- +# Summary +# ----------------------------------------------------------------------------- + +echo "" +echo "========================================================================" +echo " SUMMARY PASS=${PASS_COUNT} FAIL=${FAIL_COUNT}" +echo "========================================================================" + +if [ "${FAIL_COUNT}" -gt 0 ] || [ "${pass_a_rc}" -ne 0 ] || [ "${pass_b_rc}" -ne 0 ]; then + echo "CACHE_EQUIV_OVERALL: FAIL" + exit 1 +else + echo "CACHE_EQUIV_OVERALL: PASS" + exit 0 +fi diff --git a/test/mpi/test_mpi.cc b/test/mpi/test_mpi.cc index 3c1ce77ecc3..dbb59f4c5c1 100644 --- a/test/mpi/test_mpi.cc +++ b/test/mpi/test_mpi.cc @@ -14,8 +14,8 @@ END_C_DECLS #include #include -static ucc_status_t oob_allgather(void *sbuf, void *rbuf, size_t msglen, - void *coll_info, void **req) +ucc_status_t oob_allgather(void *sbuf, void *rbuf, size_t msglen, + void *coll_info, void **req) { MPI_Comm comm = (MPI_Comm)(uintptr_t)coll_info; MPI_Request request; @@ -25,7 +25,7 @@ static ucc_status_t oob_allgather(void *sbuf, void *rbuf, size_t msglen, return UCC_OK; } -static ucc_status_t oob_allgather_test(void *req) +ucc_status_t oob_allgather_test(void *req) { MPI_Request request = (MPI_Request)(uintptr_t)req; int completed; @@ -33,7 +33,7 @@ static ucc_status_t oob_allgather_test(void *req) return completed ? UCC_OK : UCC_INPROGRESS; } -static ucc_status_t oob_allgather_free(void *req) +ucc_status_t oob_allgather_free(void *req) { return UCC_OK; } @@ -161,14 +161,23 @@ void UccTestMpi::create_teams(std::vector &test_teams, } } -UccTestMpi::~UccTestMpi() +/* Destroy every harness team; safe to call more than once. The team-cache + suite runs on a context that must not hold unrelated live teams. */ +void UccTestMpi::destroy_teams() { for (auto &t : teams) { destroy_team(t); } + teams.clear(); for (auto &t : onesided_teams) { destroy_team(t); } + onesided_teams.clear(); +} + +UccTestMpi::~UccTestMpi() +{ + destroy_teams(); if (onesided_buffers[0]) { for (auto i = 0; i < UCC_TEST_N_MEM_SEGMENTS; i++) { ucc_free(onesided_buffers[i]); @@ -180,18 +189,21 @@ UccTestMpi::~UccTestMpi() UCC_CHECK(ucc_finalize(lib)); } -ucc_team_h UccTestMpi::create_ucc_team(MPI_Comm comm, bool is_onesided) +ucc_team_h ucc_test_create_team(ucc_context_h ctx, MPI_Comm comm, + const ucc_ep_map_t *ep_map, uint64_t ext_id, + bool work_buffer) { - ucc_context_h team_ctx = ctx; - int rank, size; + int rank, size, tmp, completed; ucc_team_h team; ucc_team_params_t team_params; ucc_status_t status; + MPI_Request req; + MPI_Comm_rank(comm, &rank); MPI_Comm_size(comm, &size); - /* Create UCC TEAM for comm world */ - team_params.mask = UCC_TEAM_PARAM_FIELD_EP | + memset(&team_params, 0, sizeof(team_params)); + team_params.mask = UCC_TEAM_PARAM_FIELD_EP | UCC_TEAM_PARAM_FIELD_EP_RANGE | UCC_TEAM_PARAM_FIELD_OOB; team_params.oob.allgather = oob_allgather; @@ -203,19 +215,25 @@ ucc_team_h UccTestMpi::create_ucc_team(MPI_Comm comm, bool is_onesided) team_params.ep = rank; team_params.ep_range = UCC_COLLECTIVE_EP_RANGE_CONTIG; - if (is_onesided) { + /* Optional fields: only the team-cache tests exercise these. */ + if (ep_map) { + team_params.mask |= UCC_TEAM_PARAM_FIELD_EP_MAP; + team_params.ep_map = *ep_map; + } + if (ext_id != 0) { + team_params.mask |= UCC_TEAM_PARAM_FIELD_ID; + team_params.id = ext_id; + } + if (work_buffer) { team_params.mask |= UCC_TEAM_PARAM_FIELD_FLAGS; team_params.flags = UCC_TEAM_FLAG_COLL_WORK_BUFFER; - team_ctx = onesided_ctx; } - UCC_CHECK(ucc_team_create_post(&team_ctx, 1, &team_params, &team)); + UCC_CHECK(ucc_team_create_post(&ctx, 1, &team_params, &team)); - MPI_Request req; - int tmp; - int completed; + /* Keep MPI progressing while UCC drives the OOB allgather. */ MPI_Irecv(&tmp, 1, MPI_INT, rank, 123, comm, &req); while (UCC_INPROGRESS == (status = ucc_team_create_test(team))) { - ucc_context_progress(team_ctx); + ucc_context_progress(ctx); MPI_Test(&req, &completed, MPI_STATUS_IGNORE); }; MPI_Send(&tmp, 1, MPI_INT, rank, 123, comm); @@ -227,6 +245,12 @@ ucc_team_h UccTestMpi::create_ucc_team(MPI_Comm comm, bool is_onesided) return team; } +ucc_team_h UccTestMpi::create_ucc_team(MPI_Comm comm, bool is_onesided) +{ + return ucc_test_create_team(is_onesided ? onesided_ctx : ctx, comm, NULL, 0, + is_onesided); +} + void UccTestMpi::create_team(ucc_test_mpi_team_t t, bool is_onesided) { ucc_team_h team; diff --git a/test/mpi/test_mpi.h b/test/mpi/test_mpi.h index 105e1b4da8d..a515f85b2d5 100644 --- a/test/mpi/test_mpi.h +++ b/test/mpi/test_mpi.h @@ -351,6 +351,7 @@ class UccTestMpi { UccTestMpi(int argc, char *argv[], ucc_thread_mode_t tm, int is_local, bool with_onesided); ~UccTestMpi(); + void destroy_teams(); void set_msgsizes(size_t min, size_t max, size_t power); void set_dtypes(std::vector &_dtypes); void set_colls(std::vector &_colls); @@ -546,5 +547,45 @@ ucc_status_t compare_buffers(void *rst, void *expected, size_t count, ucc_status_t divide_buffer(void *expected, size_t divider, size_t count, ucc_datatype_t dt); +/* OOB allgather callbacks over an MPI communicator, shared by the context, + the harness teams and the team-cache suite. @coll_info is the MPI_Comm. */ +ucc_status_t oob_allgather(void *sbuf, void *rbuf, size_t msglen, + void *coll_info, void **req); +ucc_status_t oob_allgather_test(void *req); +ucc_status_t oob_allgather_free(void *req); + +/** + * ucc_test_create_team - blocking ucc team create over @comm. + * + * Drives ucc_team_create_test to completion while keeping MPI progressing, and + * aborts the job on error. @ep_map and @ext_id are optional (NULL / 0) and are + * only used by the team-cache suite; @work_buffer requests the onesided flag. + */ +ucc_team_h ucc_test_create_team(ucc_context_h ctx, MPI_Comm comm, + const ucc_ep_map_t *ep_map, uint64_t ext_id, + bool work_buffer); + +/* Per-test outcome tally for a suite that is not a collective type. */ +typedef struct ucc_test_suite_result { + int passed; + int failed; + int skipped; +} ucc_test_suite_result_t; + +/** + * run_team_cache_tests - multi-rank team-cache correctness test suite. + * + * Run when UCC_TEAM_CACHE_CORRECTNESS_TESTS=y is set in the environment (and + * UCC_TEAM_CACHE_ENABLE=y). These tests manage team create/destroy lifecycle + * directly, so they must run after the harness teams have been destroyed. + * + * @param ctx UCC context for this process (same across all ranks). + * @param world_rank MPI rank in MPI_COMM_WORLD. + * @param world_size Total number of MPI ranks. + * + * @return per-test tally; a disabled cache reports every test as skipped. + */ +ucc_test_suite_result_t run_team_cache_tests(ucc_context_h ctx, int world_rank, + int world_size); #endif diff --git a/test/mpi/test_team_cache.cc b/test/mpi/test_team_cache.cc new file mode 100644 index 00000000000..6448f008faa --- /dev/null +++ b/test/mpi/test_team_cache.cc @@ -0,0 +1,591 @@ +/** + * Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * See file LICENSE for terms. + */ + +/* + * Multi-rank team-cache correctness tests. These cover behaviors that need real + * multi-rank OOB collectives: dormant-reuse hit counting, freed-callback safety, + * cross-rank agreement, and singleton teams. Single-process behavior is covered + * by the gtest suite. + * + * Teams are created and destroyed directly so the recreate cycle is controlled + * per test. The suite therefore runs only after the harness teams have been + * destroyed, on a context with no unrelated live teams: a live world team would + * otherwise divert a create onto the derived-from-live path and occupy cache + * entries that the eviction tests reason about. + * + * Each test returns a verdict that run_team_cache_tests reduces across ranks and + * tallies, so a skipped suite is reported as skipped rather than being + * indistinguishable from a pass. A test that detects a failure records it and + * still completes its collective sequence, so peers never block on a rank that + * left early. + */ + +#include "test_mpi.h" +#include "core/ucc_context.h" +#include "core/ucc_team_cache.h" +#include "utils/ucc_coll_utils.h" +#include +#include + +/* create/destroy/recreate cycles used by the reuse tests. */ +static const int kReuseIters = 5; + +/* Verdict of one test; run_team_cache_tests reduces these across ranks. */ +typedef enum { + TC_PASS = 0, + TC_FAIL = 1, + TC_SKIP = 2 +} tc_verdict_t; + +/* The context's team cache, or NULL when caching is disabled. */ +static ucc_team_cache_t *cache_of(ucc_context_h ctx) +{ + return ((ucc_context_t *)ctx)->team_cache; +} + +/* Report a failed check. The caller records TC_FAIL and keeps going so that its + remaining collectives stay matched with the peers'. */ +static void tc_report_fail(const char *name, int rank, const char *what) +{ + std::cerr << "*** UCC TEST FAIL: " << name << " rank " << rank << ": " + << what << "\n"; +} + +static tc_verdict_t tc_skip(const char *name, int rank, const char *why) +{ + if (0 == rank) { + std::cout << "SKIP " << name << ": " << why << "\n"; + } + return TC_SKIP; +} + +/* Pump the context until a single collective request completes; abort on error + (a failed collective leaves the ranks out of step, so there is nothing left + to report cooperatively). */ +static void progress_until(ucc_context_h ctx, ucc_coll_req_h req, + const char *what) +{ + ucc_status_t st; + + while (UCC_OK != (st = ucc_collective_test(req))) { + if (st < 0) { + std::cerr << "*** UCC TEST FAIL: " << what << " (" + << ucc_status_string(st) << ")\n"; + MPI_Abort(MPI_COMM_WORLD, -1); + } + ucc_context_progress(ctx); + } +} + +static ucc_ep_map_t ep_map_full(int size) +{ + ucc_ep_map_t m; + + memset(&m, 0, sizeof(m)); + m.type = UCC_EP_MAP_FULL; + m.ep_num = size; + return m; +} + +/* Full-membership world team (EP_MAP FULL over MPI_COMM_WORLD). */ +static ucc_team_h create_world_team(ucc_context_h ctx, int size, + uint64_t ext_id = 0) +{ + ucc_ep_map_t m = ep_map_full(size); + + return ucc_test_create_team(ctx, MPI_COMM_WORLD, &m, ext_id, false); +} + +/* Subset team over @comm using an explicit team-idx -> ctx-rank array map. */ +static ucc_team_h create_array_team(ucc_context_h ctx, MPI_Comm comm, + ucc_rank_t *map, int nmembers) +{ + ucc_ep_map_t m; + + memset(&m, 0, sizeof(m)); + m.type = UCC_EP_MAP_ARRAY; + m.ep_num = nmembers; + m.array.map = map; + m.array.elem_size = sizeof(ucc_rank_t); + return ucc_test_create_team(ctx, comm, &m, 0, false); +} + +/* destroy_ucc_team - blocking ucc_team_destroy with a context progress pump. + The harness's UccTestMpi::destroy_team spins without progressing, which is + only safe for teardown at exit. */ +static void destroy_ucc_team(ucc_team_h team, ucc_context_h ctx) +{ + ucc_status_t status; + + while (UCC_INPROGRESS == (status = ucc_team_destroy(team))) { + ucc_context_progress(ctx); + } + if (UCC_OK != status) { + std::cerr << "*** UCC TEST FAIL: ucc_team_destroy failed (" + << ucc_status_string(status) << ")\n"; + MPI_Abort(MPI_COMM_WORLD, -1); + } +} + +/* Drain dormant entries left by earlier tests, with barriers so that every rank + drains before any rank starts creating again. */ +static void drain_cache(ucc_context_h ctx) +{ + MPI_Barrier(MPI_COMM_WORLD); + ucc_team_cache_drain((ucc_context_t *)ctx); + MPI_Barrier(MPI_COMM_WORLD); +} + +/* run_barrier_on_team - blocking barrier on @team; aborts on failure. */ +static void run_barrier_on_team(ucc_team_h team, ucc_context_h ctx) +{ + ucc_coll_args_t args; + ucc_coll_req_h req; + + memset(&args, 0, sizeof(args)); + args.coll_type = UCC_COLL_TYPE_BARRIER; + + UCC_CHECK(ucc_collective_init(&args, &req, team)); + UCC_CHECK(ucc_collective_post(req)); + progress_until(ctx, req, "barrier"); + UCC_CHECK(ucc_collective_finalize(req)); +} + +/* ========================================================================== + * dormant_reuse_stats: create/destroy/recreate kReuseIters times; from iter 1 + * the dormant team is re-adopted (HIT) every time. + * ========================================================================== */ +static tc_verdict_t test_dormant_reuse_stats(ucc_context_h ctx, int world_rank, + int world_size) +{ + const char *name = "dormant_reuse_stats"; + ucc_team_cache_t *cache = cache_of(ctx); + tc_verdict_t v = TC_PASS; + uint64_t hits0, hitsN; + + /* Drain dormant squatters from prior tests (a leftover would skew hits). */ + drain_cache(ctx); + hits0 = cache->stats.hits; + + for (int i = 0; i < kReuseIters; i++) { + ucc_team_h team = create_world_team(ctx, world_size); + + run_barrier_on_team(team, ctx); + MPI_Barrier(MPI_COMM_WORLD); + destroy_ucc_team(team, ctx); + MPI_Barrier(MPI_COMM_WORLD); + } + + hitsN = cache->stats.hits; + /* The first create is a miss+insert; the recreates must all be hits. */ + if (hitsN - hits0 < (uint64_t)(kReuseIters - 1)) { + std::cerr << "*** UCC TEST FAIL: " << name << " rank " << world_rank + << ": expected >=" << (kReuseIters - 1) + << " dormant hits, got " << (hitsN - hits0) << "\n"; + v = TC_FAIL; + } + if (v == TC_PASS && 0 == world_rank) { + std::cout << "PASS " << name << "\n"; + } + return v; +} + +/* ========================================================================== + * ep_map_cb_freed_after_cache: a cached team must not retain the caller's ep_map + * callback context past its lifetime. OMPI coll/ucc passes a UCC_EP_MAP_CB whose + * cb_ctx is the MPI communicator; after the comm is freed, any deref of that + * cb_ctx is a use-after-free. Here cb_ctx is a heap box that is POISONED+FREED + * after the team goes dormant; re-adopting the team and evaluating its + * operational map must never call back into the freed box. + * ========================================================================== */ + +struct cb_ctx_box { + uint64_t magic; /* CB_CTX_MAGIC while live, poisoned after free */ + ucc_rank_t ranks[1]; /* flexible: team ep -> ctx rank (identity here) */ +}; +static const uint64_t CB_CTX_MAGIC = 0xC0FFEE5AULL; + +static uint64_t poisonable_rank_cb(uint64_t ep, void *cb_ctx) +{ + struct cb_ctx_box *box = (struct cb_ctx_box *)cb_ctx; + + /* If the cache retained this (freed) context, magic no longer matches - + fail loudly instead of silently reading poisoned memory. */ + if (box->magic != CB_CTX_MAGIC) { + std::cerr << "*** UCC TEST FAIL: use-after-free - ep_map callback " + "invoked on a freed communicator context\n"; + MPI_Abort(MPI_COMM_WORLD, -1); + } + return box->ranks[ep]; +} + +static ucc_team_h create_cb_team(ucc_context_h ctx, int size, + struct cb_ctx_box *box) +{ + ucc_ep_map_t m; + + memset(&m, 0, sizeof(m)); + m.type = UCC_EP_MAP_CB; + m.ep_num = size; + m.cb.cb = poisonable_rank_cb; + m.cb.cb_ctx = (void *)box; + return ucc_test_create_team(ctx, MPI_COMM_WORLD, &m, 0, false); +} + +static struct cb_ctx_box *alloc_cb_box(int world_size) +{ + size_t box_sz = sizeof(struct cb_ctx_box) + + (world_size - 1) * sizeof(ucc_rank_t); + struct cb_ctx_box *box = (struct cb_ctx_box *)malloc(box_sz); + + if (box == NULL) { + std::cerr << "*** UCC TEST FAIL: cb_ctx_box allocation failed\n"; + MPI_Abort(MPI_COMM_WORLD, -1); + } + box->magic = CB_CTX_MAGIC; + for (int i = 0; i < world_size; i++) { + box->ranks[i] = (ucc_rank_t)i; /* world identity mapping */ + } + return box; +} + +static tc_verdict_t test_ep_map_cb_freed_after_cache(ucc_context_h ctx, + int world_rank, + int world_size) +{ + const char *name = "ep_map_cb_freed_after_cache"; + ucc_team_cache_t *cache = cache_of(ctx); + tc_verdict_t v = TC_PASS; + uint64_t hits_before; + struct cb_ctx_box *old_box, *box; + ucc_team_h team, team2; + ucc_team_t *t; + + drain_cache(ctx); + hits_before = cache->stats.hits; + old_box = alloc_cb_box(world_size); + + /* First create + use + destroy -> the team goes dormant. */ + team = create_cb_team(ctx, world_size, old_box); + run_barrier_on_team(team, ctx); + destroy_ucc_team(team, ctx); + MPI_Barrier(MPI_COMM_WORLD); + + /* Poison but keep the box: freeing it could hand its address to the new + box, so a stale deref would read a valid magic and go unnoticed */ + old_box->magic = 0xDEADDEADULL; /* poisonable_rank_cb aborts on this */ + + /* Re-create; a hit re-adopts the dormant team with its UCC-owned map */ + box = alloc_cb_box(world_size); + team2 = create_cb_team(ctx, world_size, box); + run_barrier_on_team(team2, ctx); + + /* The re-adopt is the whole point: without a hit, nothing below exercises a + retained cb_ctx and a pass would be meaningless. */ + if (cache->stats.hits <= hits_before) { + tc_report_fail(name, world_rank, + "re-create did not hit the dormant cache, so the " + "retained-cb_ctx path was never exercised"); + v = TC_FAIL; + } + + /* Evaluate the re-adopted team's operational ctx_map for every endpoint - + the exact access TL/UCP performs when resolving a peer. With the fix the + map is UCC-owned, so this resolves correctly without calling back into the + freed box; without it, poisonable_rank_cb aborts on the poison. */ + t = (ucc_team_t *)team2; + for (ucc_rank_t e = 0; e < (ucc_rank_t)world_size; e++) { + ucc_rank_t got = ucc_ep_map_eval(UCC_TEAM_CTX_MAP(t), e); + + if (got != e) { + std::cerr << "*** UCC TEST FAIL: " << name << " rank " << world_rank + << ": operational ctx_map endpoint " << (int)e + << " resolved to " << (int)got << " (expected " << (int)e + << ")\n"; + v = TC_FAIL; + } + } + + destroy_ucc_team(team2, ctx); + free(box); + free(old_box); + + if (v == TC_PASS && 0 == world_rank) { + std::cout << "PASS " << name << "\n"; + } + return v; +} + +/* ========================================================================== + * overlap_agreement: overlapping subcommunicator scopes plus a small cache force + * DIVERGENT per-rank eviction, which previously deadlocked (one rank re-adopts a + * dormant team while a peer that evicted it enters a fresh collective build and + * waits forever). The cross-rank agreement must reconcile the split hit/miss to + * a consistent fresh build. + * + * Requires UCC_TEAM_CACHE_MAX_SIZE<=2 and >=3 ranks. Without the small cache no + * eviction happens, so the test would construct no divergence at all and still + * print PASS; the max_size guard below makes that configuration an explicit + * skip instead. + * ========================================================================== */ +static tc_verdict_t test_overlap_agreement(ucc_context_h ctx, int world_rank, + int world_size) +{ + const char *name = "overlap_agreement"; + ucc_team_cache_t *cache = cache_of(ctx); + ucc_rank_t ranksA[2] = {0, 1}; + ucc_rank_t ranksB[2] = {1, 2}; + ucc_rank_t ranksD[3] = {0, 1, 2}; + MPI_Comm commA, commB, commD; + ucc_team_h t; + + if (world_size < 3 || cache->max_size > 2) { + return tc_skip(name, world_rank, "needs >=3 ranks and MAX_SIZE<=2"); + } + + drain_cache(ctx); + + /* Overlapping member sets: A{0,1}, B{1,2}, D{0,1,2}. */ + MPI_Comm_split(MPI_COMM_WORLD, (world_rank <= 1) ? 0 : MPI_UNDEFINED, + world_rank, &commA); + MPI_Comm_split(MPI_COMM_WORLD, + (world_rank >= 1 && world_rank <= 2) ? 0 : MPI_UNDEFINED, + world_rank, &commB); + MPI_Comm_split(MPI_COMM_WORLD, (world_rank <= 2) ? 0 : MPI_UNDEFINED, + world_rank, &commD); + + /* 1) A dormant on {0,1}; 2) B dormant on {1,2} (rank 1's cache is now full + at MAX_SIZE=2); 3) D on {0,1,2} evicts the oldest dormant (A) on rank 1 + but not on rank 0 -> divergence; 4) re-create A: rank 0 re-adopts, rank 1 + missed. Pre-agreement this deadlocks; the vote must reconcile to a fresh + build on both. */ + if (commA != MPI_COMM_NULL) { + t = create_array_team(ctx, commA, ranksA, 2); + destroy_ucc_team(t, ctx); + } + if (commB != MPI_COMM_NULL) { + t = create_array_team(ctx, commB, ranksB, 2); + destroy_ucc_team(t, ctx); + } + if (commD != MPI_COMM_NULL) { + t = create_array_team(ctx, commD, ranksD, 3); + destroy_ucc_team(t, ctx); + } + if (commA != MPI_COMM_NULL) { + t = create_array_team(ctx, commA, ranksA, 2); + run_barrier_on_team(t, ctx); /* must complete, not deadlock */ + destroy_ucc_team(t, ctx); + MPI_Comm_free(&commA); + } + if (commB != MPI_COMM_NULL) { + MPI_Comm_free(&commB); + } + if (commD != MPI_COMM_NULL) { + MPI_Comm_free(&commD); + } + MPI_Barrier(MPI_COMM_WORLD); + + if (0 == world_rank) { + std::cout << "PASS " << name << "\n"; + } + return TC_PASS; +} + +/* ========================================================================== + * nonblocking_create_post: ucc_team_create_post must return promptly (post the + * vote, not block) even if one rank is late entering it. Rank 0 sleeps briefly; + * the other ranks call create_post and must return before rank 0 arrives. + * + * The threshold is the full rank-0 delay rather than a fraction of it: the + * failure being detected is create_post blocking until rank 0 shows up, and the + * CI legs run oversubscribed, where a tighter bound measures scheduling noise. + * ========================================================================== */ +static tc_verdict_t test_nonblocking_create_post(ucc_context_h ctx, + int world_rank, int world_size) +{ + const char *name = "nonblocking_create_post"; + const int kSleepMs = 2000; + tc_verdict_t v = TC_PASS; + ucc_team_params_t p; + ucc_team_h team; + ucc_status_t status; + double t_start, elapsed_ms; + + if (world_size < 2) { + return tc_skip(name, world_rank, "needs >=2 ranks"); + } + + /* Drain so this is a clean fresh create (the vote is posted regardless). */ + drain_cache(ctx); + + t_start = MPI_Wtime(); + if (world_rank == 0) { + /* Delay entering create_post: peers must not block waiting on us. */ + usleep(kSleepMs * 1000); + } + + /* Posted inline rather than through ucc_test_create_team, which blocks to + completion; only the post itself is being timed. */ + memset(&p, 0, sizeof(p)); + p.mask = UCC_TEAM_PARAM_FIELD_EP | UCC_TEAM_PARAM_FIELD_EP_RANGE | + UCC_TEAM_PARAM_FIELD_OOB | UCC_TEAM_PARAM_FIELD_EP_MAP; + p.oob.allgather = oob_allgather; + p.oob.req_test = oob_allgather_test; + p.oob.req_free = oob_allgather_free; + p.oob.coll_info = (void *)(uintptr_t)MPI_COMM_WORLD; + p.oob.n_oob_eps = world_size; + p.oob.oob_ep = world_rank; + p.ep = world_rank; + p.ep_range = UCC_COLLECTIVE_EP_RANGE_CONTIG; + p.ep_map = ep_map_full(world_size); + + UCC_CHECK(ucc_team_create_post(&ctx, 1, &p, &team)); + elapsed_ms = (MPI_Wtime() - t_start) * 1000.0; + + /* On a non-zero rank, create_post must have returned before rank 0's sleep + elapsed - proving it posted (did not block on) the vote. */ + if (world_rank != 0 && elapsed_ms >= (double)kSleepMs) { + std::cerr << "*** UCC TEST FAIL: " << name << " rank " << world_rank + << ": create_post blocked " << elapsed_ms + << "ms (rank 0 delay " << kSleepMs << "ms)\n"; + v = TC_FAIL; + } + + while (UCC_INPROGRESS == (status = ucc_team_create_test(team))) { + ucc_context_progress(ctx); + } + if (status < 0) { + std::cerr << "*** UCC TEST FAIL: " << name << " create (" + << ucc_status_string(status) << ")\n"; + MPI_Abort(MPI_COMM_WORLD, -1); + } + + run_barrier_on_team(team, ctx); + MPI_Barrier(MPI_COMM_WORLD); + destroy_ucc_team(team, ctx); + MPI_Barrier(MPI_COMM_WORLD); + + if (v == TC_PASS && 0 == world_rank) { + std::cout << "PASS " << name << "\n"; + } + return v; +} + +/* ========================================================================== + * singleton_team: a size-1 cacheable team creates + reuses correctly with no + * network vote (self-membership; the size>1 gate is not taken). Each rank builds + * its own {self} team independently over MPI_COMM_SELF, so the recreates must + * still be served from the local cache. + * ========================================================================== */ +static tc_verdict_t test_singleton_team(ucc_context_h ctx, int world_rank, + int world_size) +{ + const char *name = "singleton_team"; + ucc_team_cache_t *cache = cache_of(ctx); + tc_verdict_t v = TC_PASS; + ucc_rank_t self_map[1]; + uint64_t hits0; + + (void)world_size; + drain_cache(ctx); + hits0 = cache->stats.hits; + + for (int i = 0; i < 3; i++) { + ucc_team_h t; + + self_map[0] = (ucc_rank_t)world_rank; /* team idx 0 -> my ctx rank */ + t = create_array_team(ctx, MPI_COMM_SELF, self_map, 1); + run_barrier_on_team(t, ctx); + destroy_ucc_team(t, ctx); + } + + /* The first create is a miss; the other two must re-adopt the dormant team. + Without this the test verified nothing beyond "did not crash". */ + if (cache->stats.hits - hits0 < 2) { + std::cerr << "*** UCC TEST FAIL: " << name << " rank " << world_rank + << ": expected >=2 singleton cache hits, got " + << (cache->stats.hits - hits0) << "\n"; + v = TC_FAIL; + } + + MPI_Barrier(MPI_COMM_WORLD); + if (v == TC_PASS && 0 == world_rank) { + std::cout << "PASS " << name << "\n"; + } + return v; +} + +/* Reduce a per-rank verdict to a suite-wide one: any FAIL makes the test a + failure, otherwise any SKIP makes it a skip. */ +static tc_verdict_t tc_reduce(tc_verdict_t local) +{ + int flags[2]; + + flags[0] = (local == TC_FAIL) ? 1 : 0; + flags[1] = (local == TC_SKIP) ? 1 : 0; + MPI_Allreduce(MPI_IN_PLACE, flags, 2, MPI_INT, MPI_MAX, MPI_COMM_WORLD); + if (flags[0]) { + return TC_FAIL; + } + return flags[1] ? TC_SKIP : TC_PASS; +} + +static void tc_tally(ucc_test_suite_result_t *r, tc_verdict_t local) +{ + switch (tc_reduce(local)) { + case TC_PASS: + r->passed++; + break; + case TC_FAIL: + r->failed++; + break; + default: + r->skipped++; + break; + } + MPI_Barrier(MPI_COMM_WORLD); +} + +ucc_test_suite_result_t run_team_cache_tests(ucc_context_h ctx, int world_rank, + int world_size) +{ + /* Must match the number of tc_tally calls below: used to report every test + as skipped when caching is off, so a disabled run is never mistaken for a + clean one. */ + const int kNumTests = 5; + ucc_test_suite_result_t r = {0, 0, 0}; + + if (0 == world_rank) { + std::cout << "\n===== UCC Team Cache Correctness Tests =====\n"; + } + + /* These tests require caching to be enabled. If UCC_TEAM_CACHE_ENABLE was + not set, report the whole suite as skipped rather than let a test + MPI_Abort the job. */ + if (cache_of(ctx) == NULL) { + if (0 == world_rank) { + std::cout << "SKIP all team-cache tests: caching disabled " + "(set UCC_TEAM_CACHE_ENABLE=y)\n" + << "===== Team Cache Tests DONE =====\n"; + } + r.skipped = kNumTests; + return r; + } + + tc_tally(&r, test_dormant_reuse_stats(ctx, world_rank, world_size)); + tc_tally(&r, test_ep_map_cb_freed_after_cache(ctx, world_rank, world_size)); + tc_tally(&r, test_overlap_agreement(ctx, world_rank, world_size)); + tc_tally(&r, test_nonblocking_create_post(ctx, world_rank, world_size)); + tc_tally(&r, test_singleton_team(ctx, world_rank, world_size)); + + ucc_assert(r.passed + r.failed + r.skipped == kNumTests); + + if (0 == world_rank) { + std::cout << "===== Team Cache Tests DONE (" << r.passed << " passed, " + << r.failed << " failed, " << r.skipped << " skipped) =====\n" + << std::endl; + } + return r; +}