diff --git a/accelforge/model/_looptree/accesses.py b/accelforge/model/_looptree/accesses.py
index d54f2999..a643a5d6 100755
--- a/accelforge/model/_looptree/accesses.py
+++ b/accelforge/model/_looptree/accesses.py
@@ -5,7 +5,7 @@
import islpy as isl
-from accelforge.model._looptree.reuse.isl import IslReuseAnalysisOutput
+from accelforge.model._looptree.reuse.isl.des import IslReuseAnalysisOutput
from accelforge.util._frozenset import oset
from accelforge.model._looptree.reuse.symbolic import (
BuffetStats,
diff --git a/accelforge/model/_looptree/latency/latency.py b/accelforge/model/_looptree/latency/latency.py
index ab1810bd..9cdfadc4 100755
--- a/accelforge/model/_looptree/latency/latency.py
+++ b/accelforge/model/_looptree/latency/latency.py
@@ -4,7 +4,7 @@
# from accelforge.model._looptree._isl.singular import get_value_from_singular_qpolynomial
from accelforge.frontend.arch import Compute
from accelforge.model._looptree.latency.processors import LATENCY_PROCESSORS
-from accelforge.model._looptree.reuse.isl import IslReuseAnalysisOutput
+from accelforge.model._looptree.reuse.isl.des import IslReuseAnalysisOutput
from accelforge.model._looptree.reuse import SymbolicAnalysisOutput
from accelforge.util._sympy.broadcast_max import max_nonzero
diff --git a/accelforge/model/_looptree/latency/memory.py b/accelforge/model/_looptree/latency/memory.py
index ee97fbff..f90c395d 100755
--- a/accelforge/model/_looptree/latency/memory.py
+++ b/accelforge/model/_looptree/latency/memory.py
@@ -8,7 +8,7 @@
from accelforge.model._looptree.accesses import isl_buffer_accesses_from_buffet_actions
from accelforge.model._looptree.mapping_utilities import get_leaves
-from accelforge.model._looptree.reuse.isl import IslReuseAnalysisOutput
+from accelforge.model._looptree.reuse.isl.des import IslReuseAnalysisOutput
from accelforge.model._looptree.reuse import SymbolicAnalysisOutput
from accelforge.model._looptree.types import Buffet
diff --git a/accelforge/model/_looptree/reuse/isl/__init__.py b/accelforge/model/_looptree/reuse/isl/__init__.py
index e295e773..e69de29b 100755
--- a/accelforge/model/_looptree/reuse/isl/__init__.py
+++ b/accelforge/model/_looptree/reuse/isl/__init__.py
@@ -1,5 +0,0 @@
-from .des import IslReuseAnalysisOutput
-
-__all__ = [
- "IslReuseAnalysisOutput",
-]
diff --git a/accelforge/model/_looptree/reuse/isl/des.py b/accelforge/model/_looptree/reuse/isl/des.py
index a7890af3..7b8d2de4 100755
--- a/accelforge/model/_looptree/reuse/isl/des.py
+++ b/accelforge/model/_looptree/reuse/isl/des.py
@@ -1,6 +1,5 @@
"""
-TODO: Is this file still necessary? It is referenced elsewhere but is no longer
- the format we are looking for.
+Deserializes LoopTree reuse-analysis output into ISL objects.
"""
from dataclasses import dataclass, field
diff --git a/accelforge/model/_looptree/reuse/isl/distributed/README.md b/accelforge/model/_looptree/reuse/isl/distributed/README.md
new file mode 100644
index 00000000..696c7f77
--- /dev/null
+++ b/accelforge/model/_looptree/reuse/isl/distributed/README.md
@@ -0,0 +1,529 @@
+# Transfer models: `HypercubeMulticastModel` and how to write a new one
+
+This directory holds **distributed / network cost models** for spatial reuse analysis. This
+document explains how `HypercubeMulticastModel` works and gives a step-by-step recipe (plus a fully
+worked example) for adding your own model.
+
+---
+
+## 1. Overview
+
+A **transfer model** estimates the on-chip data-movement cost of a mapping: given where data
+*lives* (occupancy) and where it is *needed* (fills), it computes how many network *hops* the
+delivery costs and which transfers are fulfilled peer-to-peer vs. read from a parent.
+
+All transfer models implement one abstract interface, `TransferModel`. The concrete
+implementations:
+
+| Model | File | Shape |
+|-------|------|-------|
+| `SimpleLinkTransferModel` | [`../spatial.py`](../spatial.py) | Neighbor-to-neighbor mesh; no constructor state |
+| `HypercubeMulticastModel` | [`distributed_buffers.py`](distributed_buffers.py) | Distance-aware worst-case multicast; takes a `dist_fn` |
+| `FullyConnectedMulticastModel` | [`distributed_buffers.py`](distributed_buffers.py) | Fully-connected fabric; 1 hop per fabric crossing (see §6) |
+| `XYRoutingMulticastModel` | [`distributed_buffers.py`](distributed_buffers.py) | XY / dimension-order routing on a 2-D mesh; X-then-Y multicast tree; per-link `EdgePressure` (see §7–§8) |
+| `StarMulticastModel` | [`distributed_buffers.py`](distributed_buffers.py) | Star / central-switch — the spokes realization of a fully-connected fabric; per-spoke `EdgePressure` (see §8). Parameter-free constant regime only, like XY |
+
+> **There is no registry or factory.** Models are constructed directly
+> (`HypercubeMulticastModel(dist_fn)`) and applied via `.apply(...)`. The only current usage is the
+> test suite — [`tests/isl/distributed/test_multicast.py`](../../../../../../tests/isl/distributed/test_multicast.py).
+> The suite lives at `tests/isl/distributed/` and runs as part of CI.
+
+> **Module layout.** This directory splits into three files: the four model classes above and
+> the shared `MulticastModel` base they all inherit (one `apply()`, one abstract `_transfer_cost`
+> hook per model) live in [`distributed_buffers.py`](distributed_buffers.py); the `EdgePressure`
+> per-link-load abstraction and its scalar-extraction helpers live in
+> [`edge_pressure.py`](edge_pressure.py); multicast-network construction (`identify_mesh_casts`
+> and the helpers built on its result) lives in [`mesh_casts.py`](mesh_casts.py).
+
+---
+
+## 2. The `TransferModel` contract
+
+Defined in [`../spatial.py`](../spatial.py). You implement exactly one abstract method:
+
+```python
+class TransferModel(ABC):
+ @abstractmethod
+ def apply(self, buff: MappingNode, fills: Fill, occs: Occupancy) -> TransferInfo:
+ ...
+```
+
+### Inputs
+
+- **`buff: MappingNode`** — the buffer being analyzed. (`HypercubeMulticastModel` ignores it and
+ relies on its `dist_fn` instead; `SimpleLinkTransferModel` uses it to find spatial dims.)
+- **`fills: Fill`** — a *tagged* `isl.Map` `{ [spacetime] -> [data] }` describing what each element
+ needs from its parent over time.
+- **`occs: Occupancy`** — a *tagged* `isl.Map` `{ [spacetime] -> [data] }` describing what each
+ element holds over time.
+
+`Fill` and `Occupancy` are `TaggedMap`s (see
+[`../mapping_to_isl/types.py`](../mapping_to_isl/types.py)): the raw relation is `.map_`, and
+`.tags` is a list of `Tag`s labeling each **input** dimension. Both constructors assert
+`len(tags) == map_.dim(isl.dim_type.in_)`. The relevant tags are:
+
+- `TemporalTag()` — that dimension spreads over time.
+- `SpatialTag(spatial_dim, buffer)` — that dimension spreads over space, in `buffer`.
+
+### Output: `TransferInfo`
+
+A frozen dataclass (in `../spatial.py`) you must fully populate:
+
+| Field | Type | Meaning |
+|-------|------|---------|
+| `fulfilled_fill` | `Transfers` | Fills satisfied by peer-to-peer transfers (a tagged map) — the fills *covered* by a matched multicast source. |
+| `unfulfilled_fill` | `Fill` | Fills *not* satisfied — no source held the datum, so it must come from higher in the hierarchy. |
+| `parent_reads` | `Reads` | Fills satisfied by parent-to-child reads. |
+| `hops` | `isl.PwQPolynomial` | The transfer cost metric across spacetime. |
+| `link_transfer` | `bool` | Metadata flag — whether this used link transfers. |
+| `edge_pressure` | `Optional[EdgePressure]` | Per-directed-edge load backing `hops` (see §8), for models that define a per-link topology. Defaults to `None`. |
+
+`fulfilled_fill` and `unfulfilled_fill` **partition the fills exactly**: with
+`covered = _covered_fills(mcs)` (the `{ dst -> data }` pairs `identify_mesh_casts` matched to a
+source), `fulfilled = fills ∩ covered` and `unfulfilled = fills − covered`. A destination requesting
+a datum that *no* source holds simply never appears in the multicast networks, so it lands in
+`unfulfilled_fill` — it is never silently treated as fulfilled.
+
+`edge_pressure` is populated by the models with an explicit per-link topology
+(`XYRoutingMulticastModel`: mesh links; `StarMulticastModel`: spokes) inside `apply()`, from the
+same `identify_mesh_casts` result as `hops`. `HypercubeMulticastModel` and
+`FullyConnectedMulticastModel` leave it `None` — a convex bounding box has no notion of individual
+links, and a contention-free full mesh (one dedicated link per pair) has no shared-link pressure to
+report.
+
+`Transfers` and `Reads` are thin `TaggedMap` subclasses; construct them as
+`Transfers(tags, map_)` / `Reads(tags, map_)`.
+
+---
+
+## 3. How `HypercubeMulticastModel` works
+
+> **Worst-case multicast.** It assumes every multicast broadcasts to the *convex hypercube* that
+> encloses all of its sources and destinations — an upper bound on the real cost.
+
+### Constructor
+
+```python
+HypercubeMulticastModel(dist_fn: isl.Map)
+```
+
+`dist_fn` is a distance function `{ [dst -> src] -> [hops] }` — note the orientation: its domain is
+a **`[dst -> src]` pair** (destination first), because `identify_mesh_casts` composes it as
+`fills_to_matches.apply_range(dist_fn)` onto `{ ... -> [dst -> src] }` maps. Two further caller
+contracts: the tuple names in `dist_fn`'s domain must match the spacetime tuple names of
+`fills`/`occs` (ISL raises on a name mismatch when `dist_fn` is applied), and its range tuple must
+be named `hops` (the fully-connected and star models filter on `{ hops[h] : h >= 1 }` to tell
+self-deliveries apart from fabric crossings). Two assumptions are baked in:
+
+1. **Orthogonal dimensions / Manhattan distance** — each unit move along a dimension costs 1 hop,
+ and dimensions are orthogonal in the metric space.
+2. **Translational invariance** — distance depends only on the displacement: if
+ `|src − dst| = |src' − dst'|` then `dist_fn(dst, src) = dist_fn(dst', src')`.
+
+These come from `calculate_extents_per_dim`, which the model relies on.
+
+### `apply(...)` data flow
+
+```
+occs.map_ { [spacetime] -> [data] } fills.map_ { [spacetime] -> [data] }
+ \ /
+ \ /
+ identify_mesh_casts(occs, fills, dist_fn)
+ |
+ v
+ mcs : { [data] -> [dst -> src] } (each datum's nearest-source multicast network)
+ |
+ _cost_mesh_cast_hypercube(mcs)
+ |
+ v
+ hops : isl.PwQPolynomial (total upper-bound hop count)
+```
+
+**Step A — `identify_mesh_casts(src_occupancy, dst_fill, dist_fn)`**
+([`mesh_casts.py`](mesh_casts.py)).
+
+For every datum, it pairs the destinations that request it with the *nearest* source that holds it.
+Conceptually:
+
+1. Reverse occupancy to `{ [data] -> [src] }` (which elements hold each datum).
+2. Match each fill `{ [dst] -> [data] }` against the sources holding that datum →
+ `{ [dst -> data] -> [dst -> src] }`.
+3. Apply `dist_fn` and take `lexmin` over distance to keep only the closest source per
+ `(dst, data)`.
+4. Regroup as `{ [data] -> [dst -> src] }` and `lexmin` again to *devolve to a single source* when
+ several are equidistant.
+
+The result is the set of **multicast networks** (`mcns`): per datum, the destinations grouped with
+their chosen source.
+
+**Step B — extents per dimension: `calculate_extents_per_dim(mcns)`**
+([`mesh_casts.py`](mesh_casts.py)).
+
+For each multicast network it unions the sources with the destinations, then for each NoC dimension
+projects away the others and takes `dim_max − dim_min`. That difference is the **extent** (the side
+length of the bounding box) along that dimension, returned as one `isl.PwAff` per dimension.
+
+**Step C — hypercube cost: `_cost_mesh_cast_hypercube(mcns)`**
+([`distributed_buffers.py`](distributed_buffers.py)).
+
+It folds the per-dimension extents into a single cost, then `.sum()`s over all networks.
+
+> ⚠️ **Known discrepancy — this is the bug currently being tracked.**
+> The docstring/comment says the cost is
+> `(∏_i extent_i) − 1` (the number of *interior* points of the bounding box minus one — the
+> standard hypercube broadcast cost). **But the code actually computes**
+>
+> ```
+> cost = (∏_i (extent_i + 1)) − 1
+> ```
+>
+> because the loop multiplies by `dim_extent.add(one)` (`extent_i + 1`), not `extent_i`. Since
+> `extent_i = max − min` is already *one less* than the number of points along a dimension, adding
+> 1 back double-counts the span. This is consistent with the observed **~3× overestimate of hops in
+> all-to-all** topologies noted in recent work. Treat the formula in the code as *not yet correct*;
+> a new template should decide deliberately whether it wants `extent_i` or `extent_i + 1`.
+
+### Return value
+
+```python
+# { dst -> data } fills actually covered by a matched source.
+covered: isl.Map = _covered_fills(mcs)
+
+return TransferInfo(
+ fulfilled_fill=Transfers(fills.tags, fills.map_.intersect(covered)),
+ parent_reads=Reads(occs.tags, mcs), # the multicast map
+ unfulfilled_fill=Fill(fills.tags, fills.map_.subtract(covered)),
+ hops=result, # the PwQPolynomial cost
+ link_transfer=True,
+ # No per-link decomposition defined for the hypercube abstraction.
+ edge_pressure=None,
+)
+```
+
+Note the partition idiom: `_covered_fills(mcs)` reshapes the multicast networks
+`{ [data] -> [dst -> src] }` into `{ dst -> data }` (via
+`range_reverse().uncurry().domain_factor_domain().reverse()` — "forget which source served it, keep
+dst and data"), which lines up with `fills.map_` so `intersect`/`subtract` split the fills into the
+covered and uncovered halves. Earlier revisions returned the *entire* fill map as `fulfilled_fill`
+and an always-empty `unfulfilled_fill` (`fills.map_.subtract(fills.map_)`); that silently
+mis-reported fills whose datum no source holds — do not copy that idiom.
+
+---
+
+## 4. Contrast: `SimpleLinkTransferModel`
+
+For comparison (the *other* valid shape — no constructor state), `SimpleLinkTransferModel`
+([`../spatial.py`](../spatial.py)):
+
+- Takes **no `dist_fn`**; constructed as `SimpleLinkTransferModel()`.
+- Asserts `fills.tags == occs.tags`.
+- Builds a **neighbor-to-neighbor mesh** via `make_mesh_connectivity` (only 1 or 2 spatial dims
+ supported — raises otherwise).
+- Data reachable from a neighbor is `fulfilled_fill`; the rest is `unfulfilled_fill`.
+- `hops` is **1 per neighbor-filled element** (`PwQPolynomial.one_on_domain(...)`), not a
+ distance-weighted sum.
+- Has an early-out: if there is no temporal dimension or no spatial dimension, nothing moves, so it
+ returns an empty/zero `TransferInfo`.
+
+So these two models bracket the design space: **stateful + distance-aware** (hypercube) vs.
+**stateless + fixed-topology** (simple link).
+
+---
+
+## 5. How to create a new template
+
+1. **Pick a home.** Add your class under `reuse/isl/`. Put distributed / NoC / mesh models in this
+ `distributed/` directory next to `distributed_buffers.py`.
+2. **Subclass `MulticastModel`** if your model is distance-driven multicast (one `dist_fn`,
+ nearest-source matching) — the base owns the constructor, the single `identify_mesh_casts`
+ call, the fill partition, and the `TransferInfo` assembly. Only subclass `TransferModel`
+ directly (and write your own `apply`) for a different shape entirely, like the stateless
+ `SimpleLinkTransferModel`.
+3. **Implement `_transfer_cost(self, mcs) -> isl.PwQPolynomial | EdgePressure`.** `mcs` is the
+ `{ [data] -> [dst -> src] }` multicast networks, computed once by the base `apply`. Reuse the
+ existing kernels:
+ - `calculate_extents_per_dim(mcns)` if you want per-dimension bounding-box extents.
+ - `_fabric_crossing(mcns, self.dist_fn)` to keep only deliveries that cross the fabric.
+ - `_edge_pressure_from_links(...)` if you report a per-link `EdgePressure` (see §8) — return
+ it directly and the base derives `hops` from its total; return a bare `isl.PwQPolynomial`
+ and the base leaves `edge_pressure` at `None`.
+4. **Honor the invariants.** The base guarantees the big ones structurally: the fill partition
+ (`fulfilled = fills ∩ covered`, `unfulfilled = fills − covered`), one `identify_mesh_casts`
+ call per `apply`, and `hops == EdgePressure.total()` whenever you return a pressure. What is
+ left to you: build the cost over the correct domain, and decide deliberately whether your
+ topology defines a per-link decomposition or not.
+5. **Add a test** mirroring
+ [`tests/isl/distributed/test_multicast.py`](../../../../../../tests/isl/distributed/test_multicast.py):
+ a YAML-driven gamut of `(dims, fill, occ, dist_fn, expected_hops)` cases.
+
+### Copy-paste skeleton
+
+```python
+import islpy as isl
+
+from accelforge.model._looptree.reuse.isl.distributed.distributed_buffers import (
+ MulticastModel,
+)
+
+
+class MyMulticastModel(MulticastModel):
+ """One-line description of the topology/assumptions this model encodes."""
+
+ def _transfer_cost(self, mcs: isl.Map) -> isl.PwQPolynomial:
+ # `mcs` is { [data] -> [dst -> src] } from `identify_mesh_casts`,
+ # computed once by `MulticastModel.apply`; the fill partition and the
+ # `TransferInfo` assembly are already handled there.
+ # TODO: your cost kernel. Return an `EdgePressure` instead (see §8) if
+ # your topology defines per-link loads -- `apply` then derives `hops`
+ # from its total automatically.
+ raise NotImplementedError
+```
+
+(A non-multicast model — different matching, no `dist_fn` — instead subclasses `TransferModel`
+directly and implements the full `apply`; use `SimpleLinkTransferModel` in
+[`../spatial.py`](../spatial.py) as the reference for that shape.)
+
+**Checklist**
+
+- [ ] Subclasses `MulticastModel` and implements `_transfer_cost` (only non-multicast shapes
+ subclass `TransferModel` and hand-roll `apply`).
+- [ ] Constructor state matches what the cost kernel needs (the base stores `dist_fn`; override
+ `__init__` only to add state).
+- [ ] `_transfer_cost` returns an `isl.PwQPolynomial` over the right domain — or an
+ `EdgePressure` for per-link topologies, making `hops == Σ_edges load` structural.
+- [ ] A YAML-driven gamut test exists.
+
+---
+
+## 6. Worked example: `FullyConnectedMulticastModel`
+
+A real, tested model for a **fully-connected fabric** (e.g. an NVSwitch-style all-to-all). It lives
+next to the hypercube model in [`distributed_buffers.py`](distributed_buffers.py).
+
+On a fully-connected fabric every cross-node delivery costs **one hop regardless of distance**, and a
+self-delivery (source node == destination node) costs **zero**. So the cost of a mapping is just the
+number of deliveries that actually traverse the fabric:
+
+```
+cost = | { (data, dst, src) ∈ mcs : dist_fn(dst, src) ≥ 1 } |
+```
+
+`dist_fn` is used *only* to tell self-deliveries (0 hops) apart from fabric-crossing ones — its hop
+*magnitude* never enters the cost. This sidesteps the hypercube extent overestimate (see §3) entirely.
+
+```python
+class FullyConnectedMulticastModel(MulticastModel):
+ """Multicast cost on a fully-connected fabric: 1 hop per fabric crossing."""
+
+ def _transfer_cost(self, mcs: isl.Map) -> isl.PwQPolynomial:
+ return self._cost_fully_connected(mcs)
+
+ def _cost_fully_connected(self, mcns: isl.Map) -> isl.PwQPolynomial:
+ """Count the deliveries in `mcns` that traverse the fabric (dist >= 1)."""
+ # [dst -> src] pairs that actually traverse the fabric (>= 1 hop).
+ crossing: isl.Map = _fabric_crossing(mcns, self.dist_fn)
+ return crossing.wrap().card()
+```
+
+How the kernel works: `_fabric_crossing` ([`mesh_casts.py`](mesh_casts.py)) computes
+`dist_fn.intersect_range({ hops[h] : h ≥ 1 }).domain()` — the set of `[dst -> src]` pairs that
+cross the fabric — and intersects `mcns`' range with it, keeping only crossing deliveries;
+`.wrap().card()` then counts the `(data, dst, src)` points as an `isl.PwQPolynomial` (constant
+when there are no parameters). Everything else — the single `identify_mesh_casts` call, the fill
+partition, `edge_pressure=None` — is `MulticastModel.apply`'s job, not this class's.
+
+**Verified numbers** (8-GPU one-hot encoding, from the test below):
+
+| Case | `FullyConnectedMulticastModel` | `HypercubeMulticastModel` |
+|------|-------------------------------|---------------------------|
+| all-to-all (64 chunks, 8 self) | **56** | 168 |
+| single unicast GPU0 → GPU3 | **1** | 3 |
+| self chunk GPU5 → GPU5 | **0** | 0 |
+
+The all-to-all column is the headline: **56 vs 168 is exactly the ~3× overestimate** the hypercube
+model incurs on a fully-connected fabric (each unicast costed as a `(1+1)(1+1) − 1 = 3` hypercube
+instead of a single crossing — see the discrepancy in §3).
+
+This example is exercised by a real test:
+[`tests/isl/distributed/test_fully_connected.py`](../../../../../../tests/isl/distributed/test_fully_connected.py)
+with cases in
+[`tests/isl/distributed/fully_connected/test_cases.yaml`](../../../../../../tests/isl/distributed/fully_connected/test_cases.yaml).
+Run it (m4 on `PATH` per the islpy-barvinok setup):
+
+```bash
+PATH="$HOME/.local/bin:$PATH" .venv/bin/python -m pytest \
+ tests/isl/distributed/test_fully_connected.py -q
+```
+
+---
+
+## 7. XY (dimension-order) routing: `XYRoutingMulticastModel`
+
+`XYRoutingMulticastModel` ([`distributed_buffers.py`](distributed_buffers.py)) models **XY routing**
+on a 2-D mesh: every packet travels **along X first, then Y**. A multicast from one source is
+therefore a rigid tree:
+
+1. an **X segment** along the source's row, reaching every column that holds a destination, then
+2. an independent **Y segment** down each of those columns, *starting from the source's row*.
+
+The cost of one tree (source `s = (xs, ys)`, destinations `D`) is:
+
+```
+x_extent({xs} ∪ {dst columns})
+ + Σ over destination columns xd of y_extent({ys} ∪ {dst y's in column xd})
+```
+
+Source selection is **per destination** (`identify_mesh_casts` pairs each destination with its
+nearest source, devolving ties); destinations sharing a source form one tree, and the model sums
+over all trees and all data.
+
+### Where it sits: `extent_DOR` (floor) ≤ XY ≤ hypercube
+
+Because each column's Y segment restarts from the source row instead of sharing a trunk, XY is an
+**upper bound on free routing** and a **lower bound on the hypercube** (which reaches every node in
+the bounding box). Note `extent_DOR_hops` in the test yamls is the **free-routing floor**, *not* the
+XY cost — it assumes routing can move freely between dimensions. Worked example (verified by the
+model): source `(1,0)` casting to `(0,2)` and `(2,2)`:
+
+```
+ y=2 D . D
+ y=1 | . |
+ y=0 +-S-+ S=(1,0); X covers cols 0..2 (2 links)
+ x=0 1 2 then Y down col 0 and col 2 from row 0 (2 + 2 links)
+
+ floor (extent_DOR) = 4 ≤ XY = 6 ≤ hypercube = 8
+```
+
+### Cost kernel
+
+There is no standalone cost method: the cost *is* the per-edge decomposition, aggregated. The
+whole model is one hook —
+
+```python
+def _transfer_cost(self, mcs: isl.Map) -> EdgePressure:
+ return _edge_pressure_from_links(self._directed_mesh_links(mcs))
+```
+
+— because it returns an `EdgePressure` (see §8), `MulticastModel.apply` sets `edge_pressure` to
+it and derives `hops = _const_pwq(pressure.total())` in one place: one aggregation path, so the
+scalar and the per-edge view can never disagree (Σ_edges load == total link count, structurally).
+
+`_directed_mesh_links` builds the X segments explicitly along each source row (out to every
+destination column, split rightward/leftward at the source column `xs`) and the Y segments down
+each destination column from the source row (split upward/downward at `ys`), keeping every link's
+identity — see §8 for the edge naming.
+
+Notes / limitations:
+- **2-D node tuples only** (no temporal dims); `apply` raises a `ValueError` ("XYRoutingMulticastModel
+ requires a 2-D node tuple (X then Y)...") for any other arity — N-D dimension-order routing is a
+ TODO. The tuple *name* is generic: it is read off the maps at `apply` time (via
+ `_mesh_node_tuple`), so `noc[x, y]`, `pe[x, y]`, etc. all work as long as `fills`, `occs`, and
+ `dist_fn` agree on it (see the caller contract in §3).
+- Each per-column/per-row link set is built explicitly and counted via `card()` rather than summing
+ a min/max polynomial — the latter trips a barvinok `summate` assertion at scale (e.g. the 8×8
+ case).
+- Returns a **parameter-free constant** (the validated regime); parametric spacetimes are future work.
+
+### Tested
+
+[`tests/isl/distributed/test_xy_routing.py`](../../../../../../tests/isl/distributed/test_xy_routing.py)
+with hand-derived cases in
+[`tests/isl/distributed/xy_routing/test_cases.yaml`](../../../../../../tests/isl/distributed/xy_routing/test_cases.yaml)
+(there is **no XY oracle in the repo**, so the expected values are hand-derived and each case carries
+its geometry). Cases: unicast `4`, the `(1,0)` discriminator `6`, three-corner `6`, 1-D column `3`,
+replicated-source `4`, and an 8×8 scale case `448`. Run:
+
+```bash
+PATH="$HOME/.local/bin:$PATH" .venv/bin/python -m pytest \
+ tests/isl/distributed/test_xy_routing.py -q
+```
+
+---
+
+## 8. Per-edge memory pressure (link load): `EdgePressure`
+
+`hops` collapses a whole routing to one scalar, which hides a real constraint: **each physical edge
+has a finite bandwidth**, so the *busiest* link is what saturates first. `EdgePressure` keeps the
+load broken out per directed physical edge — `{ edge -> number-of-trees-crossing-it }` — which is
+exactly the quantity the production symbolic path calls `max_traffic` and feeds into the `Network`
+latency formula `max(max_hops·latency, max_link_traffic / throughput)`
+([`frontend/arch/components.py`](../../../../frontend/arch/components.py)). It is exposed as the
+`TransferInfo.edge_pressure` field: `model.apply(buff, fills, occs).edge_pressure`, an
+`Optional[EdgePressure]` populated by `XYRoutingMulticastModel` and `StarMulticastModel` inside
+`apply()` (from the same `identify_mesh_casts` result as `hops`, so the two are always consistent)
+and `None` for `HypercubeMulticastModel` / `FullyConnectedMulticastModel` (see §2). There is no
+standalone `edge_pressure(fills, occs)` method.
+
+### The primitive
+
+After `identify_mesh_casts` fixes the `[dst -> src]` pairs, build, per directed edge type, a map
+`M = { [data -> src] -> edge }` associating each multicast **tree** with every edge its route
+crosses. Then (`_edge_pressure_from_links` does exactly this, unioning the per-type results)
+
+```python
+load = M.reverse().card() # { edge -> #trees crossing it }
+```
+
+The key is the tree `(data, src)`, **not** the destination: within one tree a link is traversed once
+regardless of how many leaves hang off it, so this counts *pressure* (distinct flows on a link), not
+summed hops. `EdgePressure.total()` sums the load over every edge (this is what XY/Star `hops` is
+built from, via `_const_pwq`); `bottleneck()` returns the max load (enumerating the finite,
+parameter-free edge domain); `eval_edge(name, coords)` looks up one edge.
+
+### XY routing → directed mesh edges
+
+`XYRoutingMulticastModel._directed_mesh_links` decomposes each tree onto four directed link types:
+`xedge_r`/`xedge_l` along the source row (split at the source column `xs`) and `yedge_u`/`yedge_d`
+down each destination column (split at the source row `ys`). This **is** the §7 cost kernel — every
+link is built explicitly with its identity kept (not via `calculate_extents_per_dim`, which only
+yields a scalar length and discards *which* links are used), and the scalar `hops` is just
+`EdgePressure.total()` over these maps.
+
+Decisive cross-check (no oracle needed): **`Σ_edges load == total XY hops`**. The per-edge loads must
+sum back to the already-trusted A–F totals (4/6/6/3/4/448), so a wrong decomposition fails. Worked F
+geometry (8×8, datum `(d0,d1)` held at node `(d0,d1)`, requested by all of column `x=d0`): every
+datum of column `c` floods all 7 vertical links of column `c`, so `yedge_u[c,t]` load = `t+1` and
+`yedge_d[c,t]` = `7−t`; the busiest directed link is `7` (the physical link total is `8`). Cases A–E
+are single trees, so every edge load is `1`.
+
+### Star / fully-connected → spokes
+
+`StarMulticastModel` **is** the spokes realization of the fully-connected fabric (NVSwitch-style: each
+node has one link to a central switch). A delivery routes `src → switch → dst`, so each datum loads
+its source's egress spoke once (multicast fans out *at the switch*) and each destination's ingress
+spoke once; self-deliveries (0 hops) load nothing. `_spoke_loads` returns `spoke_in[n]` (ingress) and
+`spoke_out[n]` (egress) loads. For an N-way all-to-all every node receives `N−1` and sources `1`, so
+the **ingress spokes are hottest at `N−1`** — this is where bandwidth actually bites, whereas
+`FullyConnectedMulticastModel`'s contention-free full-mesh view has no hotspot. The two are tied by
+**`Σ_nodes ingress == FullyConnected crossing count`** (every crossing delivery is one node's
+ingress; e.g. 8-GPU all-to-all: `56`). The star scalar `hops` = `Σ egress + Σ ingress` (injections +
+deliveries).
+
+### Tested
+
+[`tests/isl/distributed/test_edge_pressure.py`](../../../../../../tests/isl/distributed/test_edge_pressure.py)
+— the XY `Σ load == hops` invariant over A–F (with the pressure taken from
+`apply(...).edge_pressure`), F's bottleneck/edge loads (`7`, `yedge_u[0,6]=7`, `yedge_d[0,1]=6`),
+single-tree unit bottlenecks, and the star spoke loads / `Σ ingress == FC count` invariant for 4-
+and 8-GPU all-to-all (star `hops` = injections + deliveries: `16`/`64`, ingress bottleneck
+`3`/`7`). Run:
+
+```bash
+PATH="$HOME/.local/bin:$PATH" .venv/bin/python -m pytest \
+ tests/isl/distributed/test_edge_pressure.py -q
+```
+
+---
+
+## 9. References
+
+- Interface + `SimpleLinkTransferModel` + `TransferInfo`: [`../spatial.py`](../spatial.py)
+- `MulticastModel` base + the four model classes, `HypercubeMulticastModel`'s
+ `_cost_mesh_cast_hypercube`: [`distributed_buffers.py`](distributed_buffers.py)
+- `identify_mesh_casts`, `calculate_extents_per_dim`, `_covered_fills`, `_mesh_node_tuple`:
+ [`mesh_casts.py`](mesh_casts.py)
+- `EdgePressure`, `_eval_const`, `_const_pwq`, `_edge_pressure_from_links`:
+ [`edge_pressure.py`](edge_pressure.py)
+- Tagged-map / tag types: [`../mapping_to_isl/types.py`](../mapping_to_isl/types.py)
+- Example test harness:
+ [`tests/isl/distributed/test_multicast.py`](../../../../../../tests/isl/distributed/test_multicast.py)
diff --git a/accelforge/model/_looptree/reuse/isl/distributed/bind.py b/accelforge/model/_looptree/reuse/isl/distributed/bind.py
index e95b2ca5..7da0a0e4 100644
--- a/accelforge/model/_looptree/reuse/isl/distributed/bind.py
+++ b/accelforge/model/_looptree/reuse/isl/distributed/bind.py
@@ -1,6 +1,6 @@
"""Applies the binding layer into one that can be used for later analysis,"""
-from accelforge.frontend.binding import Binding
+from accelforge.frontend._binding import Binding
from accelforge.frontend.mapping import Mapping
from accelforge.frontend.workload import Workload
diff --git a/accelforge/model/_looptree/reuse/isl/distributed/distributed_buffers.py b/accelforge/model/_looptree/reuse/isl/distributed/distributed_buffers.py
index 35d87cb0..92614f56 100644
--- a/accelforge/model/_looptree/reuse/isl/distributed/distributed_buffers.py
+++ b/accelforge/model/_looptree/reuse/isl/distributed/distributed_buffers.py
@@ -1,15 +1,22 @@
"""
Models for handling calculating the cost of a Workload on distributed buffer
architectures.
+
+The shared shape (one `dist_fn`, one `identify_mesh_casts` call per `apply`,
+one `TransferInfo` assembly) lives on the `MulticastModel` base class; each
+concrete model below contributes only its cost kernel via `_transfer_cost`.
+The underlying primitives -- `EdgePressure` and its scalar-extraction helpers,
+and the multicast-network construction (`identify_mesh_casts` and friends) --
+live in `edge_pressure.py` and `mesh_casts.py` respectively.
"""
-import logging
+from abc import abstractmethod
+
+from typing import Optional
import islpy as isl
from accelforge.frontend.mapping import MappingNode
-from accelforge.model._looptree.reuse.isl.isl_functions import dim_projector_mask
-from accelforge.model._looptree.reuse.isl.mapping_to_isl import DUMP_ISL_IR
from accelforge.model._looptree.reuse.isl.mapping_to_isl.types import Fill, Occupancy
from accelforge.model._looptree.reuse.isl.spatial import (
Reads,
@@ -18,174 +25,74 @@
TransferModel,
)
+from accelforge.model._looptree.reuse.isl.distributed.edge_pressure import (
+ EdgePressure,
+ _const_pwq,
+ _edge_pressure_from_links,
+ _union_pwqs,
+)
+from accelforge.model._looptree.reuse.isl.distributed.mesh_casts import (
+ identify_mesh_casts,
+ calculate_extents_per_dim,
+ _covered_fills,
+ _mesh_node_tuple,
+ _per_src,
+ _fabric_crossing,
+)
-def identify_mesh_casts(
- src_occupancy: isl.Map, dst_fill: isl.Map, dist_fn: isl.Map
-) -> isl.Map:
- """
- Given srcs with data, fills to destinations, and a distance function, identify per data
- the srcs delivering that data to dsts.
-
- Parameters
- ----------
- src_occupancy:
- An isl.Map of the form { [src] -> [data] } corresponding to the data held
- at the buffer at space `src`.
- dst_fill:
- An isl.Map of the form { [dst] -> [data] } corresponding to the data requested
- at the element at space `dst`.
- dist_fn:
- A distance function { [src -> dst] -> [hops] } that accepts two points in
- space, corresponding to the `src` and `dst`, and returns the distance
- between the two points in terms of `hops`, a quantized atomic distance of
- data transmission cost.
-
- Returns
- -------
- { [data] -> [dst -> src] } where { [dst] -> [data] } and { [src] -> [data] } are in
- `src_occupancy` and `dst_fill` respectively, and where `[dst -> src]` is the infimum of
- `dst_fn(src, dst), ∀ src, dst s.t. { [src] -> [data] } ∈ `src_occupancy` and
- `{ [dst] -> [data] }` ∈ `dst_fill`.
- """
- # Makes { [dst -> data] -> [dst -> data] }
- fill_to_fill: isl.Map = dst_fill.wrap().identity()
- if DUMP_ISL_IR:
- logging.info(f"fill_to_fill: {fill_to_fill}")
-
- # Inverts src_occupancy s.t. data -> src.
- # i.e. { [xs, ys] -> [d0, d1] } to { [d0, d1] -> [xs, ys] }
- data_presence: isl.Map = src_occupancy.reverse()
-
- # { [dst -> data] -> [dst -> src] } where src contains data.
- fills_to_matches: isl.Map = (
- fill_to_fill.uncurry() # { [[dst -> data] -> dst] -> data }
- .apply_range(data_presence) # { [[dst -> data] -> dst] -> src }
- .curry()
- ) # { [[dst -> data] -> [dst -> src] }
- if DUMP_ISL_IR:
- logging.info(f"fills_to_matches: {fills_to_matches}")
-
- # Calculates the distance of a fill to the nearest src satisfying the fill.
- # { [dst -> data] -> [dist] }
- fill_min_dist: isl.Map = fills_to_matches.apply_range(dist_fn).lexmin()
- # Isolates the relevant minimal pairs.
- # { [dst -> data] -> [dst -> src] :.dst -> src is minimized distance }
- minimal_pairs: isl.Map = (
- fill_min_dist.apply_range(
- # Note: Need to match fill -> min_dist with min_dist -> [fill -> match] as lexmin over
- # fill and match will minimize distance over the tuple (src, dst, data), but that
- # overconstrains the optimization as we want to minimize over distance (dst, data)
- # only for all src.
- fills_to_matches.range_map()
- .apply_range(dist_fn)
- .reverse()
- )
- .range()
- .unwrap()
- )
- if DUMP_ISL_IR:
- logging.info(f"minimal_pairs: {minimal_pairs}")
-
- # Isolates the multicast networks.
- # { [data] -> [dst -> src] : dst -> src is minimized distance }
- multicast_networks: isl.Map = minimal_pairs.curry().range().unwrap()
- # Devolves to a single source if multiple sources per domain point.
- multicast_networks = multicast_networks.uncurry().lexmin().curry()
-
- return multicast_networks
-
-
-def calculate_extents_per_dim(mcns: isl.Map) -> list[isl.PwAff]:
- """
- Parameters
- ----------
- mcns:
- Mesh cast-networks, or networks in which all dsts per data are grouped with
- the closest src containing the data.
-
- Returns
- -------
- A list of `isl.PwAff` that gives the max extent (length) along dim_i per mcn,
- where i is the i-th `isl.PwAff`.
-
- Preconditions
- -------------
- `mcns` were generated with a Manhattan distance `dst_fn` by `identify_mesh_casts`
- s.t. all dimensions are orthogonal to each other in a metric space, where each
- unit movement in a dimension counts as 1 hop.
- We also assume `dst_fn` is translationally invariant (i.e., ∀src, dst,
- src', dst' ∈ space, if |src - dst| = |src' - dst'|,
- dst_fn(src, dst) = dst_fn(src', dst').
+class MulticastModel(TransferModel):
"""
- # Makes mcns from { [data] -> [dst -> src] } to { [data -> src] -> [dst] }
- potential_srcs: isl.Map = mcns.range_reverse().uncurry()
- # Sources are part of the extents, so we union it with the destinations.
- # { [data -> src] -> [src] }
- srcs: isl.Map = potential_srcs.domain().unwrap().range_map()
- # { [data -> src] -> [spacetime] }
- casting_extents: isl.Map = srcs.union(potential_srcs)
-
- # Projects away all dimensions but one to find their extent for hypercube.
- dims: int = potential_srcs.range_tuple_dim()
- # Creates a mask of what to project out.
- project_out_mask: list[bool] = [True] * dims
- dim_extents: list[isl.PwAff] = [None] * dims
-
- # Gets the extents of all dimensions
- for noc_dim in range(dims):
- # Project out all the dimensions of the output besides noc_dim.
- project_out_mask[noc_dim] = False
- # { [spacetime] -> [dimension] }
- extent_mapper: isl.Map = dim_projector_mask(
- casting_extents.range().get_space(), project_out_mask
- ).reverse()
- dim_extent_space: isl.Map = casting_extents.apply_range(extent_mapper)
- project_out_mask[noc_dim] = True
-
- # Finds max(noc_dim) - min(noc_dim) for each [data -> src]
- max_extent: isl.PwAff = dim_extent_space.dim_max(0)
- min_extent: isl.PwAff = dim_extent_space.dim_min(0)
-
- # Subtracts the max from the min to get the extent per [data -> src]
- dim_extents[noc_dim] = max_extent.sub(min_extent).coalesce()
-
- return dim_extents
-
-
-class HypercubeMulticastModel(TransferModel):
- """
- Does distributed multicasting a mesh using worst-case multicasting
- behavior by assuming all multicasts are broadcasting to the convex
- hypercube that encapsulates all their destinations and sources.
+ Common shape shared by every distance-driven multicast transfer model.
+
+ Every concrete model below (`HypercubeMulticastModel`,
+ `FullyConnectedMulticastModel`, `XYRoutingMulticastModel`,
+ `StarMulticastModel`) is built from one `dist_fn` and computes `apply`
+ identically up to its cost kernel: call `identify_mesh_casts` exactly
+ once, derive the cost from that single result, partition the fills, and
+ assemble a `TransferInfo`. This base class owns that shared shape; a
+ subclass need only implement `_transfer_cost`.
+
+ Because `hops` and `edge_pressure` are both derived here from the same
+ `_transfer_cost` return value, `Σ_edges load == hops` (for models that
+ report an `EdgePressure`) is structural rather than an invariant each
+ subclass has to maintain by hand: whenever `_transfer_cost` returns an
+ `EdgePressure` `p`, `hops` is *always* `_const_pwq(p.total())`.
"""
def __init__(self, dist_fn: isl.Map):
"""
- Initializes the HypercubeMulticastModel with the distance function
- over the metric space.
-
- Because we are using calculate_extents_per_dim(mcns), we inherit the
- following requirements:
- `dst_fn` holds all dimensions are orthogonal to each other in a metric space,
- where each unit movement in a dimension counts as 1 hop.
-
- We also assume `dst_fn` is translationally invariant (i.e., ∀src, dst,
- src', dst' ∈ space, if |src - dst| = |src' - dst'|,
- dst_fn(src, dst) = dst_fn(src', dst').
+ Parameters
+ ----------
+ dist_fn:
+ A distance function { [dst -> src] -> [hops] } that accepts two
+ points in space, corresponding to `dst` and `src`, and returns
+ the distance between them in `hops`, a quantized atomic distance
+ of data transmission cost. Stored and passed straight through to
+ `identify_mesh_casts` on every `apply` call.
+
+ Caller contract (inherited from `identify_mesh_casts`): the
+ tuple names of `fills`/`occs`'s domains (the spacetime/node
+ tuple, e.g. `noc[x, y]`) must match the corresponding tuple
+ names in `dist_fn`'s domain -- ISL raises on a name mismatch
+ when `dist_fn` is applied. `dist_fn`'s range tuple must be named
+ `hops` (`_fabric_crossing`, used by `FullyConnectedMulticastModel`
+ and `StarMulticastModel`, filters on `{ hops[h] : h >= 1 }` to
+ distinguish self-deliveries from fabric-crossing ones).
"""
self.dist_fn = dist_fn
def apply(self, buff: MappingNode, fills: Fill, occs: Occupancy) -> TransferInfo:
"""
Given a buffer, its fills across time, and its occupancies across time,
- calculate the spatial transfers."
+ calculate the spatial transfers.
Parameters
----------
buff:
- The buffer whose spatial analysis is being considered. Currently,
- we rely on dist_fn to deal with this rather than buffer.
+ The buffer whose spatial analysis is being considered. Not used
+ by any current subclass (they all rely on `dist_fn` instead);
+ kept for interface symmetry with `TransferModel.apply`.
fills:
The fill of `buffer` across time from parents.
occs:
@@ -193,24 +100,88 @@ def apply(self, buff: MappingNode, fills: Fill, occs: Occupancy) -> TransferInfo
Returns
-------
- Fills that were fulfilled, Fills that were unfilled, and parent reads per
- position in spacetime. Then, gets hops per timestep.
+ A `TransferInfo` whose `fulfilled_fill`/`unfulfilled_fill` partition
+ `fills` by whether `identify_mesh_casts` found a matched source for
+ that (dst, data) pair -- a fill with no source holding its datum is
+ `unfulfilled_fill`, never silently treated as fulfilled. `hops` and
+ `edge_pressure` both come from `_transfer_cost(mcs)`: if it returns
+ an `EdgePressure` `p`, `edge_pressure=p` and
+ `hops=_const_pwq(p.total())`; otherwise `edge_pressure=None` and
+ `hops` is the returned polynomial directly (the topology has no
+ per-link decomposition to report).
"""
+ # `identify_mesh_casts` is called exactly once; every output below is
+ # derived from this single `mcs`, so they cannot fall out of sync.
mcs: isl.Map = identify_mesh_casts(occs.map_, fills.map_, self.dist_fn)
- result: isl.PwQPolynomial = self._cost_mesh_cast_hypercube(mcs)
+ cost: isl.PwQPolynomial | EdgePressure = self._transfer_cost(mcs)
+ edge_pressure: Optional[EdgePressure]
+ hops: isl.PwQPolynomial
+ if isinstance(cost, EdgePressure):
+ edge_pressure = cost
+ hops = _const_pwq(cost.total())
+ else:
+ edge_pressure = None
+ hops = cost
+ # { dst -> data } fills actually covered by a matched source.
+ covered: isl.Map = _covered_fills(mcs)
- # TODO: Read once from all buffers, assert that
- # card(mcs) == tensor_size * duplication factor
- n_meshcasts: int = mcs.card()
return TransferInfo(
- fulfilled_fill=Transfers(fills.tags, fills.map_),
+ fulfilled_fill=Transfers(fills.tags, fills.map_.intersect(covered)),
parent_reads=Reads(occs.tags, mcs),
- unfulfilled_fill=Fill(fills.tags, fills.map_.subtract(fills.map_)),
- hops=result,
+ unfulfilled_fill=Fill(fills.tags, fills.map_.subtract(covered)),
+ hops=hops,
link_transfer=True,
+ edge_pressure=edge_pressure,
)
- def _cost_mesh_cast_hypercube(self, mcns: isl.Map) -> int:
+ @abstractmethod
+ def _transfer_cost(self, mcs: isl.Map) -> isl.PwQPolynomial | EdgePressure:
+ """
+ Compute this topology's transfer cost from a fixed multicast-network map.
+
+ Parameters
+ ----------
+ mcs:
+ Multicast networks { [data] -> [dst -> src] } from
+ `identify_mesh_casts`, computed once by `apply` and shared with
+ every other output `apply` derives.
+
+ Returns
+ -------
+ Either a bare `isl.PwQPolynomial` (a topology with no per-link
+ decomposition -- `apply` sets `hops` to exactly this polynomial and
+ leaves `edge_pressure` at `None`), or an `EdgePressure` (a topology
+ with one -- `apply` sets `hops` to `_const_pwq(edge_pressure.total())`,
+ structurally keeping the scalar and the per-edge view consistent).
+ """
+ raise NotImplementedError
+
+
+class HypercubeMulticastModel(MulticastModel):
+ """
+ Does distributed multicasting a mesh using worst-case multicasting
+ behavior by assuming all multicasts are broadcasting to the convex
+ hypercube that encapsulates all their destinations and sources.
+
+ `edge_pressure` is always `None` on the resulting `TransferInfo`: the
+ hypercube model costs a convex bounding box, which has no notion of
+ individual links to report pressure on.
+
+ Preconditions
+ -------------
+ Because the cost kernel uses `calculate_extents_per_dim(mcns)`, `dist_fn`
+ must hold all dimensions orthogonal to each other in a metric space,
+ where each unit movement in a dimension counts as 1 hop, and must be
+ translationally invariant (i.e., ∀src, dst, src', dst' ∈ space, if
+ |src - dst| = |src' - dst'|, dist_fn(src, dst) = dist_fn(src', dst')).
+ """
+
+ def _transfer_cost(self, mcs: isl.Map) -> isl.PwQPolynomial:
+ # TODO: Read once from all buffers, assert that
+ # card(mcs) == tensor_size * duplication factor
+ return self._cost_mesh_cast_hypercube(mcs)
+
+ def _cost_mesh_cast_hypercube(self, mcns: isl.Map) -> isl.PwQPolynomial:
"""
Given a multicast_network, calculate the hypercube.
@@ -259,3 +230,358 @@ def _cost_mesh_cast_hypercube(self, mcns: isl.Map) -> int:
# Return the hypercube cost as a piecewise polynomial.
return hypercube_costs.sum()
+
+
+class FullyConnectedMulticastModel(MulticastModel):
+ """
+ Multicast cost model for a fully-connected fabric (e.g. an NVSwitch-style
+ all-to-all interconnect).
+
+ Every cross-node delivery costs exactly one hop regardless of distance, and
+ a self-delivery (source node == destination node) costs zero:
+
+ cost = | { (data, dst, src) in mcs : dist_fn(dst, src) >= 1 } |
+
+ This is distance-independent in magnitude (one hop per crossing); `dist_fn`
+ is used only to tell self-deliveries (0 hops) apart from fabric-crossing ones.
+
+ `edge_pressure` is always `None` on the resulting `TransferInfo`: this
+ model treats the fabric as a contention-free full mesh (one dedicated
+ link per pair), so there is no per-link pressure to report.
+ """
+
+ def _transfer_cost(self, mcs: isl.Map) -> isl.PwQPolynomial:
+ return self._cost_fully_connected(mcs)
+
+ def _cost_fully_connected(self, mcns: isl.Map) -> isl.PwQPolynomial:
+ """
+ Count the deliveries in `mcns` that traverse the fabric.
+
+ Parameters
+ ----------
+ mcns:
+ Multicast networks { [data] -> [dst -> src] } from `identify_mesh_casts`.
+
+ Returns
+ -------
+ The number of (data, dst, src) deliveries with dist_fn(dst, src) >= 1, as a
+ piecewise quasi-polynomial (constant when there are no parameters).
+ """
+ # [dst -> src] pairs that actually traverse the fabric (>= 1 hop).
+ crossing: isl.Map = _fabric_crossing(mcns, self.dist_fn)
+ return crossing.wrap().card()
+
+
+# Directed-edge vocabulary + the four name-independent link-emitting maps for
+# XY routing (parsed once at import). `_XY_EDGE_TYPES` is the single source
+# of truth for the edge names: it feeds both the maps below and
+# `XYRoutingMulticastModel.EDGE_TYPES`, so the class constant and the parsed
+# maps can never name the edges differently.
+_XY_EDGE_TYPES: tuple[str, str, str, str] = (
+ "xedge_r",
+ "xedge_l",
+ "yedge_u",
+ "yedge_d",
+)
+_XEDGE_R, _XEDGE_L, _YEDGE_U, _YEDGE_D = _XY_EDGE_TYPES
+
+# { [[col[x] -> span[t]] -> ysv[ys]] -> yedge_u/d[x, t] : ... } -- re-key a
+# column's spanned vertical links (see `_span_links`) onto the up/down edge
+# tuples, splitting direction at the source row `ys`.
+_YEDGE_U_MAP: isl.Map = isl.Map.read_from_str(
+ isl.DEFAULT_CONTEXT,
+ "{ [[col[x] -> span[t]] -> ysv[ys]] -> %s[x, t] : t >= ys }" % _YEDGE_U,
+)
+_YEDGE_D_MAP: isl.Map = isl.Map.read_from_str(
+ isl.DEFAULT_CONTEXT,
+ "{ [[col[x] -> span[t]] -> ysv[ys]] -> %s[x, t] : t < ys }" % _YEDGE_D,
+)
+# { [span[t] -> xy[xs, ys]] -> xedge_r/l[t, ys] : ... } -- re-key the source
+# row's spanned horizontal links onto the right/left edge tuples, splitting
+# direction at the source column `xs`.
+_XEDGE_R_MAP: isl.Map = isl.Map.read_from_str(
+ isl.DEFAULT_CONTEXT,
+ "{ [span[t] -> xy[xs, ys]] -> %s[t, ys] : t >= xs }" % _XEDGE_R,
+)
+_XEDGE_L_MAP: isl.Map = isl.Map.read_from_str(
+ isl.DEFAULT_CONTEXT,
+ "{ [span[t] -> xy[xs, ys]] -> %s[t, ys] : t < xs }" % _XEDGE_L,
+)
+
+
+def _span_links(coords: isl.Map) -> isl.Map:
+ """
+ Build the set of link indices spanned by a per-key coordinate map.
+
+ Parameters
+ ----------
+ coords:
+ { key -> name[c] }, one point per (key, coordinate) pair to span --
+ e.g. every destination-plus-source y within one tree's column, or
+ every destination-plus-source column along one tree's row. `name` is
+ read directly off `coords`'s range tuple, so this works for any
+ single-dimension range tuple.
+
+ Returns
+ -------
+ { key -> span[t] : min(key) <= t < max(key) } -- the contiguous run of
+ link indices between each key's minimum and maximum coordinate. Built as
+ an explicit lexmin/lexmax intersection (a `card()`-friendly link set)
+ rather than by summing a min/max polynomial, which trips a barvinok
+ `summate` assertion at scale (e.g. the 8x8 case).
+ """
+ name: str = coords.get_tuple_name(isl.dim_type.out)
+ lo: isl.Map = coords.lexmin().apply_range(
+ isl.Map.read_from_str(
+ isl.DEFAULT_CONTEXT, "{ %s[lo] -> span[t] : t >= lo }" % name
+ )
+ )
+ hi: isl.Map = coords.lexmax().apply_range(
+ isl.Map.read_from_str(
+ isl.DEFAULT_CONTEXT, "{ %s[hi] -> span[t] : t < hi }" % name
+ )
+ )
+ return lo.intersect(hi)
+
+
+class XYRoutingMulticastModel(MulticastModel):
+ """
+ Multicast cost model for XY (dimension-order) routing on a 2-D mesh.
+
+ XY routing constrains every packet to travel along the X dimension first and
+ only then along the Y dimension, so a multicast from one source forms a rigid
+ tree:
+
+ 1. an X segment along the source's row, reaching every column that holds a
+ destination, and
+ 2. an independent Y segment down each of those columns, starting from the
+ source's row.
+
+ The hop cost of one such tree (source `s = (xs, ys)` with destination set
+ `D`) is therefore::
+
+ x_extent({xs} u {xd : (xd, yd) in D})
+ + sum over destination columns xd of
+ y_extent({ys} u {yd : (xd, yd) in D})
+
+ Because the Y segments restart from the source row in every column rather than
+ sharing a trunk, this is an upper bound on free (any-monotone-path) routing and
+ a lower bound on the hypercube model, giving the ordering::
+
+ extent_DOR (floor) <= XY routing <= hypercube
+
+ For example, source `(1, 0)` casting to `(0, 2)` and `(2, 2)` costs 4
+ (floor), 6 (XY), and 8 (hypercube) respectively.
+
+ Source selection is per destination: `identify_mesh_casts` pairs each
+ destination with its nearest source (devolving ties), and destinations that
+ share a source form one tree; the cost sums over all such trees and all data.
+ `dist_fn` is expected to be Manhattan and translationally invariant, like the
+ hypercube model (see `HypercubeMulticastModel`'s Preconditions), since nearest-
+ source selection relies on the same distance shape.
+
+ Its four directed link types are declared once as the class constant
+ `EDGE_TYPES = ("xedge_r", "xedge_l", "yedge_u", "yedge_d")`.
+
+ Preconditions
+ -------------
+ The NoC is two-dimensional (no temporal dimensions in the spacetime), with
+ the first coordinate routed before the second (X then Y). The node tuple's
+ name is generic -- read off `fills`/`occs` at `apply` time, so any name
+ works (e.g. `noc[x, y]` or `pe[x, y]`) as long as `fills`, `occs`, and
+ `self.dist_fn` all agree on it (see the caller contract on
+ `identify_mesh_casts`'s `dist_fn` parameter). The tuple must be exactly
+ 2-D; `apply` raises `ValueError` otherwise. N-dimensional dimension-order
+ routing is a future extension. The returned cost is a parameter-free
+ constant (the validated regime); parametric spacetimes are not yet
+ supported.
+ """
+
+ EDGE_TYPES = _XY_EDGE_TYPES
+
+ def _transfer_cost(self, mcs: isl.Map) -> EdgePressure:
+ return _edge_pressure_from_links(self._directed_mesh_links(mcs))
+
+ def _directed_mesh_links(self, mcns: isl.Map) -> list[isl.Map]:
+ """
+ Decompose every multicast tree in `mcns` onto the directed mesh links it
+ traverses under XY routing.
+
+ Each returned map is { [data -> src] -> edge } for one directed edge
+ type, associating a tree with every link of that type on its route. The
+ X phase runs along the source row out to every destination column; the Y
+ phase runs down each destination column from the source row. Directions
+ split at the source: rightward/leftward in X (at the source column `xs`)
+ and upward/downward in Y (at the source row `ys`).
+
+ Parameters
+ ----------
+ mcns:
+ Multicast networks { [data] -> [dst -> src] } from `identify_mesh_casts`.
+
+ Returns
+ -------
+ `[xedge_r, xedge_l, yedge_u, yedge_d]` maps (see `EDGE_TYPES`). An
+ edge `xedge_r[t, ys]` is the rightward link in row `ys` between
+ columns `t` and `t + 1`; `yedge_u[x, t]` is the upward link in
+ column `x` between rows `t` and `t + 1` (and `_l` / `_d` the
+ opposite directions).
+
+ Preconditions
+ -------------
+ The node tuple embedded in `mcns` must be exactly 2-D; every map string
+ below assumes a 2-D `name[x, y]` node tuple. Raises `ValueError`
+ otherwise.
+ """
+ ctx = isl.DEFAULT_CONTEXT
+ # Reads the node tuple's name (and validates its dimensionality) off
+ # `mcns` instead of assuming the literal 'noc', so this method works
+ # for any 2-D node tuple.
+ name, dims = _mesh_node_tuple(mcns)
+ if dims != 2:
+ raise ValueError(
+ "XYRoutingMulticastModel requires a 2-D node tuple "
+ f"(X then Y); got tuple '{name}' with {dims} dimensions. "
+ "N-dimensional dimension-order routing is not yet supported."
+ )
+
+ # { [data -> src] -> dst name[x, y] } and a handle on the source per tree.
+ per_src, keymap = _per_src(mcns) # keymap: { [data->src] -> src }
+
+ # --- Y phase: vertical links per (tree, destination column). ---
+ # Key each destination's y by its column: { [data->src->col] -> yv[y] }.
+ dst_y: isl.Map = per_src.apply_range(
+ isl.Map.read_from_str(
+ ctx,
+ "{ %s[x, y] -> [col[x'] -> yv[y']] : x' = x and y' = y }" % name,
+ )
+ ).uncurry()
+ # Inject the source row ys into every destination column so each Y segment
+ # starts from the source.
+ src_y: isl.Map = keymap.apply_range(
+ isl.Map.read_from_str(ctx, "{ %s[xs, ys] -> yv[ys] }" % name)
+ )
+ src_row: isl.Map = dst_y.domain().unwrap().range_product(src_y).uncurry()
+ col_ys: isl.Map = dst_y.union(src_row)
+ # Every link spanned in a column, keyed by the column (see
+ # `_span_links`): { [data->src->col] -> span[t] }.
+ ylinks: isl.Map = _span_links(col_ys)
+ # Re-key links by the tree and carry ys so direction splits at the source:
+ # { [data->src] -> [[col[x] -> span[t]] -> ysv[ys]] }.
+ y_with_ys: isl.Map = ylinks.curry().range_product(
+ keymap.apply_range(
+ isl.Map.read_from_str(ctx, "{ %s[xs, ys] -> ysv[ys] }" % name)
+ )
+ )
+ yedge_u: isl.Map = y_with_ys.apply_range(_YEDGE_U_MAP)
+ yedge_d: isl.Map = y_with_ys.apply_range(_YEDGE_D_MAP)
+
+ # --- X phase: horizontal links along the source row. ---
+ # Columns spanned per tree = {src column} u {destination columns}. (Built
+ # explicitly, not from calculate_extents_per_dim, which keeps only the
+ # extent length and discards which links are used.)
+ col_x: isl.Map = per_src.apply_range(
+ isl.Map.read_from_str(ctx, "{ %s[x, y] -> cx[x] }" % name)
+ ).union(
+ keymap.apply_range(
+ isl.Map.read_from_str(ctx, "{ %s[xs, ys] -> cx[xs] }" % name)
+ )
+ )
+ # Every link spanned along the row, keyed by the tree (see
+ # `_span_links`): { [data->src] -> span[t] }.
+ xlinks: isl.Map = _span_links(col_x)
+ # Carry the source (xs, ys) so direction splits at xs and the row ys is
+ # part of the edge identity: { [data->src] -> [span[t] -> xy[xs, ys]] }.
+ x_with_src: isl.Map = xlinks.range_product(
+ keymap.apply_range(
+ isl.Map.read_from_str(ctx, "{ %s[xs, ys] -> xy[xs, ys] }" % name)
+ )
+ )
+ xedge_r: isl.Map = x_with_src.apply_range(_XEDGE_R_MAP)
+ xedge_l: isl.Map = x_with_src.apply_range(_XEDGE_L_MAP)
+
+ return [xedge_r, xedge_l, yedge_u, yedge_d]
+
+
+# Spoke vocabulary for the star model, declared once and shared with
+# `StarMulticastModel.EDGE_TYPES`.
+_STAR_EDGE_TYPES: tuple[str, str] = ("spoke_in", "spoke_out")
+
+
+class StarMulticastModel(MulticastModel):
+ """
+ Multicast cost model for a star / central-switch fabric -- the spokes
+ realization of a fully-connected interconnect (e.g. an NVSwitch, where every
+ GPU connects to a shared switch rather than to a full mesh of peers).
+
+ Every node has exactly one spoke (its bidirectional link to the switch). A
+ delivery routes `src -> switch -> dst`: the source injects each datum once
+ up its egress spoke (multicast fan-out happens at the switch, so one copy per
+ datum regardless of how many destinations want it), and every destination
+ receives its datum down its ingress spoke. Self-deliveries (a node already
+ holding the datum) never cross the fabric and so load no spoke; `dist_fn`'s
+ hop magnitude otherwise never enters the spoke load -- on a star every
+ crossing is one switch hop each way, so `dist_fn` is used only to pick each
+ destination's nearest source and to classify self- vs. fabric-crossing
+ deliveries (`_fabric_crossing`).
+
+ Its two spoke directions are declared once as the class constant
+ `EDGE_TYPES = ("spoke_in", "spoke_out")`.
+ """
+
+ EDGE_TYPES = _STAR_EDGE_TYPES
+
+ def _transfer_cost(self, mcs: isl.Map) -> EdgePressure:
+ return EdgePressure(_union_pwqs([*self._spoke_loads(mcs)]))
+
+ def _spoke_loads(
+ self, mcns: isl.Map
+ ) -> tuple[isl.PwQPolynomial, isl.PwQPolynomial]:
+ """
+ Per-spoke ingress and egress load for the multicast networks `mcns`.
+
+ Parameters
+ ----------
+ mcns:
+ Multicast networks { [data] -> [dst -> src] } from `identify_mesh_casts`.
+
+ Returns
+ -------
+ `(ingress, egress)` where `ingress` is { spoke_in[n] -> #data the node
+ receives } and `egress` is { spoke_out[n] -> #data the node sources },
+ counting only fabric-crossing (>= 1 hop) deliveries.
+ """
+ # Reads the node tuple's name off `mcns` instead of assuming the
+ # literal 'noc'; unlike XY, the star model has no dimensionality
+ # requirement.
+ name, _dims = _mesh_node_tuple(mcns)
+
+ # Keep only deliveries that actually cross the fabric (>= 1 hop); a node
+ # already holding its datum loads no spoke.
+ crossing: isl.Map = _fabric_crossing(mcns, self.dist_fn)
+
+ # { dst -> [src -> data] }: regroup so each delivery is keyed by destination.
+ cur: isl.Map = crossing.reverse().curry()
+ # Distinct (node, data) pairs per direction. A source injects each datum
+ # once (multicast fans out at the switch); a destination receives each once.
+ ingress_nodes: isl.Map = cur.range_factor_range() # { name[dst] -> data }
+ egress_nodes: isl.Map = cur.range().unwrap() # { name[src] -> data }
+
+ # Relabel the node tuple to the directed spoke edge, then count data per
+ # spoke. { spoke_in[n] -> #data } and { spoke_out[n] -> #data }.
+ dims: int = ingress_nodes.dim(isl.dim_type.in_)
+ idx: str = ", ".join(f"i{k}" for k in range(dims))
+ spoke_in, spoke_out = self.EDGE_TYPES
+ ingress: isl.PwQPolynomial = ingress_nodes.apply_domain(
+ isl.Map.read_from_str(
+ isl.DEFAULT_CONTEXT,
+ "{ %s[%s] -> %s[%s] }" % (name, idx, spoke_in, idx),
+ )
+ ).card()
+ egress: isl.PwQPolynomial = egress_nodes.apply_domain(
+ isl.Map.read_from_str(
+ isl.DEFAULT_CONTEXT,
+ "{ %s[%s] -> %s[%s] }" % (name, idx, spoke_out, idx),
+ )
+ ).card()
+ return ingress, egress
diff --git a/accelforge/model/_looptree/reuse/isl/distributed/edge_pressure.py b/accelforge/model/_looptree/reuse/isl/distributed/edge_pressure.py
new file mode 100644
index 00000000..629a7cfc
--- /dev/null
+++ b/accelforge/model/_looptree/reuse/isl/distributed/edge_pressure.py
@@ -0,0 +1,251 @@
+"""
+Per-edge memory pressure (link load) for distributed transfer models:
+`EdgePressure` and the scalar-extraction / accumulation helpers that build and
+consume it.
+
+Leaf module -- depends only on `islpy`, `functools`, `dataclasses`, and
+`typing`, so every other module in this package (`mesh_casts.py`,
+`distributed_buffers.py`, and `../spatial.py`) can import from here without
+risk of a cycle.
+"""
+
+import functools
+
+from dataclasses import dataclass
+
+from typing import Optional
+
+import islpy as isl
+
+
+@dataclass(frozen=True)
+class EdgePressure:
+ """
+ Per-edge memory pressure (link load) for a spatial transfer.
+
+ Where a `TransferModel`'s `hops` collapses a whole routing to a single
+ scalar, `EdgePressure` keeps the load broken out per physical edge: how many
+ multicast trees cross each individual link -- the quantity a per-link
+ bandwidth limit acts on, and what the symbolic network model calls
+ `max_traffic`.
+
+ "Edge" is a directed physical link, identified by the tuple name and
+ coordinates of `load`'s domain, e.g. `yedge_u[x, t]` (the upward vertical
+ link in column `x` between rows `t` and `t + 1`, XY mesh) or `spoke_in[n]` /
+ `spoke_out[n]` (a node's ingress / egress link to the switch, star model).
+
+ The load is keyed on the multicast tree `(data, src)`, not on individual
+ destinations: within one tree a link is traversed once regardless of how
+ many leaves hang off it, so `load` measures pressure (distinct flows over a
+ link) rather than summed hops -- e.g. for XY routing, sum over edges of load
+ == total hops.
+
+ Preconditions
+ -------------
+ `bottleneck` and `eval_edge` assume the load is piecewise constant (the
+ parameter-free regime the distributed models are validated in); there is no
+ clean ISL "max of a quasi-polynomial over its domain" primitive, so the
+ bottleneck is obtained by enumerating pieces.
+ """
+
+ # { edge -> number-of-trees } spanning every directed edge type the model
+ # emits.
+ load: isl.UnionPwQPolynomial
+
+ @functools.cached_property
+ def _pieces(self) -> list[isl.PwQPolynomial]:
+ """
+ Cache `load`'s decomposition into per-edge-type pieces.
+
+ Returns
+ -------
+ The list of `isl.PwQPolynomial` pieces `load` unions together, in
+ `foreach_pw_qpolynomial` enumeration order.
+ """
+ # Note: `cached_property` works on this frozen dataclass only because
+ # it is not `slots=True` -- the cache is written straight into the
+ # instance `__dict__`, bypassing the frozen `__setattr__`.
+ pieces: list[isl.PwQPolynomial] = []
+ self.load.foreach_pw_qpolynomial(pieces.append)
+ return pieces
+
+ def total(self) -> int:
+ """
+ Sum the load over every edge.
+
+ For a single-direction (or single edge-type) pressure this is the total
+ traffic; for the full mesh pressure it equals the model's total `hops`
+ (sum over edges of load == total hops), which is the primary cross-check
+ that the per-edge decomposition is correct.
+ """
+ total = 0
+ for pwq in self._pieces:
+ # `.sum()` collapses the edge-indexed domain away (summing over every
+ # edge of this piece), leaving a 0-set-dim, parameter-free
+ # polynomial -- exactly `_eval_const`'s precondition.
+ total += _eval_const(pwq.sum())
+ return total
+
+ def bottleneck(self) -> int:
+ """
+ Return the load on the single most-pressured edge.
+
+ This is the bandwidth-binding quantity: with a uniform per-link
+ bandwidth, the most-congested edge saturates first, so its load sets the
+ transfer's bandwidth-bound latency.
+
+ The per-edge load is generally not constant across an edge type (e.g. a
+ flooded column's upward link `yedge_u[x, t]` carries `t + 1` trees), so
+ the maximum is found by enumerating the (finite, parameter-free) edge
+ domain and evaluating the load at each edge -- sampling one point per
+ piece would under-report a monotone load.
+ """
+ best = 0
+ for pwq in self._pieces:
+ edges: list[isl.Point] = []
+ pwq.domain().foreach_point(edges.append)
+ for edge in edges:
+ best = max(best, _eval_at(pwq, edge))
+ return best
+
+ def eval_edge(self, name: str, coords: list[int]) -> int:
+ """
+ Look up the load on one named edge, e.g. `eval_edge("yedge_u", [0, 6])`.
+
+ Returns 0 if no flow crosses that edge (it is outside the load's support).
+
+ Parameters
+ ----------
+ name:
+ The edge tuple name (`xedge_r`/`xedge_l`/`yedge_u`/`yedge_d`
+ for the mesh model, `spoke_in`/`spoke_out` for the spokes model).
+ coords:
+ The integer edge coordinates within that tuple.
+ """
+ for pwq in self._pieces:
+ if pwq.domain().get_space().get_tuple_name(isl.dim_type.set) != name:
+ continue
+ point: isl.Point = isl.Set.read_from_str(
+ isl.DEFAULT_CONTEXT,
+ "{ %s[%s] }" % (name, ", ".join(str(c) for c in coords)),
+ ).sample_point()
+ return _eval_at(pwq, point)
+ return 0
+
+
+def _eval_at(pwq: isl.PwQPolynomial, point: isl.Point) -> int:
+ """
+ Evaluate a piecewise quasi-polynomial at one point and return a Python `int`.
+
+ Parameters
+ ----------
+ pwq:
+ The polynomial to evaluate.
+ point:
+ The point to evaluate at. Must lie in a space compatible with `pwq`'s
+ domain (matching tuple name and dimensionality).
+
+ Returns
+ -------
+ `pwq`'s value at `point`, as a Python `int`.
+
+ Note: isl.Val has no direct int() conversion in this islpy build (it
+ raises TypeError); round-tripping through str() is the working idiom
+ used throughout this module.
+ """
+ return int(str(pwq.eval(point)))
+
+
+def _eval_const(pwq: isl.PwQPolynomial) -> int:
+ """
+ Evaluate a parameter-free, already-reduced piecewise quasi-polynomial to its
+ scalar value.
+
+ Preconditions
+ -------------
+ `pwq` has no free set dimensions and no parameters -- e.g. the output of
+ `.card()` on a parameter-free map/set, `.sum()` on a parameter-free
+ polynomial, or a constant built by `_const_pwq`. Evaluated only at the
+ space's zero point, so a `pwq` that still varies over domain points or
+ parameters gives a meaningless result; callers must reduce to a true
+ constant first (via `.sum()`/`.card()`).
+
+ Returns
+ -------
+ The polynomial's constant value as a Python `int`.
+ """
+ return _eval_at(pwq, isl.Point.zero(pwq.domain().get_space()))
+
+
+def _const_pwq(value: int) -> isl.PwQPolynomial:
+ """
+ Build a 0-dimensional, parameter-free constant `isl.PwQPolynomial` equal to
+ `value`.
+
+ Parameters
+ ----------
+ value:
+ The integer the returned polynomial evaluates to everywhere (it has no
+ domain dimensions or parameters to vary over).
+
+ Returns
+ -------
+ An `isl.PwQPolynomial` over the empty (0-dim, 0-param) space, suitable
+ wherever a parameter-free `hops` cost is expected. Round-trips through
+ `_eval_const` back to `value`.
+ """
+ ctx = isl.DEFAULT_CONTEXT
+ zero_dim: isl.Space = isl.Space.set_alloc(ctx, 0, 0)
+ return isl.PwQPolynomial.from_qpolynomial(
+ isl.QPolynomial.val_on_domain(zero_dim, isl.Val(value, ctx))
+ )
+
+
+def _union_pwqs(pwqs: list[isl.PwQPolynomial]) -> isl.UnionPwQPolynomial:
+ """
+ Accumulate a list of piecewise quasi-polynomials into one union.
+
+ Parameters
+ ----------
+ pwqs:
+ The polynomials to sum, e.g. one per directed edge type. May be
+ empty.
+
+ Returns
+ -------
+ The union-add of every `pwqs` entry (each promoted to an
+ `isl.UnionPwQPolynomial` first). An empty `pwqs` yields the empty union
+ `{ }`, so a caller with zero edge types (or a degenerate empty mesh) does
+ not need to special-case the accumulation.
+ """
+ acc: Optional[isl.UnionPwQPolynomial] = None
+ for pwq in pwqs:
+ contribution = isl.UnionPwQPolynomial.from_pw_qpolynomial(pwq)
+ acc = contribution if acc is None else acc.add(contribution)
+ if acc is None:
+ acc = isl.UnionPwQPolynomial.read_from_str(isl.DEFAULT_CONTEXT, "{ }")
+ return acc
+
+
+def _edge_pressure_from_links(edge_maps: list[isl.Map]) -> EdgePressure:
+ """
+ Turn directed flow maps into an `EdgePressure`.
+
+ Parameters
+ ----------
+ edge_maps:
+ A list of { [data -> src] -> edge } maps, one per directed edge type,
+ each associating a multicast tree with every edge its route traverses.
+
+ Returns
+ -------
+ An `EdgePressure` whose `load` is { edge -> number-of-trees }: for each map
+ we reverse it and take the cardinality (`reverse().card()` counts, per
+ edge, how many distinct `(data, src)` trees cross it), then union the
+ per-type results into one `UnionPwQPolynomial` (via `_union_pwqs`).
+ """
+ # { edge -> #trees crossing it }, one per edge type.
+ per_edge_loads: list[isl.PwQPolynomial] = [
+ edge_map.reverse().card() for edge_map in edge_maps
+ ]
+ return EdgePressure(_union_pwqs(per_edge_loads))
diff --git a/accelforge/model/_looptree/reuse/isl/distributed/mesh_casts.py b/accelforge/model/_looptree/reuse/isl/distributed/mesh_casts.py
new file mode 100644
index 00000000..c6318a7f
--- /dev/null
+++ b/accelforge/model/_looptree/reuse/isl/distributed/mesh_casts.py
@@ -0,0 +1,268 @@
+"""
+Multicast-network construction for distributed transfer models: matching each
+requested datum to its nearest source (`identify_mesh_casts`) and the
+shared-shape helpers built on top of its result (fill-partition recovery,
+per-dimension extents, tree re-keying, and fabric-crossing filtering).
+
+Imports `islpy`, `isl_functions.dim_projector_mask`, and
+`mapping_to_isl.DUMP_ISL_IR` only -- no dependency on `edge_pressure.py`,
+`spatial.py`, or `distributed_buffers.py`, so this module cannot participate
+in the import cycle those would otherwise risk.
+"""
+
+import logging
+
+import islpy as isl
+
+from accelforge.model._looptree.reuse.isl.isl_functions import dim_projector_mask
+from accelforge.model._looptree.reuse.isl.mapping_to_isl import DUMP_ISL_IR
+
+# [dst -> src] pairs whose distance is >= 1 hop, i.e. every fabric-crossing
+# delivery (as opposed to a self-delivery, which never traverses a link).
+# Parsed once at import rather than per `_fabric_crossing` call -- the string
+# encodes `identify_mesh_casts`'s `dist_fn` contract (range tuple named
+# `hops`), not anything instance- or call-specific.
+_CROSSING_HOPS: isl.Set = isl.Set.read_from_str(
+ isl.DEFAULT_CONTEXT, "{ hops[h] : h >= 1 }"
+)
+
+
+def identify_mesh_casts(
+ src_occupancy: isl.Map, dst_fill: isl.Map, dist_fn: isl.Map
+) -> isl.Map:
+ """
+ Given srcs with data, fills to destinations, and a distance function, identify per data
+ the srcs delivering that data to dsts.
+
+ Parameters
+ ----------
+ src_occupancy:
+ An isl.Map of the form { [src] -> [data] } corresponding to the data held
+ at the buffer at space `src`.
+ dst_fill:
+ An isl.Map of the form { [dst] -> [data] } corresponding to the data requested
+ at the element at space `dst`.
+ dist_fn:
+ A distance function { [dst -> src] -> [hops] } that accepts two points in
+ space, corresponding to the `dst` and `src`, and returns the distance
+ between the two points in terms of `hops`, a quantized atomic distance of
+ data transmission cost.
+
+ Caller contract: the tuple names of `dst_fill`'s and `src_occupancy`'s
+ domains (the spacetime/node tuple, e.g. `noc[x, y]`) must match the
+ corresponding tuple names in `dist_fn`'s domain -- ISL will raise on a
+ name mismatch when `dist_fn` is applied. `dist_fn`'s range tuple must be
+ named `hops` (consumers such as `FullyConnectedMulticastModel` and
+ `StarMulticastModel` filter on `{ hops[h] : h >= 1 }` to distinguish
+ self-deliveries from fabric-crossing ones).
+
+ Returns
+ -------
+ { [data] -> [dst -> src] } where { [dst] -> [data] } and { [src] -> [data] } are in
+ `src_occupancy` and `dst_fill` respectively, and where `[dst -> src]` is the infimum of
+ `dst_fn(src, dst), ∀ src, dst s.t. { [src] -> [data] } ∈ `src_occupancy` and
+ `{ [dst] -> [data] }` ∈ `dst_fill`.
+ """
+ # Makes { [dst -> data] -> [dst -> data] }
+ fill_to_fill: isl.Map = dst_fill.wrap().identity()
+ if DUMP_ISL_IR:
+ logging.info(f"fill_to_fill: {fill_to_fill}")
+
+ # Inverts src_occupancy s.t. data -> src.
+ # i.e. { [xs, ys] -> [d0, d1] } to { [d0, d1] -> [xs, ys] }
+ data_presence: isl.Map = src_occupancy.reverse()
+
+ # { [dst -> data] -> [dst -> src] } where src contains data.
+ fills_to_matches: isl.Map = (
+ fill_to_fill.uncurry() # { [[dst -> data] -> dst] -> data }
+ .apply_range(data_presence) # { [[dst -> data] -> dst] -> src }
+ .curry()
+ ) # { [[dst -> data] -> [dst -> src] }
+ if DUMP_ISL_IR:
+ logging.info(f"fills_to_matches: {fills_to_matches}")
+
+ # Calculates the distance of a fill to the nearest src satisfying the fill.
+ # { [dst -> data] -> [dist] }
+ fill_min_dist: isl.Map = fills_to_matches.apply_range(dist_fn).lexmin()
+ # Isolates the relevant minimal pairs.
+ # { [dst -> data] -> [dst -> src] :.dst -> src is minimized distance }
+ minimal_pairs: isl.Map = (
+ fill_min_dist.apply_range(
+ # Note: Need to match fill -> min_dist with min_dist -> [fill -> match] as lexmin over
+ # fill and match will minimize distance over the tuple (src, dst, data), but that
+ # overconstrains the optimization as we want to minimize over distance (dst, data)
+ # only for all src.
+ fills_to_matches.range_map()
+ .apply_range(dist_fn)
+ .reverse()
+ )
+ .range()
+ .unwrap()
+ )
+ if DUMP_ISL_IR:
+ logging.info(f"minimal_pairs: {minimal_pairs}")
+
+ # Isolates the multicast networks.
+ # { [data] -> [dst -> src] : dst -> src is minimized distance }
+ multicast_networks: isl.Map = minimal_pairs.curry().range().unwrap()
+ # Devolves to a single source if multiple sources per domain point.
+ multicast_networks = multicast_networks.uncurry().lexmin().curry()
+
+ return multicast_networks
+
+
+def _covered_fills(mcs: isl.Map) -> isl.Map:
+ """
+ Recover the { dst -> data } fills actually covered by a matched multicast
+ source, from `identify_mesh_casts`'s result.
+
+ Parameters
+ ----------
+ mcs:
+ { [data] -> [dst -> src] }, the multicast-network map returned by
+ `identify_mesh_casts`.
+
+ Returns
+ -------
+ { dst -> data } -- every (destination, datum) pair that has a matched
+ source in `mcs`. A model's fill partition is then
+ `fulfilled = fills.map_.intersect(covered)` and
+ `unfulfilled = fills.map_.subtract(covered)`.
+ """
+ return (
+ mcs.range_reverse() # { data -> [src -> dst] }
+ .uncurry() # { [data -> src] -> dst }
+ .domain_factor_domain() # Drops src, keeps -> dst. { data -> dst }
+ .reverse() # { dst -> data }
+ )
+
+
+def _mesh_node_tuple(mcns: isl.Map) -> tuple[str, int]:
+ """
+ Read the spacetime/node tuple's name and dimensionality off a multicast
+ network map, instead of assuming a hardcoded name.
+
+ Parameters
+ ----------
+ mcns:
+ { [data] -> [dst -> src] }, the multicast-network map returned by
+ `identify_mesh_casts`. `dst` and `src` share the same node tuple, so
+ it suffices to read the name/dims off one side (`dst`, via
+ `.range().unwrap()`'s domain).
+
+ Returns
+ -------
+ `(name, dims)`: the node tuple's ISL tuple name (e.g. "noc" or "pe") and
+ its dimensionality.
+ """
+ # Note: correct even when mcns is empty -- ISL preserves space/tuple-name
+ # information on empty relations, so this does not require data to be
+ # present.
+ node_space: isl.Map = mcns.range().unwrap()
+ name: str = node_space.get_tuple_name(isl.dim_type.in_)
+ dims: int = node_space.dim(isl.dim_type.in_)
+ return name, dims
+
+
+def _per_src(mcns: isl.Map) -> tuple[isl.Map, isl.Map]:
+ """
+ Re-key multicast networks by tree `(data, src)`, plus a source lookup.
+
+ Parameters
+ ----------
+ mcns:
+ Multicast networks { [data] -> [dst -> src] } from `identify_mesh_casts`.
+
+ Returns
+ -------
+ `(per_src, keymap)`:
+ - `per_src`: { [data -> src] -> dst } -- every destination grouped under
+ its tree's key.
+ - `keymap`: { [data -> src] -> src } -- the tree's own source point,
+ recovered from `per_src`'s (wrapped) domain. Used to inject/carry the
+ source's own coordinates (e.g. its row/column) into computations keyed
+ on the same tree.
+ """
+ per_src: isl.Map = mcns.range_reverse().uncurry() # { [data -> src] -> dst }
+ keymap: isl.Map = per_src.domain().unwrap().range_map() # { [data -> src] -> src }
+ return per_src, keymap
+
+
+def _fabric_crossing(mcns: isl.Map, dist_fn: isl.Map) -> isl.Map:
+ """
+ Filter multicast networks down to fabric-crossing deliveries (>= 1 hop).
+
+ Parameters
+ ----------
+ mcns:
+ Multicast networks { [data] -> [dst -> src] } from `identify_mesh_casts`.
+ dist_fn:
+ The same distance function { [dst -> src] -> [hops] } passed to
+ `identify_mesh_casts` -- see its caller contract on the `hops` range
+ tuple name, which this filters on.
+
+ Returns
+ -------
+ { [data] -> [dst -> src] }, restricted to entries whose `[dst -> src]`
+ pair has `dist_fn(dst, src) >= 1`, i.e. every delivery that actually
+ crosses the fabric. Self-deliveries (0 hops) are dropped, since they load
+ no fabric link.
+ """
+ crossing_pairs: isl.Set = dist_fn.intersect_range(_CROSSING_HOPS).domain()
+ return mcns.intersect_range(crossing_pairs)
+
+
+def calculate_extents_per_dim(mcns: isl.Map) -> list[isl.PwAff]:
+ """
+ Parameters
+ ----------
+ mcns:
+ Mesh cast-networks, or networks in which all dsts per data are grouped with
+ the closest src containing the data.
+
+ Returns
+ -------
+ A list of `isl.PwAff` that gives the max extent (length) along dim_i per mcn,
+ where i is the i-th `isl.PwAff`.
+
+ Preconditions
+ -------------
+ `mcns` were generated with a Manhattan distance `dst_fn` by `identify_mesh_casts`
+ s.t. all dimensions are orthogonal to each other in a metric space, where each
+ unit movement in a dimension counts as 1 hop.
+
+ We also assume `dst_fn` is translationally invariant (i.e., ∀src, dst,
+ src', dst' ∈ space, if |src - dst| = |src' - dst'|,
+ dst_fn(src, dst) = dst_fn(src', dst').
+ """
+ # Makes mcns from { [data] -> [dst -> src] } to { [data -> src] -> [dst] }
+ potential_srcs, srcs = _per_src(mcns)
+ # Sources are part of the extents, so we union it with the destinations.
+ # { [data -> src] -> [spacetime] }
+ casting_extents: isl.Map = srcs.union(potential_srcs)
+
+ # Projects away all dimensions but one to find their extent for hypercube.
+ dims: int = potential_srcs.range_tuple_dim()
+ # Creates a mask of what to project out.
+ project_out_mask: list[bool] = [True] * dims
+ dim_extents: list[isl.PwAff] = [None] * dims
+
+ # Gets the extents of all dimensions
+ for noc_dim in range(dims):
+ # Project out all the dimensions of the output besides noc_dim.
+ project_out_mask[noc_dim] = False
+ # { [spacetime] -> [dimension] }
+ extent_mapper: isl.Map = dim_projector_mask(
+ casting_extents.range().get_space(), project_out_mask
+ ).reverse()
+ dim_extent_space: isl.Map = casting_extents.apply_range(extent_mapper)
+ project_out_mask[noc_dim] = True
+
+ # Finds max(noc_dim) - min(noc_dim) for each [data -> src]
+ max_extent: isl.PwAff = dim_extent_space.dim_max(0)
+ min_extent: isl.PwAff = dim_extent_space.dim_min(0)
+
+ # Subtracts the max from the min to get the extent per [data -> src]
+ dim_extents[noc_dim] = max_extent.sub(min_extent).coalesce()
+
+ return dim_extents
diff --git a/accelforge/model/_looptree/reuse/isl/mapping_to_isl/tiling.py b/accelforge/model/_looptree/reuse/isl/mapping_to_isl/tiling.py
index 7da28fed..122d4219 100644
--- a/accelforge/model/_looptree/reuse/isl/mapping_to_isl/tiling.py
+++ b/accelforge/model/_looptree/reuse/isl/mapping_to_isl/tiling.py
@@ -345,8 +345,9 @@ def consumer_based_tile_shape_inference(
# For each tensor read by this einsum, tile that tensor's producers.
for tensor in workload.einsums[einsum].input_tensor_names:
+ # The einsums that write `tensor` -- its producers in this workload.
producer_einsums: oset[EinsumName] = oset(
- [e for e in workload.einsums[einsum].output_tensor_names]
+ [e.name for e in workload.einsums if tensor in e.output_tensor_names]
)
if len(producer_einsums) > 1:
raise NotImplementedError(
diff --git a/accelforge/model/_looptree/reuse/isl/spatial.py b/accelforge/model/_looptree/reuse/isl/spatial.py
index 47d6fa42..7c6ad5a5 100644
--- a/accelforge/model/_looptree/reuse/isl/spatial.py
+++ b/accelforge/model/_looptree/reuse/isl/spatial.py
@@ -9,6 +9,9 @@
import islpy as isl
from accelforge.frontend.mapping import MappingNode
+from accelforge.model._looptree.reuse.isl.distributed.edge_pressure import (
+ EdgePressure,
+)
from accelforge.model._looptree.reuse.isl.isl_functions import (
insert_equal_dims_map,
reorder_projector,
@@ -37,9 +40,17 @@ class TransferInfo:
# Crucial information to transfer info.
fulfilled_fill: Transfers
- """Fills done by peer-to-peer transfers."""
+ """Fills done by peer-to-peer transfers.
+
+ Restricted to fills covered by a matched multicast source (see
+ `identify_mesh_casts`); see `unfulfilled_fill`.
+ """
unfulfilled_fill: Fill
- """Fills not performed."""
+ """Fills not performed.
+
+ `fills - covered`, the complement of `fulfilled_fill` within the original
+ fill set (the two partition `fills` exactly).
+ """
parent_reads: Reads
"""Fills done by parent-to-child transfers."""
hops: isl.PwQPolynomial
@@ -48,6 +59,18 @@ class TransferInfo:
# Metadata on what is occurring.
link_transfer: bool
+ # Optional, model-specific per-link decomposition -- last field / defaulted
+ # so this addition does not disturb any existing positional or keyword
+ # construction site (`grep`-verified: every `TransferInfo(...)` call in the
+ # tree already uses keyword arguments).
+ edge_pressure: Optional[EdgePressure] = None
+ """Per-directed-edge load backing `hops`, for models that define one.
+
+ Populated by `XYRoutingMulticastModel` (mesh links) and
+ `StarMulticastModel` (spokes); `None` where a model has no per-link
+ decomposition (`HypercubeMulticastModel`, `FullyConnectedMulticastModel`).
+ """
+
class TransferModel(ABC):
"""
diff --git a/notebooks/astrasim2_correlation/correlation.ipynb b/notebooks/astrasim2_correlation/correlation.ipynb
new file mode 100644
index 00000000..014fd35f
--- /dev/null
+++ b/notebooks/astrasim2_correlation/correlation.ipynb
@@ -0,0 +1,680 @@
+{
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "id": "3aa122f8",
+ "metadata": {},
+ "source": [
+ "# Network Sim Correlation\n",
+ "\n",
+ "The following are attempts to correlate the network model in accelforge to actual networks measured in real architectures."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 1,
+ "id": "e3a04b7c",
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-06-11T18:41:47.222682Z",
+ "iopub.status.busy": "2026-06-11T18:41:47.222573Z",
+ "iopub.status.idle": "2026-06-11T18:41:48.167242Z",
+ "shell.execute_reply": "2026-06-11T18:41:48.166287Z"
+ }
+ },
+ "outputs": [],
+ "source": [
+ "import accelforge as af\n",
+ "import matplotlib.pyplot as plt"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "95cf92d9",
+ "metadata": {},
+ "source": [
+ "## We are testing an 8 GPU All-to-All to Correlate Later, simulating an NVLink Switch\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "195ecb91",
+ "metadata": {},
+ "source": [
+ "### Encoding: one-hot GPU coordinates on a fully-connected fabric\n",
+ "\n",
+ "GPU $i$ sits at one-hot coordinate $e_i$ of an $N$-dimensional `noc` space. `data[s, d]` is the chunk sent by GPU $s$ to GPU $d$: each GPU *holds* the chunks it sends (occupancy) and *requests* the chunks addressed to it (fill).\n",
+ "\n",
+ "Every $s \\neq d$ cast then has extent 1 along exactly the src and dst dimensions, so `HypercubeMulticastModel`'s bounding-box cost is $(1+1)(1+1)-1 = 3$ **uniformly for every pair** — the fully-connected property (no pair privileged). The distance function is unit-cost (0 if same GPU, 1 otherwise) and only influences source matching in `identify_mesh_casts`. Self-chunks match at distance 0 and cost 0 hops: they never cross the fabric."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 2,
+ "id": "c5b2c7bf",
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-06-11T18:41:48.169250Z",
+ "iopub.status.busy": "2026-06-11T18:41:48.169045Z",
+ "iopub.status.idle": "2026-06-11T18:41:48.177262Z",
+ "shell.execute_reply": "2026-06-11T18:41:48.176277Z"
+ }
+ },
+ "outputs": [],
+ "source": [
+ "import math\n",
+ "\n",
+ "import islpy as isl\n",
+ "import pandas as pd\n",
+ "\n",
+ "from accelforge.model._looptree.reuse.isl.distributed.distributed_buffers import (\n",
+ " HypercubeMulticastModel,\n",
+ ")\n",
+ "from accelforge.model._looptree.reuse.isl.mapping_to_isl.types import (\n",
+ " Fill,\n",
+ " Occupancy,\n",
+ " SpatialTag,\n",
+ ")\n",
+ "\n",
+ "CTX = isl.DEFAULT_CONTEXT"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 3,
+ "id": "e343cac3",
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-06-11T18:41:48.178504Z",
+ "iopub.status.busy": "2026-06-11T18:41:48.178391Z",
+ "iopub.status.idle": "2026-06-11T18:41:48.183047Z",
+ "shell.execute_reply": "2026-06-11T18:41:48.182522Z"
+ }
+ },
+ "outputs": [],
+ "source": [
+ "def onehot_constraints(prefix: str, n: int) -> str:\n",
+ " \"\"\"One-hot constraints over dims ``{prefix}0..{prefix}{n-1}``.\"\"\"\n",
+ " bounds = \" and \".join(f\"0 <= {prefix}{i} <= 1\" for i in range(n))\n",
+ " hot = \" + \".join(f\"{prefix}{i}\" for i in range(n)) + \" = 1\"\n",
+ " return f\"{bounds} and {hot}\"\n",
+ "\n",
+ "\n",
+ "def linear_id(prefix: str, n: int) -> str:\n",
+ " \"\"\"Affine recovery of the GPU id from a one-hot vector: id = sum i*g_i.\"\"\"\n",
+ " return \" + \".join(f\"{i}*{prefix}{i}\" for i in range(1, n))\n",
+ "\n",
+ "\n",
+ "def all_to_all_maps(n: int) -> tuple[isl.Map, isl.Map, isl.Map]:\n",
+ " \"\"\"Build (occupancy, fill, dist_fn) for an N-GPU fully-connected all-to-all.\n",
+ "\n",
+ " data[s, d] is the chunk sent by GPU s to GPU d. Each GPU holds the\n",
+ " chunks it sends (occ) and requests the chunks addressed to it (fill).\n",
+ " \"\"\"\n",
+ " gs = \", \".join(f\"gs{i}\" for i in range(n))\n",
+ " gd = \", \".join(f\"gd{i}\" for i in range(n))\n",
+ " occ = isl.Map.read_from_str(\n",
+ " CTX,\n",
+ " f\"{{ noc[{gs}] -> data[s, d] : {onehot_constraints('gs', n)} \"\n",
+ " f\"and s = {linear_id('gs', n)} and 0 <= d < {n} }}\",\n",
+ " )\n",
+ " fill = isl.Map.read_from_str(\n",
+ " CTX,\n",
+ " f\"{{ noc[{gd}] -> data[s, d] : {onehot_constraints('gd', n)} \"\n",
+ " f\"and d = {linear_id('gd', n)} and 0 <= s < {n} }}\",\n",
+ " )\n",
+ "\n",
+ " xd = \", \".join(f\"xd{i}\" for i in range(n))\n",
+ " xs = \", \".join(f\"xs{i}\" for i in range(n))\n",
+ " same = \" and \".join(f\"xd{i} = xs{i}\" for i in range(n))\n",
+ " diff = \" or \".join(f\"(xd{i} < xs{i}) or (xd{i} > xs{i})\" for i in range(n))\n",
+ " dist_fn = isl.Map.read_from_str(\n",
+ " CTX,\n",
+ " f\"{{ [noc[{xd}] -> noc[{xs}]] -> hops[0] : {same}; \"\n",
+ " f\" [noc[{xd}] -> noc[{xs}]] -> hops[1] : {diff} }}\",\n",
+ " )\n",
+ " return occ, fill, dist_fn"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "a2a2df5f",
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-06-11T18:41:48.184434Z",
+ "iopub.status.busy": "2026-06-11T18:41:48.184302Z",
+ "iopub.status.idle": "2026-06-11T18:41:48.188107Z",
+ "shell.execute_reply": "2026-06-11T18:41:48.187196Z"
+ }
+ },
+ "outputs": [],
+ "source": [
+ "def eval_total(poly: isl.PwQPolynomial) -> int:\n",
+ " \"\"\"Evaluate a fully-summed PwQPolynomial (zero-dimensional domain).\"\"\"\n",
+ " return int(poly.eval(isl.Point.zero(poly.domain().get_space())).to_python())\n",
+ "\n",
+ "\n",
+ "def model_hops_per_element(n: int) -> int:\n",
+ " \"\"\"Run the real HypercubeMulticastModel on the all-to-all at V=1.\"\"\"\n",
+ " occ_map, fill_map, dist_fn = all_to_all_maps(n)\n",
+ " tags = [SpatialTag(i, 0) for i in range(n)]\n",
+ " model = HypercubeMulticastModel(dist_fn)\n",
+ " info = model.apply(0, Fill(tags, fill_map), Occupancy(tags, occ_map))\n",
+ " return eval_total(info.hops)\n",
+ "\n",
+ "\n",
+ "def single_chunk_hops(n: int, src: int, dst: int) -> int:\n",
+ " \"\"\"Hops for one (src, dst) chunk — uniformity / self-chunk probe.\"\"\"\n",
+ " occ_map, fill_map, dist_fn = all_to_all_maps(n)\n",
+ " chunk = isl.Set.read_from_str(CTX, f\"{{ data[{src}, {dst}] }}\")\n",
+ " tags = [SpatialTag(i, 0) for i in range(n)]\n",
+ " model = HypercubeMulticastModel(dist_fn)\n",
+ " info = model.apply(\n",
+ " 0,\n",
+ " Fill(tags, fill_map.intersect_range(chunk)),\n",
+ " Occupancy(tags, occ_map.intersect_range(chunk)),\n",
+ " )\n",
+ " return eval_total(info.hops)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 5,
+ "id": "06d40ac3",
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-06-11T18:41:48.189890Z",
+ "iopub.status.busy": "2026-06-11T18:41:48.189717Z",
+ "iopub.status.idle": "2026-06-11T18:41:48.193103Z",
+ "shell.execute_reply": "2026-06-11T18:41:48.192370Z"
+ }
+ },
+ "outputs": [],
+ "source": [
+ "NODES = 8 # GPUs on the fabric\n",
+ "LINK_BW_GBPS = 150.0 # per-GPU per-direction NVLink bandwidth, GB/s (V100 generation)\n",
+ "ALPHA_S = 0.0 # per-operation latency overhead, seconds\n",
+ "MIN_MIB, MAX_MIB = 1, 1024 # collective size sweep bounds, powers of 2\n",
+ "PER_RANK = False # interpret swept sizes as per-rank NCCL sizes instead of totals"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 6,
+ "id": "60364739",
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-06-11T18:41:48.194288Z",
+ "iopub.status.busy": "2026-06-11T18:41:48.194178Z",
+ "iopub.status.idle": "2026-06-11T18:41:49.888588Z",
+ "shell.execute_reply": "2026-06-11T18:41:49.888067Z"
+ }
+ },
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "ISL network model: 8-GPU fully-connected all-to-all (one-hot)\n",
+ " hops/element : 168 (uniform 3 per pair)\n",
+ " fabric-crossing chunks: 56 of 64\n"
+ ]
+ }
+ ],
+ "source": [
+ "n = NODES\n",
+ "\n",
+ "# --- Run the tool at V=1 and validate the fully-connected structure. ---\n",
+ "hops_per_elem = model_hops_per_element(n)\n",
+ "expected = 3 * n * (n - 1)\n",
+ "assert hops_per_elem == expected, (\n",
+ " f\"model returned {hops_per_elem} hops/element, expected {expected} \"\n",
+ " f\"(= 3 per off-diagonal pair x {n*(n-1)} pairs)\"\n",
+ ")\n",
+ "probe_off = single_chunk_hops(n, 0, min(3, n - 1))\n",
+ "probe_self = single_chunk_hops(n, n - 1, n - 1)\n",
+ "assert probe_off == 3, f\"off-diagonal chunk cost {probe_off} != 3\"\n",
+ "assert probe_self == 0, f\"self chunk cost {probe_self} != 0\"\n",
+ "crossing_chunks = hops_per_elem // 3 # = N(N-1) fabric-crossing chunks\n",
+ "\n",
+ "print(f\"ISL network model: {n}-GPU fully-connected all-to-all (one-hot)\")\n",
+ "print(f\" hops/element : {hops_per_elem} (uniform 3 per pair)\")\n",
+ "print(f\" fabric-crossing chunks: {crossing_chunks} of {n*n}\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "c89a8e17",
+ "metadata": {},
+ "source": [
+ "### Size and latency conventions\n",
+ "\n",
+ "`collective_size` $S$ = **total** bytes moved by the collective across all ranks (ASTRA-sim style). Per-rank NCCL buffer size = $S/N$; each $(s, d)$ pair exchanges a chunk of $S/N^2$ bytes. Set `PER_RANK = True` to interpret the swept sizes as per-rank NCCL sizes instead.\n",
+ "\n",
+ "Latency derives from the model's hop count: hops are uniform 3 per fabric-crossing chunk, so crossing chunks/element $= \\mathrm{hops}/3 = N(N-1)$, i.e. each GPU receives $N-1$ chunks $\\Rightarrow$ per-GPU wire bytes $= S(N-1)/N^2$. With a non-blocking switch and full-duplex ports (send overlaps receive), port drain time gives\n",
+ "\n",
+ "$$ t = \\alpha + \\frac{S\\,(N-1)}{N^2 \\cdot BW} $$\n",
+ "\n",
+ "NCCL conventions: $\\mathrm{algbw} = (S/N)/t$, $\\mathrm{busbw} = \\mathrm{algbw} \\cdot (N-1)/N$. With $\\alpha = 0$, busbw $\\equiv$ BW exactly — used as a per-row sanity invariant below.\n",
+ "\n",
+ "**Correlation notes** — this is a flat bandwidth model with no launch/protocol overhead, so expect empirical underprediction at small sizes; calibrate `ALPHA_S` (per-operation latency) and `LINK_BW_GBPS` (effective port bandwidth) against nccl-tests `alltoall_perf`, and overlay the ASTRA-sim 2.0 (FullyConnected) and EC2 series on the plot below."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "6d5e39c9",
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-06-11T18:41:49.890046Z",
+ "iopub.status.busy": "2026-06-11T18:41:49.889928Z",
+ "iopub.status.idle": "2026-06-11T18:41:49.909301Z",
+ "shell.execute_reply": "2026-06-11T18:41:49.908706Z"
+ }
+ },
+ "outputs": [
+ {
+ "data": {
+ "text/html": [
+ "
\n",
+ "\n",
+ "
\n",
+ " \n",
+ "
\n",
+ "
\n",
+ "
nodes
\n",
+ "
collective_size_bytes
\n",
+ "
per_rank_bytes
\n",
+ "
per_gpu_wire_bytes
\n",
+ "
hops_per_element
\n",
+ "
fabric_chunks_per_element
\n",
+ "
model_latency_s
\n",
+ "
algbw_GBps
\n",
+ "
busbw_GBps
\n",
+ "
\n",
+ " \n",
+ " \n",
+ "
\n",
+ "
0
\n",
+ "
8
\n",
+ "
1048576
\n",
+ "
131072
\n",
+ "
114688
\n",
+ "
168
\n",
+ "
56
\n",
+ "
7.645867e-07
\n",
+ "
171.428571
\n",
+ "
150.0
\n",
+ "
\n",
+ "
\n",
+ "
1
\n",
+ "
8
\n",
+ "
2097152
\n",
+ "
262144
\n",
+ "
229376
\n",
+ "
168
\n",
+ "
56
\n",
+ "
1.529173e-06
\n",
+ "
171.428571
\n",
+ "
150.0
\n",
+ "
\n",
+ "
\n",
+ "
2
\n",
+ "
8
\n",
+ "
4194304
\n",
+ "
524288
\n",
+ "
458752
\n",
+ "
168
\n",
+ "
56
\n",
+ "
3.058347e-06
\n",
+ "
171.428571
\n",
+ "
150.0
\n",
+ "
\n",
+ "
\n",
+ "
3
\n",
+ "
8
\n",
+ "
8388608
\n",
+ "
1048576
\n",
+ "
917504
\n",
+ "
168
\n",
+ "
56
\n",
+ "
6.116693e-06
\n",
+ "
171.428571
\n",
+ "
150.0
\n",
+ "
\n",
+ "
\n",
+ "
4
\n",
+ "
8
\n",
+ "
16777216
\n",
+ "
2097152
\n",
+ "
1835008
\n",
+ "
168
\n",
+ "
56
\n",
+ "
1.223339e-05
\n",
+ "
171.428571
\n",
+ "
150.0
\n",
+ "
\n",
+ "
\n",
+ "
5
\n",
+ "
8
\n",
+ "
33554432
\n",
+ "
4194304
\n",
+ "
3670016
\n",
+ "
168
\n",
+ "
56
\n",
+ "
2.446677e-05
\n",
+ "
171.428571
\n",
+ "
150.0
\n",
+ "
\n",
+ "
\n",
+ "
6
\n",
+ "
8
\n",
+ "
67108864
\n",
+ "
8388608
\n",
+ "
7340032
\n",
+ "
168
\n",
+ "
56
\n",
+ "
4.893355e-05
\n",
+ "
171.428571
\n",
+ "
150.0
\n",
+ "
\n",
+ "
\n",
+ "
7
\n",
+ "
8
\n",
+ "
134217728
\n",
+ "
16777216
\n",
+ "
14680064
\n",
+ "
168
\n",
+ "
56
\n",
+ "
9.786709e-05
\n",
+ "
171.428571
\n",
+ "
150.0
\n",
+ "
\n",
+ "
\n",
+ "
8
\n",
+ "
8
\n",
+ "
268435456
\n",
+ "
33554432
\n",
+ "
29360128
\n",
+ "
168
\n",
+ "
56
\n",
+ "
1.957342e-04
\n",
+ "
171.428571
\n",
+ "
150.0
\n",
+ "
\n",
+ "
\n",
+ "
9
\n",
+ "
8
\n",
+ "
536870912
\n",
+ "
67108864
\n",
+ "
58720256
\n",
+ "
168
\n",
+ "
56
\n",
+ "
3.914684e-04
\n",
+ "
171.428571
\n",
+ "
150.0
\n",
+ "
\n",
+ "
\n",
+ "
10
\n",
+ "
8
\n",
+ "
1073741824
\n",
+ "
134217728
\n",
+ "
117440512
\n",
+ "
168
\n",
+ "
56
\n",
+ "
7.829367e-04
\n",
+ "
171.428571
\n",
+ "
150.0
\n",
+ "
\n",
+ " \n",
+ "
\n",
+ "
"
+ ],
+ "text/plain": [
+ " nodes collective_size_bytes per_rank_bytes per_gpu_wire_bytes \\\n",
+ "0 8 1048576 131072 114688 \n",
+ "1 8 2097152 262144 229376 \n",
+ "2 8 4194304 524288 458752 \n",
+ "3 8 8388608 1048576 917504 \n",
+ "4 8 16777216 2097152 1835008 \n",
+ "5 8 33554432 4194304 3670016 \n",
+ "6 8 67108864 8388608 7340032 \n",
+ "7 8 134217728 16777216 14680064 \n",
+ "8 8 268435456 33554432 29360128 \n",
+ "9 8 536870912 67108864 58720256 \n",
+ "10 8 1073741824 134217728 117440512 \n",
+ "\n",
+ " hops_per_element fabric_chunks_per_element model_latency_s algbw_GBps \\\n",
+ "0 168 56 7.645867e-07 171.428571 \n",
+ "1 168 56 1.529173e-06 171.428571 \n",
+ "2 168 56 3.058347e-06 171.428571 \n",
+ "3 168 56 6.116693e-06 171.428571 \n",
+ "4 168 56 1.223339e-05 171.428571 \n",
+ "5 168 56 2.446677e-05 171.428571 \n",
+ "6 168 56 4.893355e-05 171.428571 \n",
+ "7 168 56 9.786709e-05 171.428571 \n",
+ "8 168 56 1.957342e-04 171.428571 \n",
+ "9 168 56 3.914684e-04 171.428571 \n",
+ "10 168 56 7.829367e-04 171.428571 \n",
+ "\n",
+ " busbw_GBps \n",
+ "0 150.0 \n",
+ "1 150.0 \n",
+ "2 150.0 \n",
+ "3 150.0 \n",
+ "4 150.0 \n",
+ "5 150.0 \n",
+ "6 150.0 \n",
+ "7 150.0 \n",
+ "8 150.0 \n",
+ "9 150.0 \n",
+ "10 150.0 "
+ ]
+ },
+ "execution_count": 7,
+ "metadata": {},
+ "output_type": "execute_result"
+ }
+ ],
+ "source": [
+ "bw = LINK_BW_GBPS * 1e9 # bytes/s\n",
+ "\n",
+ "rows = []\n",
+ "mib = MIN_MIB\n",
+ "while mib <= MAX_MIB:\n",
+ " size = mib * (1 << 20)\n",
+ " total = size * n if PER_RANK else size\n",
+ " per_rank = total / n\n",
+ " chunk_bytes = total / (n * n)\n",
+ " # Latency derived from the model output: hops/3 crossing chunks\n",
+ " # spread evenly, (N-1) into each GPU's port.\n",
+ " per_gpu_wire = (crossing_chunks / n) * chunk_bytes\n",
+ " t = ALPHA_S + per_gpu_wire / bw\n",
+ " # Closed form cross-check: t = alpha + S(N-1)/(N^2 BW)\n",
+ " assert math.isclose(t, ALPHA_S + total * (n - 1) / (n * n * bw), rel_tol=1e-12)\n",
+ " algbw = per_rank / t / 1e9\n",
+ " busbw = algbw * (n - 1) / n\n",
+ " rows.append((n, int(total), int(per_rank), int(per_gpu_wire),\n",
+ " hops_per_elem, crossing_chunks, t, algbw, busbw))\n",
+ " mib *= 2\n",
+ "\n",
+ "sweep = pd.DataFrame(rows, columns=[\n",
+ " \"nodes\", \"collective_size_bytes\", \"per_rank_bytes\", \"per_gpu_wire_bytes\",\n",
+ " \"hops_per_tile_coordinate\", \"fabric_chunks_per_tile_coordinate\",\n",
+ " \"model_latency_s\", \"algbw_GBps\", \"busbw_GBps\",\n",
+ "])\n",
+ "sweep\n",
+ "\n",
+ "# TODO: Needs to account for links sharing the same \\\n",
+ "# Currently: doing an ALL GATHER, not all to all.\n",
+ "# All to all: reshufflling tile shards to be on 1 gpu.\n",
+ "# Hops per element * per rank byte (document somewhere)\n",
+ "\n",
+ "# NOTE: Can think of nvlink as one hop. 0,0 as router. Every time you want to send something, have to go through 0,0.\n",
+ "# NOTE: How do we route for star routing? Instead of figuring out all routings, should we have a set of routing choices?\n",
+ "# NOTE: We support various different topologie. However, you have to specify the routing chosen for multicast?\n",
+ "# NOTE: It is a simple multicast.\n",
+ " # What if that's the contribution? For regular workloads, superimposing multicast trees would balance traffic.\n",
+ " # Don't need fancy things (cycle accurate simulation), and still account for congestion (bw pressure on links).\n",
+ " # Hard to frame contribution in terms of coordinates"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 8,
+ "id": "cbdaefb9",
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-06-11T18:41:49.910815Z",
+ "iopub.status.busy": "2026-06-11T18:41:49.910686Z",
+ "iopub.status.idle": "2026-06-11T18:41:49.913503Z",
+ "shell.execute_reply": "2026-06-11T18:41:49.913097Z"
+ }
+ },
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "busbw invariant holds: busbw == 150 GB/s on every row\n"
+ ]
+ }
+ ],
+ "source": [
+ "if ALPHA_S == 0:\n",
+ " assert (sweep[\"busbw_GBps\"] - LINK_BW_GBPS).abs().max() < 1e-9 * LINK_BW_GBPS\n",
+ " print(f\"busbw invariant holds: busbw == {LINK_BW_GBPS:g} GB/s on every row\")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 9,
+ "id": "11751aab",
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-06-11T18:41:49.914710Z",
+ "iopub.status.busy": "2026-06-11T18:41:49.914592Z",
+ "iopub.status.idle": "2026-06-11T18:41:50.198575Z",
+ "shell.execute_reply": "2026-06-11T18:41:50.197933Z"
+ }
+ },
+ "outputs": [
+ {
+ "data": {
+ "image/png": "iVBORw0KGgoAAAANSUhEUgAAAmkAAAHbCAYAAACQmw0xAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjksIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvJkbTWQAAAAlwSFlzAAAPYQAAD2EBqD+naQAAmoVJREFUeJzs3XlcVFX/B/DPzLArDLKjIpCYgiiYuOWGiqKZZu5tmtoq5kKb/ipxqSzNtYcy60lLK81My7Xcl9RcccNdXFKBUWRfZ+b+/qCZh3FYZoM7A5/36+Wr5t5z7jl35svly517zpEIgiCAiIiIiKyKVOwOEBEREZE+JmlEREREVohJGhEREZEVYpJGREREZIWYpBERERFZISZpRERERFaISRoRERGRFWKSRkRERGSFmKQRERERWSEmaWRRSqUS77zzDgICAiCVSjFo0CCj6s+YMQMSiURnW1BQEF588UXLdZIsorzPylqZ01dN3Xv37lm4V/+zZ88eSCQS7NmzR7vtxRdfRFBQULW1aUnm9l+tViM8PBwfffRR9XSQ9ERHRyM8PFzsbphkxYoVkEgkuH79utF1H74WlJSUICAgAF988YUFe2g5TNJq0OXLlzFy5Eg0btwYLi4uaNGiBWbNmoX8/HyjjpOSkoIJEybg0UcfhYuLC1xcXBAWFoa4uDicPn1ap6wmIDX/NGXff/99ZGdn65Wr6BdReHg4oqOjq+zbt99+i3nz5mHo0KH47rvvMGXKFKPOjSzrzp07mDFjBpKSkkTrw4svvgiJRILWrVujvFXoJBIJJkyYAABYsGABJBIJduzYUeHxvv76a0gkEvz+++8AbPuXjSVYw2dsrp9++gm3bt3SxgHwv1/Ex44d0yl74MAB9OvXD40aNYKTkxOaNGmCAQMG4Mcff9QpVzauxPbFF19gxYoVYneDymFvb4/4+Hh89NFHKCwsFLs7epik1ZBbt26hffv2OHz4MCZMmIBFixahU6dOSEhIwDPPPGPwcTZt2oTw8HCsXLkSMTExWLhwIRYvXox+/fphy5YtiIyMxI0bN/Tqffnll1i5ciUWLFiAFi1a4KOPPkLfvn3L/aVpjl27dqFRo0ZYuHAhXnjhBXTv3t2ixyfj3LlzBzNnzrSKX+BnzpzBr7/+WmmZkSNHQiqV6v3CLevHH3+Ep6cn+vXrZ3Db77//PgoKCgwub0us6TM21bx58zBy5EjI5fJKy61duxbdunVDWloaJk2ahM8//xzPP/88Hjx4gK+//rqGems8JmnWbcyYMbh3716l1x2x2Indgbpi5cqVyMzMxIEDB9CyZUsAwCuvvAK1Wo3vv/8eDx48QIMGDSo9xtWrVzFy5EgEBgZi586d8Pf319n/6aef4osvvoBUqp97Dx06FF5eXgCA1157DUOGDMGvv/6Kw4cPo1OnThY6SyA9PR3u7u4WOx7VDs7OzggICMCsWbMwePDgCr96bNiwIXr06IFff/0VX375JRwdHXX23759G/v27cMrr7wCe3t7g9u3s7ODnR0vd9bo5MmTOHXqFObPn19l2RkzZiAsLAyHDx+Gg4ODzr709PTq6mKtkpeXh3r16ondDavi7u6OPn36YMWKFRg7dqzY3dHBO2k1RPPVoq+vr852f39/SKVSvQtOeebOnYu8vDwsX75cL0EDSn8RTZw4EQEBAVUeq2fPngBKvzq1hOvXr0MikWD37t04d+6c9uvVPXv2lPu8Stk6xvyFee3aNUgkEixcuFBv38GDByGRSPDTTz9VeZytW7eie/fucHV1hZubG9q1a6f3V9TatWvRtm1bODs7w8vLC88//zxu376tU+bFF19E/fr1cfv2bQwaNAj169eHt7c33nrrLahUKr1z/eyzz7Bs2TI0bdoUjo6OaNeuHY4eParXvwsXLmDo0KHw8PCAk5MToqKitF/vlZWZmYkpU6YgKCgIjo6OaNy4MUaNGoV79+5hz549aNeuHYDSvxQ1n0nZ9/vvv/9G3759IZfL4eLigu7du+Ovv/7Sa+fAgQNo164dnJyc0LRpU3z11VdVvsdlSaVSvP/++zh9+jTWr19fadnnn38eWVlZ2Lx5s96+1atXQ61W47nnnjOq/fKeSdN8HbZhwwaEh4fD0dERLVu2xLZt26o83o0bNxASEoLw8HCkpaVVWm78+PFo3rw5nJ2d4enpiWHDhpn0LE15DPmMDYljsfoPABs2bICDgwO6detWZdmrV6+iXbt25V4vfXx8LNIfQ3+mgdJn6RYtWoSWLVvCyckJvr6+ePXVV/HgwQNtmaCgIJw7dw579+7Vfj7R0dHIzMyETCbDkiVLtGXv3bsHqVQKT09PnW85Xn/9dfj5+em0bcz16erVq3jiiSfg6upa6c/On3/+CRcXFzzzzDNQKpUVltM8YnD69Gl0794dLi4uCAkJwS+//AIA2Lt3Lzp06ABnZ2c0b9683McXTp48iX79+sHNzQ3169dHr169cPjwYb1y586dQ8+ePeHs7IzGjRvjww8/hFqtLrdfW7duRdeuXVGvXj24urqif//+OHfuXIXnUVbv3r1x4MABZGRkGFS+pjBJqyGa57nGjRuHpKQk3Lp1C2vWrMGXX36JiRMnGvSXzaZNmxASEoIOHTqY3Z+rV68CADw9Pc0+FgB4e3tj5cqVaNGiBRo3boyVK1di5cqVCA0NtcjxNR555BF07twZP/zwg96+H374Aa6urnjqqacqPcaKFSvQv39/ZGRkYNq0afjkk08QGRmp88t5xYoVGD58OGQyGebMmYOXX34Zv/76K7p06YLMzEyd46lUKsTGxsLT0xOfffYZunfvjvnz52PZsmV6bf/444+YN28eXn31VXz44Ye4fv06Bg8ejJKSEm2Zc+fOoWPHjjh//jymTp2K+fPno169ehg0aJBOgpObm4uuXbvi888/R58+fbB48WK89tpruHDhAv755x+EhoZi1qxZAErv2mo+E80vw127dqFbt27Izs5GQkICPv74Y2RmZqJnz544cuSItp0zZ86gT58+SE9Px4wZMzBmzBgkJCRUmWw97Nlnn0WzZs0wa9asSr9mHzx4MJycnMr96uHHH39EYGAgOnfubFTbFTlw4ADGjx+PkSNHYu7cuSgsLMSQIUNw//79CutcvXoV3bp1g6urK/bs2aP3h1dZR48excGDBzFy5EgsWbIEr732Gnbu3Ino6Gijn0UtT1WfsTFxLEb/gdI/rsLDww26M6r5FuGff/6xSNsVMfRn+tVXX8Xbb7+Nzp07Y/HixRgzZgx++OEHxMbGan+mFy1ahMaNG6NFixbaz+e9996Du7s7wsPDsW/fPu3xDhw4AIlEgoyMDCQnJ2u379+/H127dtW+NuZzVSqViI2NhY+PDz777DMMGTKk3HPetGkTBg4ciGHDhmHVqlVV3nl+8OABnnzySXTo0AFz586Fo6MjRo4ciTVr1mDkyJF44okn8MknnyAvLw9Dhw5FTk6Otu65c+fQtWtXnDp1Cu+88w4++OADpKSkIDo6Gn///be2XGpqKnr06IGkpCRMnToVkydPxvfff4/Fixfr9WflypXo378/6tevj08//RQffPABkpOT0aVLF4P+qGjbti0EQcDBgwerLFujBKoxs2fPFpydnQUA2n/vvfeeQXWzsrIEAMKgQYP09j148EBQKBTaf/n5+dp9CQkJAgDh4sWLgkKhEFJSUoSvvvpKcHR0FHx9fYW8vDydcgqFotz2W7ZsKXTv3r3Kfnbv3l1o2bKlzrbdu3cLAITdu3frbE9JSREACMuXL9frb1mBgYHC6NGjta+/+uorAYBw/vx57bbi4mLBy8tLp1x5MjMzBVdXV6FDhw5CQUGBzj61Wq09lo+PjxAeHq5TZtOmTQIAYfr06dpto0ePFgAIs2bN0jlWmzZthLZt2+qdq6enp5CRkaHd/ttvvwkAhI0bN2q39erVS2jVqpVQWFio07fHH39caNasmXbb9OnTBQDCr7/+qneemnM5evSo3nus2d+sWTMhNjZWW1YQBCE/P18IDg4Wevfurd02aNAgwcnJSbhx44Z2W3JysiCTyfQ+q/KMHj1aqFevniAIgvDdd9/p9RmAEBcXp1Nn2LBhgpOTk5CVlaXdduHCBQGAMG3aNJ2y5cXcw8qLKwCCg4ODcOXKFe22U6dOCQCEzz//XK+uQqEQzp8/LzRs2FBo166dzudYkbI/ixqHDh0SAAjff/+9dlt5PyOjR48WAgMDq2yjos/YmDgWs/+NGzcWhgwZord9+fLlAgDh6NGj2m3//e9/tZ9bjx49hA8++EDYv3+/oFKp9OqXF1eGMPRnev/+/QIA4YcfftApt23bNr3tFV0/4+LiBF9fX+3r+Ph4oVu3boKPj4/w5ZdfCoIgCPfv3xckEomwePFiQRBMuz5NnTpVr+2yPzfr1q0T7O3thZdffrnc97K8ugCEH3/8UbtN8/MplUqFw4cPa7f/8ccfevE5aNAgwcHBQbh69ap22507dwRXV1ehW7du2m2TJ08WAAh///23dlt6erogl8sFAEJKSoogCIKQk5MjuLu7Cy+//LJOP1NTUwW5XK6zvbxrgaZ9AMKnn35a5fnXJN5Jq0FBQUHo1q0bli1bhnXr1mHs2LH4+OOP8Z///KfKupqvS+vXr6+3Lzo6Gt7e3tp/iYmJemWaN28Ob29vBAcH49VXX0VISAg2b94MFxcX80+shg0fPhxOTk46d9P++OMP3Lt3D88//3yldbdv346cnBxMnToVTk5OOvs0X4cdO3YM6enpGD9+vE6Z/v37o0WLFuV+Dffaa6/pvO7atSuuXbumV27EiBE6zx5q/jrWlM3IyMCuXbswfPhw5OTk4N69e7h37x7u37+P2NhYXL58WfuVxrp16xAREYGnn35ar52qpptISkrC5cuX8eyzz+L+/fvadvLy8tCrVy/s27cParUaKpUKf/zxBwYNGoQmTZpo64eGhiI2NrbSNsrz3HPPGXQ37fnnn0dhYaHOQAPNnTVjv+qsTExMDJo2bap93bp1a7i5uZX72Z09exbdu3dHUFAQduzYUeUzpEDps3gaJSUluH//PkJCQuDu7o4TJ05Y5iQqYEocP6wm+n///n2D3ksAGDt2LLZt24bo6GgcOHAAs2fPRteuXdGsWTOL3wGp6md67dq1kMvl6N27t/bn5969e2jbti3q16+P3bt3V9lG165dkZaWhosXLwIovWPWrVs3dO3aFfv37wdQendNEATttcKUz/X111+vsA8//fQTRowYgVdffRVfffVVuc80l6d+/foYOXKk9nXz5s3h7u6O0NBQnW97NP+vee9UKhX+/PNPDBo0CI888oi2nL+/P5599lkcOHBA+/tuy5Yt6NixI9q3b68t5+3trXcN2L59OzIzM/HMM8/ofBYymQwdOnQw6LPQxGB1TrVjCj5JW0NWr16NV155BZcuXULjxo0BlH6to1ar8e677+KZZ56Bp6cnMjIyUFxcrK3n7OwMuVwOV1dXAKVfcT3sq6++Qk5ODtLS0ipMUtatWwc3NzfY29ujcePGOr+YDGUtc2K5u7trh9zPnj0bQOlXnY0aNdI+a5ebm6vzXslkMnh7e2u/5q1sygbN6NjmzZvr7WvRogUOHDigs83JyQne3t462xo0aKDzXIpG2URHUw6AtuyVK1cgCAI++OADfPDBB+X2Lz09HY0aNcLVq1cr/OqiKpcvXwYAjB49usIyWVlZKCoqQkFBAZo1a6a3v3nz5tiyZYtR7cpkMrz//vsYPXo0NmzYUG6CCQD9+vWDh4cHfvzxR+0ceT/99BMiIiK0A28s4eHPA6j4sxswYAB8fX3xxx9/lPvHUnkKCgowZ84cLF++HLdv39ZJTLOysozqq0Kh0Hkmqn79+pX2w9A4VqlUUCgUOvs9PDzg4OBg0f5XprKE/WGxsbGIjY1Ffn4+jh8/jjVr1mDp0qV48sknceHCBYs8m2bIz/Tly5eRlZVVYXuGDGTQJF779+9H48aNcfLkSXz44Yfw9vbGZ599pt3n5uaGiIgIAMZfn+zs7LS/cx6WkpKC559/HsOGDcPnn39eZX/Laty4sd7vBLlcrvdMtGbErua9UygUyM/PL7f/oaGhUKvVuHXrFlq2bIkbN26U+3jPw3U11zPN9f9hbm5uVZ6PJgat5fecBpO0GvLFF1+gTZs2ej8sAwcOxIoVK3Dy5EnExMRg8ODB2Lt3r3b/6NGjsWLFCsjlcvj7++Ps2bN6x9YEcWXfu3fr1k07urM8mr/IKpqmID8/X+/Ok6EqCvqHH8I1xqhRo7B27VocPHgQrVq1wu+//47x48dr/wr87LPPMHPmTG35wMBAiz7sXJZMJjO7rOYCoXkg9q233qrwTlVISIiRPdSnaWfevHmIjIwst0z9+vVRVFRkdlsPe+655zB79mzMmjWrwsmO7e3tMXz4cHz99ddIS0vDzZs3cfnyZcydO9eifanq8yhryJAh+O677/DDDz/g1VdfNej4b7zxBpYvX47JkyejU6dOkMvlkEgkGDlyZIUPP1ekXbt2OtPrJCQkYMaMGUYdozy3bt1CcHCwzrbdu3cjOjraov2viKenZ7lJcVVcXFzQtWtXdO3aFV5eXpg5cya2bt1a6R8ehjLkZ1qtVsPHx6fc52MB6CV55WnYsCGCg4Oxb98+BAUFQRAEdOrUCd7e3pg0aRJu3LiB/fv34/HHHzf4DtfDHB0dK6zr7+8Pf39/bNmyBceOHUNUVJTBx63oPTLmZ8pSNLG4cuVKvQEWAAwa2a2Jwcp+T4qBSVoNSUtLK/eWvubhUs1Imvnz5+tcsBo2bKj9//79++Obb77BkSNHdG7/WkJgYCAA4OLFi3p/CeXn5+PWrVvo06ePScfWnPfDD7SWN5+bofr27Qtvb2/88MMP6NChA/Lz8/HCCy9o948aNQpdunTRvtZ8baO5g3j27NkKk52y78XDf5ldvHhRu786aG7/29vbIyYmptKyTZs2LTdpL6uiBFnzPri5uVXajre3N5ydnbV/qZal+YrGWJq7aS+++CJ+++23Css999xzWLp0KdasWYOUlBRIJBKj5hS0tHnz5sHOzg7jx4+Hq6srnn322Srr/PLLLxg9erTO9BKFhYUGPbT/sB9++EHnjyhNrFT0GRsax35+fti+fbvOfs1dG0v2vyItWrQwe5S5Jrm4e/euJbpkkKZNm2LHjh3o3LmzztfC5ans7kzXrl2xb98+BAcHIzIyEq6uroiIiIBcLse2bdtw4sQJvT84Actcn5ycnLBp0yb07NkTffv2xd69ey16p7o83t7ecHFxKff6ceHCBUilUu3voMDAQIOuPZrrmY+PT5XXzYpoYtDSg93MxWfSasijjz6KkydP4tKlSzrbf/rpJ0ilUrRu3RpA6QiTmJgY7b+wsDBt2XfeeQcuLi4YO3ZsucP+zflLpVevXnBwcMCXX36p9xfysmXLoFQqjZo8tKzAwEDIZDKdUUwAzFqGw87ODs888wx+/vlnrFixAq1atdK+h0DpL7Cy76NmNGCfPn3g6uqKOXPm6M0urXn/oqKi4OPjg6VLl+rcSdq6dSvOnz+P/v37m9zvqvj4+CA6OhpfffVVub9wyn4tNWTIEJw6darcUZaac9GMGn74l2rbtm3RtGlTfPbZZ+V+ha5pRyaTITY2Fhs2bMDNmze1+8+fP48//vjD+BP81/PPP4+QkBCdXz4P69y5M4KCgrBq1SqsWbMG3bt3r/Brm5ogkUiwbNkyDB06FKNHjy53SpSHyWQyvZ/Lzz//3KS7yJ07d9aJaU2SVtFnbGgcOzk56Rw3JiZG+4eVJftfkU6dOuHs2bMG3bXduXNnuds1X7uX9xVadRk+fDhUKpX2kYuylEqlzudRr169ChPbrl274vr161izZo3260+pVIrHH38cCxYsQElJic7ITktfn+RyOf744w/4+Pigd+/e2kdCqotMJkOfPn3w22+/6Xy7kZaWhh9//BFdunTRfj35xBNP4PDhwzqjzRUKhd7dy9jYWLi5ueHjjz/WGSlftk5Vjh8/DolEYtF5Qy2Bd9JqyNtvv62dw2XChAnw9PTEpk2bsHXrVrz00ks6d8wq0qxZM/z444945pln0Lx5czz33HOIiIiAIAhISUnBjz/+CKlUatIvMh8fH0yfPh3vv/8+unXrhoEDB8LFxQUHDx7ETz/9hD59+mDAgAGmnDrkcrn2mQeJRIKmTZti06ZNZk8+OWrUKCxZsgS7d+/Gp59+alAdNzc3LFy4EC+99BLatWuHZ599Fg0aNMCpU6eQn5+P7777Dvb29vj0008xZswYdO/eHc888wzS0tKwePFiBAUFVftSV4mJiejSpQtatWqFl19+GY888gjS0tJw6NAh/PPPPzh16hSA0pj65ZdfMGzYMIwdOxZt27ZFRkYGfv/9dyxduhQRERFo2rQp3N3dsXTpUri6uqJevXro0KEDgoOD8c0336Bfv35o2bIlxowZg0aNGuH27dvYvXs33NzcsHHjRgDAzJkzsW3bNnTt2hXjx4+HUqnE559/jpYtW+otQ2YomUyG9957D2PGjKmwjEQiwbPPPouPP/4YALRTTZRHoVDgww8/1NseHBxs0YEGUqkUq1atwqBBgzB8+HBs2bKlwudgAODJJ5/EypUrIZfLERYWhkOHDmHHjh0Wm/oGQKWfsblxXBP9f+qppzB79mzs3bu3yrv1Tz31FIKDgzFgwAA0bdoUeXl52LFjBzZu3Ih27drpXaOOHTtWblxER0fr3Gk3Rffu3fHqq69izpw5SEpKQp8+fWBvb4/Lly9j7dq1WLx4MYYOHQqg9I+iL7/8Eh9++CFCQkLg4+OjjRtNAnbx4kVtrAOlj6hs3bpVO5+iRnVcn7y8vLB9+3Z06dIFMTExOHDgABo1amTO21OpDz/8UNve+PHjYWdnh6+++gpFRUU6jzS88847WLlyJfr27YtJkyahXr16WLZsGQIDA3WuPW5ubvjyyy/xwgsv4LHHHsPIkSPh7e2NmzdvYvPmzejcuXOVA/S2b9+Ozp07WzS2LaLGx5PWYX///bfQr18/wc/PT7C3txceffRR4aOPPhJKSkqMOs6VK1eE119/XQgJCRGcnJwEZ2dnoUWLFsJrr70mJCUl6ZStamqNh61atUro2LGjUK9ePcHR0VFo0aKFMHPmTJ3pICpT0XQICoVCGDJkiODi4iI0aNBAePXVV4WzZ8+aNAVHWS1bthSkUqnwzz//GNQ/jd9//114/PHHBWdnZ8HNzU1o37698NNPP+mUWbNmjdCmTRvB0dFR8PDwEJ577jm9dspOL1HWw+ehmYJj3rx5emUBCAkJCTrbrl69KowaNUobK40aNRKefPJJ4ZdfftEpd//+fWHChAlCo0aNBAcHB6Fx48bC6NGjhXv37mnL/Pbbb0JYWJhgZ2en936fPHlSGDx4sODp6Sk4OjoKgYGBwvDhw4WdO3fqtLN3716hbdu2goODg/DII48IS5curXAo+8Mqeo9KSkqEpk2bVjpVwrlz5wQAgqOjo/DgwYNyy2imAyjvX69evQRBqHgKjvLafTjeyvsZys/PF7p37y7Ur19fZ7qBhz148EAYM2aM4OXlJdSvX1+IjY0VLly4oNeGOVNYCELln7EhcSx2/1u3bi2MGzdOZ1t5U3D89NNPwsiRI4WmTZsKzs7OgpOTkxAWFia89957QnZ2tk79imICgDB79uwK+2Loz7TGsmXLhLZt2wrOzs6Cq6ur0KpVK+Gdd94R7ty5oy2Tmpoq9O/fX3B1dRUA6E3H4ePjIwAQ0tLStNsOHDggABC6du1abj/NuT4JQvnX6itXrgj+/v5CaGhopb8zKrrOBwYGCv3799fbXt7P2okTJ4TY2Fihfv36gouLi9CjRw/h4MGDenVPnz4tdO/eXXBychIaNWokzJ49WzsVi2YKDo3du3cLsbGxglwuF5ycnISmTZsKL774onDs2DFtmfI+x8zMTMHBwUH45ptvKjxnsUgEoRqf5iOqZm3atIGHh0eFX4MQkfVbuXIl4uLicPPmTS4rRzVu0aJFmDt3Lq5evVrl84U1jc+kkc06duwYkpKSMGrUKLG7QkRmeO6559CkSZNy53gkqk4lJSVYsGAB3n//fatL0ACAd9LI5pw9exbHjx/H/Pnzce/ePVy7ds3k6UGIiIisFe+kkc355ZdfMGbMGJSUlOCnn35igkZERLUS76QRERERWSHeSSMiIiKyQkzSiIiIiKyQzU9mm5mZiZiYGCiVSiiVSkyaNAkvv/yywfXVajXu3LkDV1dXq1tYlYiIiGoXQRCQk5ODhg0bVrkmq80/k6ZSqVBUVAQXFxfk5eUhPDwcx44dM3jW4H/++UdvrUoiIiKi6nTr1q0qVwiy+TtpMpkMLi4uAICioiIIgmDUGpaurq4ASt8szXpharUaCoUC3t7e5Wa55u63BWKcg6XbNPd4ptQ3po6hZRmPjEdT6zMeqwfj0Tri0dwyYsVidnY2AgICtPlHZURP0vbt24d58+bh+PHjuHv3LtavX49BgwbplElMTMS8efOQmpqKiIgIfP7552jfvr12f2ZmJrp3747Lly9j3rx58PLyMrh9zVecbm5uOklaYWEh3NzcKrzImLPfFohxDpZu09zjmVLfmDqGlmU8Mh5Nrc94rB6MR+uIR3PLiB2LhjxiJfpPSF5eHiIiIiqcaXrNmjWIj49HQkICTpw4gYiICMTGxuoszu3u7o5Tp05pFxlPS0urqe4TERERVQvR76T169cP/fr1q3D/ggUL8PLLL2PMmDEAgKVLl2Lz5s349ttvMXXqVJ2yvr6+iIiIwP79+zF06NByj1dUVISioiLt6+zsbAClGbVardb+vyAI2tcPM3e/LRDjHCzdprnHM6W+MXUMLct4ZDyaWp/xWD0Yj9YRj+aWESsWjWlP9CStMsXFxTh+/DimTZum3SaVShETE4NDhw4BANLS0uDi4gJXV1dkZWVh3759eP311ys85pw5czBz5ky97QqFAoWFhQBK38CsrCwIglDh7Xpz9tsCMc7B0m2aezxT6htTx9CyjEfGo6n1GY/Vg/FoHfFobhmxYjEnJ8fgsladpN27dw8qlQq+vr462319fXHhwgUAwI0bN/DKK69oBwy88cYbaNWqVYXHnDZtGuLj47WvNQ/weXt76zyTJpFItA8TqtVqFBcXa+uo1WoolcpKn7mobL8tEOMcLN2mucczpb4xdQwta2681cZ4tLe3h0wmq/Y2y14HxD6eKfWNqWNo2arKmbvfFohxDoxHy5cRKxaNWcrQqpM0Q7Rv3x5JSUkGl3d0dISjo6PedqlUqvMhSSQSSKVSKJVKpKSk6Nye1Nwezc3NLffBv6r22wIxzsHSbZp7PFPqG1PH0LLmxlttjUd3d3f4+flV6zlprgOWuoCbezxT6htTx9CyVZUzd78tEOMcGI+WLyPG52hMW1adpHl5eUEmk+kNBEhLS4Ofn1+1ty8IAu7evQuZTIaAgADtGysIApRKJezs7Cr8pVjZflsgxjlYuk1zj2dKfWPqGFrW3HirbfEIAPn5+drBQ/7+/mJ2jYio2lh1kubg4IC2bdti586d2mk51Go1du7ciQkTJlR7+0qlEvn5+WjYsKF2Ljag7v1SZJLGJE1sD5+Ds7MzACA9PR0+Pj7V/tUnEZEYRE/ScnNzceXKFe3rlJQUJCUlwcPDA02aNEF8fDxGjx6NqKgotG/fHosWLUJeXp52tKelqNX6ozuVSiUEQYC9vb3eBLma1xVNnFvVflsgxjlYuk1zj2dKfWPqGFrW3HirjfHo7OwMQRBQVFRk1DMehuJoOtPKcXSnbbRZG+KRoztrwLFjx9CjRw/ta81D/aNHj8aKFSswYsQIKBQKTJ8+HampqYiMjMS2bdv0BhMYKzExEYmJiVCpVADKH91ZUlICtVoNlUoFpVKprSsIgrZeRXcuKttvC8Q4B0u3ae7xTKlvTB1Dy5obb7U1HlUqFdRqNe7fvw97e3uLt8nRdKaV4+hO22izNsQjR3fWgOjo6Cr/up8wYYLFv96Mi4tDXFwcsrOzIZfLyx3d6erqitzcXNjZ2WmfhSmrql8M1fGLo6aJcQ6WbtPc45lS35g6hpY1N95qWzza2dlBKpXC09Oz2u6kcTSd8eU4utM22qwN8cjRnXVIRaM7JRKJ9p+GIAja1xXduahsvy0Q4xwqajMoKAiTJ0/G5MmTDTrOjBkzsGHDBpw8ebLSc/jggw+QlpaGZcuWGdUfU87BnLLmxpsp56F5DzUjp1988UVkZmZiw4YNBtXXtFXeMm+GGDlyJNq1a4c333wTQPnnoPm5rM6RWRxNZ1o5ju60jTZrQzxacnSnSi3gSEoG0nMK4ePqhPbBHpBJLf/7z5j3y3Z/QmyISi3g0NX7+C3pNg5dvQ+VunqfC3rxxRd1fjEqFAq8/vrraNKkCRwdHeHn54fY2Fj89ddf2jJBQUFYtGhRtfbLmqSmpmLx4sV47733tNv27duHAQMGoGHDhpBIJOUmJC+++KJO4i6RSNC3b1+dMhkZGXj++efh5uYGd3d3jBs3Drm5uVX26eTJkxgxYgT8/f3h6OiIwMBAPPnkk9i4caP2bvP169d12nZwcECzZs3w8ccfl3tHeubMmXjhhReMfHf0LV68GCtWrDD7OIZ6//338dFHHyErK6vG2iSiumvb2bvo8ukuPPP1YUxanYRnvj6MLp/uwrazd0XtF++kVbNtZ+9i5sZk3M0q1G7zlzshYUAY+obXzNQBQ4YMQXFxMb777js88sgjSEtLw86dO3H//v0aad8affPNN3j88ccRGBio3aZZR3bs2LEYPHhwhXX79u2L5cuXa18/PO/e6NGjkZqaiu3bt6OkpARjxozBK6+8gh9//LHCY/72228YPnw4YmJi8N133yEkJARFRUU4ePAgPvjgA3Tq1AleXl7a8jt27EDLli1RVFSE/fv34+WXX0ajRo3w0ksv6R333XffNfh9qYhcLjf7GMYIDw9H06ZNsWrVKsTFxdVo20RUt2w7m4q4H0/i4T9zU7MK8fqqE/jy+cdq7Pf1w3gn7V+a0Z2af5oRH5qVDMr+A3RHmlW0f9vZVLy+6oROggb874PfeuZuufXN/achCAIePHiA/fv345NPPkF0dDSaNGmCdu3aYerUqRgwYIBenw09x6CgIMyePRujRo1C/fr1ERgYiN9++w3p6el46qmnUL9+fbRu3RpHjx7VqffLL7+gZcuWcHR0RFBQED777DOdNtPS0jBw4EA4OzsjODgYq1at0uvDgwcPMG7cOO1zhD179kRSUlKF70F557B69Wo8+eSTOtv69u2L2bNn69yFfLg+UJqU+fr6av+5u7tr958/fx5//PEHvv76a7Rv3x6dO3fGkiVLsHr1aty+fbvc9zgvLw/jxo1D//79sWnTJvTu3RvBwcFo0aIFxo4di6SkJO3zkpp6Hh4e8PX1RZMmTfDcc8+hU6dOOHHihM7xb968iXPnzqFv375Qq9WYMWOG9m5qw4YN8cYbb1QZR5r/19yd1byOjo7GG2+8gbfffhseHh7w8/NDQkJChfUFQcD06dPh7++PU6dOQRAEJCYmolmzZnBycoKvry+GDh2qU/7JJ5/E6tWrq4zHh392LfnP0sc393im1DemjqFlqypn7n5b+CfGOTAeLV9GqVJj1qZkvQQNgHbbzI3JKFGqLPreG6rO3kkzZnSnUqnUTsdRUKKCIJSOLJMVq1DeIz6CABSXKDHj93MVfvASADM2nkOHILlB33k728sMfp5IG3xKJZycnFC/fn2sX78eUVFR5a62ULaeZhSrIFQ9InDRokWYPXs2pk6diiVLlmDUqFHo1KkTRo8ejY8//hj/93//h1GjRuHUqVOQSCQ4ceIERowYgQ8++ADDhg3D4cOH8cYbb6BBgwYYNWoUBEHAuHHjtHeh7O3tMWXKFKSnp+v0bdiwYXBycsLGjRvh5uaGb775BjExMTh37hw8PDy0P5QlJSXlnkNGRgaSk5PRpk0bnVG7D1MqlXr11Wo19uzZo03OevTogZkzZ8LT0xMAcODAAbi7uyMiIkJ77OjoaEilUhw8eFAvAVSpVNi6dSvu37+P+Pj4cvujucCUlJRo92tiEigdIX3y5Ek8//zzOvU3bNiA7t27w9nZGb/88gsWLVqEVatWISwsDGlpaTh9+nSF5695DzX7y8aUpk/ff/89Jk2ahAMHDuDw4cN46aWX0LFjR8TExGiPo1KpUFJSgilTpmDLli3YtWsXQkJC8Pfff2PSpElYvnw5OnXqhIyMDPz11186/Wnbti0+/vhj5OXlwcHBQe+zUCqVUKs5utNSdQwtW1U5c/fbAjHOgfFo+TJqtRp/XUpFanZRhf0QANzNKsSfJ6+hbYBr5SdoIJsa3SkWU0Z35hcrETF7l0XaFwCkZRfhsY92G1T+3Mw+cLE37OPSPASp6ffy5cvxyiuvYNmyZXjsscfQrVs3jBw5Eq1bt9ar9/Ao1sp++T3xxBPaxewTEhLw1VdfoV27dhg5ciQAYOrUqXj88cdx//59+Pn5YcmSJejVqxcSEhIAAGFhYbhw4QIWLFiAsWPH4tKlS/jzzz/x999/o127dgCA//73vwgLC9P27cCBAzh69CjS0tK0Cef8+fPx+++/Y8OGDXjllVe0Az40fX/4HO7cuQNBEBAQEFDuqF0NOzs7yGQynfr9+vXDkCFDEBwcjKtXr+K9997DwIEDcfDgQchkMigUCnh7e+uNQvTw8IBCoSi3vatXr2rfD83+o0ePomfPntoyK1euxKBBg7T7u3fvDqlUiuLiYpSUlGDcuHF6cwdu2rQJAwcOhL29Pf755x/ts4j29vZ45JFH0KlTpwrPXfMeatorG1NAaaLUunVrzJw5EwAQGhqKpUuXYs+ePTrP6AmCgDFjxuDkyZPYv38/GjVqBAC4ffs26tWrh6eeegqurq5o2rSp9jPXCAgIQHFxMe7du6f9WpqjOzm60xqIcQ6MR8uXUavVKLz0wKC+l9g5w8fHx6CyVeHoThMYMrpTzJGaprSvKT906FA8+eST2L9/Pw4fPoytW7di3rx5+Oabb/Diiy+W24YgVD0isHXr1tp9mmW6ytumUCjg7++P8+fP46mnntI5XpcuXbB48WKo1WqcP38ednZ2aNu2rbZMaGgo3N3dtX07ffo0cnNzdZ7PAoCCggJcu3ZN730q7xw0d0ydnZ2rfE8frv/MM8/onH9ERASaNm2KvXv3olevXuXWLfu6olHCD++PiIjQjqxs1qyZ3l2kNWvWIDQ0FCUlJThz5gwmTpyIqVOn4tNPPwUAZGdnY+/evfjvf/8LoDQG/vOf/6Bp06bo27cvnnjiCQwYMKDCJLWiz77s67KfNVC6PJNCodDZFh8fD0dHRxw+fFjnM+vTpw8CAwO1/enbty+efvppnZU9NP9fUFBQbp84utPydTi603Ac3Wkd8WhumfxilUH99nVzttj7bsxxmKQZwdlehuRZsdqvgSpbhufQFQVeWnmyymOuGNMO7YM9DGrbHE5OTujduzd69+6NDz74AC+99BISEhJ0kjRjlb2roXkfyttmzPfvVcnNzYW/vz/27Nmjt8/d3d2gY2iShQcPHsDb29us/jzyyCPw8vLClStX0KtXL/j5+UGhUOiUUSqVyMjIqHC92WbNmgEALl68iI4dOwIofe4tJCSkwnYDAgK0+1u0aIHLly9jxowZmDlzJpycnLB161aEhYUhICBAe9fwwoUL2LlzJ7Zv347x48dj3rx52Lt3r8lfFT5cTyKR6H3WvXv3xk8//YQ//vgDzz33nHa7q6srTpw4gT179uDPP//E9OnTMWPGDBw9elT7OWZkZACA2Z8REdHDVGoBX+29is9236q0nASAn9zJoN/T1cF2/4wRgUQigYuDnUH/uoR4wU/uhIru00hQOsqzazNvg45n6bt4YWFhyMvLs+gxqxIaGqoz7QcA/PXXX3j00Uchk8nQokULKJVKHD9+XLv/4sWLyMzM1L5+7LHHkJqaCjs7O4SEhOj8e/juWkWaNm0KNzc3JCcnm31O//zzD+7fv69d5LtTp07IzMzUOYddu3ZBrVajQ4cO5R6jT58+8PDw0N4FM4VMJoNSqURxcTGA0lGdTz31lE4ZZ2dnDBgwAEuWLMGePXtw6NAhnDlzxuQ2DTFw4ED8+OOPeOmll7B69WqdfXZ2doiJicHcuXNx+vRpXL9+Hbt2/e9xgrNnz6Jx48YGf65ERIa4k1mA5745jLl/XIJaACID5JAAer+vNa8TBoRVy3xphuCdtH+VHXFRdjRIeSMGgarXQpRKgIQnwzD+hxOQADoDCDQf9fQnwyCVVN96ioIg4P79+xg+fDjGjBmD1q1bw9XVFceOHcPcuXMxcOBAnbb/+ecfnDz5v7t/SqUSTZs2RYMGDSo8fnnvS0Wj8eLj49G+fXvMmjULI0aMwKFDh/Cf//wHiYmJEAQBzZs3R58+ffDaa6/hiy++gJ2dHaZMmaJdo1EQBPTq1QudOnXCoEGD8Omnn+LRRx/FnTt3sHnzZjz99NOIioqqcHSnhkQiQUxMDPbv36+TyFS0jqyPjw+aNGmC3NxczJw5E0OGDIGfnx+uXr2Kd999FyEhIejTpw8EQUBoaCj69OmDV155BV9++SVKSkowYcIEjBw5Ev7+/uW+X/Xr18fXX3+NkSNHon///njjjTfQrFkz5ObmYtu2bQD+d3tcU//evXu4e/culEolzpw5g//85z/o0aMHXF1dUVJSgq1bt+LNN9/Ulv/uu+8AAB06dICLiwtWrlwJZ2dnNGnSpNz4qyi+H35vK6ur+f9Bgwbh+++/x6hRoyCTyTB06FBs2rQJ165dQ7du3dCgQQNs2bIFarUajz76qLb+/v370bt373LjqWz7xo6WMlTZ64A1HM+U+sbUMbRsVeXM3W8LxDgHxqNlymw+fRfvbTiL7EIlXBxkmNytEcZ0b4Ht5xWYtek8UrP/NxuDn9wJH/QPRZ8wX4t+1hzdaQBTRndqVDXyUbO/V3NPfD4yAh9uuaAzesRP7oj3+rVATAuvSkcXmurh0Z1RUVFYuHAhrl27hpKSEjRu3Bhjx47F1KlTddqfP38+5s+fr3Os5cuX63xN9XA7D/e/7DqnD49EbN26NX788UfMnDkTH374Ifz9/ZGQkKAdlSgIApYuXYq4uDhER0fD19cXM2bMwK1bt3Ta+u233zB9+nSMHTsWCoUCfn5+6NKlCzw9PbUj/iob3QmUTkr7+uuv4+OPP9YmQH///Td69+6tLaOZ7f6FF17Af//7XwiCgNOnT+P7779HZmYmGjZsiJiYGMyYMUN7J0sQBHz77bd48803ERMTA6lUiqeffhoLFy7Ue6/KxtGAAQOwb98+zJs3D6NHj0ZGRgbkcjnatm2LVatWoW/fvjqjOzX9lMlk8Pf3R58+fTB79mwolUrs2rVLOwWKpk+urq5YsGAB3nzzTahUKoSHh2P9+vWQy+XlxqAhozvL7i+vTNl4GDRoEP773/9qR/H6+Pjg119/xcyZM1FYWIiQkBCsXLkSzZs3h1KpRGFhITZs2IBNmzZpz4GjO8UfTWdIOY7utI02a0M8GlMmt1CJhftuY8v50vlBw3xdkBAbCDdJERQKBR7zkWLdi2FIup2L+3kl8Kxnj8hG9SGTSpCenm7Q+RnKmNGdEqG6buPYCM3ozgcPHuiM7lQoFHB1dcXNmzcRHBysNxqjpKSk0l8MZfer1AKOXs9AenYRfNwc0S6oepaasLSqztEW2qzoeIIgoGPHjpg8ebLOYABL9MeYOoaWNSbeJk6cCKVSiS+++MKkPlmDL7/8Ehs2bMAff/yh3fbwORQWFiIlJQVBQUHVNrpTM1rXUr8UzTmeKfWNqWNo2arKmbvfFohxDoxH08vsOp2C2dtv4daDAkglwPjopnijZwhkEogSi9nZ2WjQoAGysrK0eUdF6uydtIcZMrpTo6qRjw/vt5NJ0KmpbT1XY8joTmtvs7LjSSQSLFu2DGfOnKmwLVP6Y0wdQ8saG2+tWrVCp06djBqpa20cHBzw+eefV3oOHN1p+Toc3Wk4ju60jnisrIxSpcbnu6/iP7uuQCUAjdydsXBEpHYQgGZ6jpr+HDm6k8gAkZGRiIyMFLsbFvfKK6+I3QWzPby8FRGRMW7ez8fkNSdx4mYmAOCpyIaYPSgcbk62840CwCSNiIiIaglBEPDridtI+P0ccouUqO9oh7d6NMaobqE2eeeWSRoRERHZvKz8Ery34Qw2nb4LAGgX1ADzh7WGQ0muyD0zHZM0A9TxsRVEVok/l0Skcejqfbz5cxLuZBVCJpVgSkwzvB4dAgkEpKczSbN55c2TJpFIIAgCiouL9UaPVTVPWlX7bYEY52DpNs09nin1jaljaFlz4602xmNeXh4EQYBMJuM8aRaow3nSDMd50qwjHtVqNYqVKny67QKW7U+BIACBni5YNDwCEQHuAIRKjyNWLHKeNAMYMk+aZuRHWlqa3rqWarVaO/rzYVXttwVinIOl2zT3eKbUN6aOoWXNjbfaFo9A6fQb6enpcHR0xP3796ulTc5LZVo5zpNmG23WhnhMuZ+P6Vuu4vL90pVWBrT0xJTuAXBxKNbObVbZccSKRWPmSauzSVpcXBzi4uK086R5e3vrzJMmkUjg7e0NT09PXL9+Hf/8849O/bK/MMpT1X5bIMY5WLpNc49nSn1j6hha1tx4q43x6OHhAV9f32pLPMteByz1S9Gc45lS35g6hpatqpy5+22BGOfAePwfQRCw+ugtfLj5EgpKVJA72+Hjp1uhX7j++siVHUesWDRmXsc6m6Q9rKJ50pycnPDoo49q10QESj/Y+/fvw9PTs8KLUGX7bYEY52DpNs09nin1jaljaFlz4602xqO9vT1kMlm1t8t5qUwrx3nSbKNNW4zH+7lFmPrrGWxPTgMARAW4YsmzUWjYwMWkNjlPWi2gSdY01Go17O3t4eTkVOEvxcr22wIxzsHSbZp7PFPqG1PH0LLmxhvjkYhqg32XFHhz7SkocopgL5Pg7T7N8eSjLvCTW37FEWvBJI2IiIisVmGJCnO3XcS3f6UAAEJ86mPxyEiE+rlafF1Na8MkjYiIiKzShdQcxP98ChdSSx+2H9UpENP6hcLZoXpGdVsbJmlERERkVQRBwJqT6Uj86zaKlWp41XfA3KGt0bOFr9hdq1FM0oiIiMhqpOcU4u21p7D30j0AQI/m3pg7NALero4i96zmMUn7V3mT2XKyRk7WKPZkjYaUYzzaRpuMR+P7ZK0Yj9UXjzvOp2HqujPIyC+Bg0yCaf1aYFSnQEgkEr16hhyvsjJixaIx7dXZJM2QyWw5WSMnaxR7skZDyjEebaNNxqPp52FtGI+Wj8fCEjUW77uF9WdK754183LGm5090TrQCQqFwuQ+VFZGrFjkZLYGMHQyW07WyMkaOXmo+BiPjEdrwni0bDyevZ2FyWtO4dq9PADAS12CMSUmBNkP7ld6fEP6UFkZsWKRk9maoKLJbDlZIydr5OSh1oHxyHi0JoxH8+NRpRbw9f5rmP/nRZSoBPi6OWL+sEh0aeYFtVqNHAOOb0gfOJktERERkYHuZBYg/uckHL6WAQDo29IPcwa3QoN6DiL3zLowSSMiIqIas+XMXfzf+rPILlTCxUGGGQNaYlhU42pbh9eWMUkjIiKiapdbpMSsP65jy/n7AICIxnIsGtkGwV71RO6Z9WKSRkRERNXq+I0HmLLmJG5mFEAqAcZHh2BSTDPYy2z3ucSawCSNiIiIqoVSpUbi7qtYsusyVGoBfq4OWDSyDTo29RK7azaBSRoRERGZRaUWcCQlA+k5hfBxdUL7YA/cflCAKT8n4fiNBwCAgRH+eKOTD5o28RC5t7aDSRoRERGZbNvZu5i5MRl3swq12+TO9igqUaFQqYarox1mDwrHwAh/pKeni9hT28MkjYiIiEyy7Wwq4n48CeGh7VkFJQCApt71sGJMewR4uNj0UmBiYZL2L67dqYtr03GtRGvCeGQ8WhPGY2l9pUqNWZuS9RK0svKLVfBzc9T+jrVkPHLtzlqMa3dWjmvTca1Ea8J4ZDxaE8Zjaf2/LqUiNbuo0nJ3swrx58lraBvgavF45NqdtRjX7qwc16bjWonWhPHIeLQmjMfS+oWXHhhUtsTOGT4+PhaPR67dWYdw7U59XJuOayVaE8Yj49Ga1PV4FAQBF9LzDSrr6+asbcPS8ci1O4mIiIj+lZFXjHd/OYXt5ysfqSkB4CcvnY6DTGO7f8YQERFRjdp3SYHYRfuw/Xw67KQSDG7TCBKUJmRlaV4nDAiDTMo1OU3FO2lERERUqcISFeZuu4hv/0oBAIR418P03k3QJTwIfVr66s2T5id3QsKAMPQN9xery7UCkzQiIiKq0MXUHExafRIXUktHJY7qFIh3Y5sjJ7N0ofS+4f7oHeant+IA76CZj0kaERER6REEASsOXsecrRdQrFTDs54D5g5tjV6hvlCr1Sg7kYRMKkGnpp6i9bW2YpJGREREOtJzCvH22tPYe0kBAOjR3Btzh0bA29VR5J7VLUzSiIiISGt7chreXXcaGXnFcLST4r3+oXihYyAkEn59WdOYpBEREREKilX4cHMyfvj7JgAg1N8NS0ZGopmvq8g9q7uYpBEREdVxZ29nYeLqk7imyAMAvNw1GG/FNoejnUzkntVtTNKIiIjqKJVawLL9VzH/z4soUQnwdXPE/GGR6NLMS+yuEZikaanVau3K9Gq1GoIgVLhSvbn7bYEY52DpNs09nin1jaljaFnGI+PR1PqMx+pRW+IxNbsIk377G3+nlK7BGdvSFx8/HY4GLg5mx4G5dQwpa24ZsWLRmPbqbJKWmJiIxMREqFQqAIBCoUBhYelEfGq1GllZWRAEocIFgs3ZbwvEOAdLt2nu8Uypb0wdQ8syHhmPptZnPFaP2hCPOy7exyc7byK3WA1neyniowPwZJgnSnIzkZ5bPf2xdDyaW0asWMzJyam60L/qbJIWFxeHuLg4ZGdnQy6Xw9vbG25ubgBKPziJRAJvb+8KLzLm7LcFYpyDpds093im1DemjqFlGY+MR1PrMx6rhy3HY26REjM3JmPdidsAgNaN3LBwRCSCvepVe38sHY/mlhErFp2cnAwuW2eTtIdJpVKdD0kikehtK8vc/bZAjHOwdJvmHs+U+sbUMbQs45HxaGp9xmP1sMV4PHHzASavTsLNjHxIJcCodn6YNiACjvampQLWEI/mlhHjczSmLSZpREREtZhSpUbi7qtYsusyVGoBjdydMX9YawTVU8JeZruJcl3AJI2IiKiWunk/H1N+TsLxG6WDA56KbIjZg8JR30GG9PR0kXtHVWGSRkREVMsIgoBfT9xGwu/nkFukhKujHWYPCsegNo0AGDfCkMTDJI2IiKgWycovwXsbzmDT6bsAgHZBDbBgeCQCPFxE7hkZi0kaERFRLXH42n3Er0nCnaxCyKQSTIlphtejQyCTct1NW8QkjYiIyMYVK9VYuOMSlu69CkEAAj1dsHhkG0QGuIvdNTIDkzQiIiIbdlWRi8mrk3DmdhYAYERUAKYPCEM9R/6Kt3X8BImIiGyQIAhYffQWZm1MRkGJCnJne3wyuBX6tfIXu2tkIUzSiIiIbExGXjHeXXca25PTAACdQzwxf1gk/OSGz2ZP1o9JGhERkQ3Zd0mBN9eegiKnCPYyCd6JbYFxXYIh5eCAWodJGhERkQ0oUqoxe/N5LP/rOgAgxKc+Fo+MRMuGcnE7RtWGSRoREZGVu5iagwmrL+DqvQIAwKhOgZjWLxTODjKRe0bViUkaERGRlRIEASsOXsecrRdQrFTDs54D5g1rjZ4tfMXuGtUAJmlERERWKD2nEG+vPY29lxQAgMeD3LDwmSj4yp1F7hnVFCZpREREIlKpBRxJyUB6TiF8XJ3QPtgDuy+k4511p5GRVwxHOymm9WuB2Eec4O3qKHZ3qQYxSSMiIhLJtrN3MXNjMu5mFWq3uTjIkF+sAgCE+rthychINPWuh/T0dLG6SSJhkkZERCSCbWfv4vVVJyA8tF2ToPUO88V/nm0DRzsZ1Gp1zXeQRMck7V9qtVr7Q6BWqyEIQoU/FObutwVinIOl2zT3eKbUN6aOoWUZj4xHU+szHquHJc5BpRYw4/dkvQStrLO3syDF/34/MR4tW0asWDSmvTqbpCUmJiIxMREqVelfLAqFAoWFpbeb1Wo1srKyIAgCpFKpXl1z99sCMc7B0m2aezxT6htTx9CyjEfGo6n1GY/VwxLncPxWDlKzCystczerEH+evIa2Aa6Mx2ooI1Ys5uTkGFy2ziZpcXFxiIuLQ3Z2NuRyOby9veHm5gag9IOTSCTw9vau8CJjzn5bIMY5WLpNc49nSn1j6hhalvHIeDS1PuOxeljiHEruKg0rZ+cMHx8fxmM1lBErFp2cDF+6q84maQ+TSqU6H5JEItHbVpa5+22BGOdg6TbNPZ4p9Y2pY2hZxiPj0dT6jMfqYc455BYp8evJ2waV9XVz1rbBeLR8GTFi0Zi2mKQRERHVkBM3H2Dy6iTczMivtJwEgJ+8dDoOqrts988YIiIiG6FUqbF4x2UMW3oINzPy0cjdGW/2fhQSlCZkZWleJwwIg4yLptdpvJNGRERUjW7ez8eUn5Nw/MYDAMCgyIaYNSgcbk72aOZbX2+eND+5ExIGhKFvuL9YXSYrwSSNiIioGgiCgPUnb2P6b+eQW6SEq6MdZg8Kx6A2jbRl+ob7o3eYn96KA7yDRgCTNCIiIovLKijB+xvOYuOpOwCAdkENsGB4JAI8XPTKyqQSdGrqWdNdJBvAJI2IiMiCDl+7j/g1SbiTVQiZVIIpMc3wenQI746R0ZikERERWUCxUo2FOy5h6d6rEAQg0NMFi0e2QWSAu9hdIxvFJI2IiMhMVxW5mLw6CWduZwEARkQFYPqAMNRz5K9ZMh2jh4iIyESCIGD10VuYtTEZBSUqyJ3t8cngVujXiiMzyXxM0oiIiEyQkVeMaevPYntyGgCgc4gn5g+LhJ/c8GV/iCrDJI2IiMhIf9/Ixoc7zkKRUwR7mQTvxLbAuC7BkHJwAFkQkzQiIiIDFZao8Om2C1j+13UAQIhPfSweGYmWDeXidoxqJSZpREREBriYmoNJq0/iQmoOAOCFjk3wf0+EwdlBJnLPqLZikkZERFQJQRCw4uB1zNl6AcVKNTzrOeD/Yprg6Q7NIJVyCWyqPkzSiIiIKpCeU4i3157G3ksKAECP5t74ZHArCAVZIveM6gImaUREROXYnpyGd9edRkZeMRztpHivfyhe6BgIQRCQXiB276guYJJGRERURkGxCh9uTsYPf98EAIT6u2HJyEg083UFUPr1J1FNYJJGRET0r7O3szBx9UlcU+QBAF7uGoy3YpvD0Y6DA6jmMUkjIqI6T6UW8PX+a5j/50WUqAT4ujli/rBIdGnmJXbXqA5jkkZERHXancwCxP+chMPXMgAAfVv6Yc7gVmhQz0HknlFdxySNiIjqrM2n72Lar6eRXaiEi4MMMwa0xLCoxpBIuHIAiY9JGhER1Tm5RUok/HYO6078AwCIaCzHopFtEOxVT+SeEf0PkzQiIqpTjt94gClrknAzIx9SCTA+OgSTYprBXsaJacm6MEkjIqI6QalSI3H3VSzZdRkqtYBG7s5YOCIS7YM9xO4aUbmYpBERUa13834+pvychOM3HgAAnopsiNmDwuHmZC9yz4gqxiSNiIhqLUEQ8OuJ20j4/Rxyi5RwdbTD7EHhGNSmkdhdI6qSzX8Bf+vWLURHRyMsLAytW7fG2rVrxe4SERFZgaz8Erzx00m8ufYUcouUaBfUAFsmdWWCRjbD5u+k2dnZYdGiRYiMjERqairatm2LJ554AvXqcYQOEVFddfjafcSvScKdrELIpBJMiWmG16NDIJNyag2yHTafpPn7+8Pf3x8A4OfnBy8vL2RkZDBJIyKq5VRqAUdSMpCeUwgfVye0D/aASi1g4Y5LWLr3KgQBCPR0waIRkWjTpIHY3SUymuhJ2r59+zBv3jwcP34cd+/exfr16zFo0CCdMomJiZg3bx5SU1MRERGBzz//HO3bt9c71vHjx6FSqRAQEFBDvSciIjFsO3sXMzcm425WoXabd30HODvY4WZGPgBgRFQApg8IQz1H0X/VEZlE9GfS8vLyEBERgcTExHL3r1mzBvHx8UhISMCJEycQERGB2NhYpKen65TLyMjAqFGjsGzZsproNhERiWTb2bt4fdUJnQQNABS5xbiZkQ8XBxm+fO4xfDq0NRM0smmiR2+/fv3Qr1+/CvcvWLAAL7/8MsaMGQMAWLp0KTZv3oxvv/0WU6dOBQAUFRVh0KBBmDp1Kh5//PFK2ysqKkJRUZH2dXZ2NgBArVZDrVZr/18QBO3rh5m73xaIcQ6WbtPc45lS35g6hpZlPDIeTa1fG+NRpRYw4/dkCJWUqe9oh5hQn2qLF8ajdcSjuWXEujYa057oSVpliouLcfz4cUybNk27TSqVIiYmBocOHQJQOrz6xRdfRM+ePfHCCy9Uecw5c+Zg5syZetsVCgUKC0v/KlOr1cjKyoIgCJBK9W82mrvfFohxDpZu09zjmVLfmDqGlmU8Mh5NrV8b4/H4rRykZhdWWiY9pwh/nryGtgGuRh/fEIxH64hHc8uIdW3MyckxuKxJSVpJSQlSU1ORn58Pb29veHhUz2zN9+7dg0qlgq+vr852X19fXLhwAQDw119/Yc2aNWjdujU2bNgAAFi5ciVatWpV7jGnTZuG+Ph47evs7GwEBATA29sbbm5uAEo/OIlEAm9v7wovMubstwVinIOl2zT3eKbUN6aOoWUZj4xHU+vXxngsuas0rJydM3x8fIw+viEYj9YRj+aWEeva6OTkZHBZg5O0nJwcrFq1CqtXr8aRI0dQXFwMQRAgkUjQuHFj9OnTB6+88gratWtnUqdN1aVLF6NuHTo6OsLR0VFvu1Qq1fmQJBKJ3rayzN1vC8Q4B0u3ae7xTKlvTB1DyzIeGY+m1q9t8VisrOyLzv/xdXOu1lhhPFpHPJpbRozP0Zi2DCq5YMECBAUFYfny5YiJicGGDRuQlJSES5cu4dChQ0hISIBSqUSfPn3Qt29fXL582eTOl+Xl5QWZTIa0tDSd7WlpafDz87NIG0REZP0EQcDyv1Lw3oYzlZaTAPCXO3E9TqoVDLqTdvToUezbtw8tW7Ysd3/79u0xduxYLF26FMuXL8f+/fvRrFkzszvn4OCAtm3bYufOndppOdRqNXbu3IkJEyaYffyyOHBAFx+MtY4HYw0px3i0jTYZj8b3SUORU4R31p3G3kv3AAAtG7rh3J1sSACdAQSaaWo/6B8KCQSo1YbddTMW49E64pEDB/71008/GXQwR0dHvPbaawY3DgC5ubm4cuWK9nVKSgqSkpLg4eGBJk2aID4+HqNHj0ZUVBTat2+PRYsWIS8vTzva01SJiYlITEyESqUCwIEDD+ODsdbxYKwh5RiPttEm49G089h3NRMf77iBzAIlHGUSvNGtMYa09saeq5lYuOcW0nNLtGV96ttjcnQAHvOR6k3TZEmMR+uIRw4cMEB2djZ27dqF5s2bIzQ01Oj6x44dQ48ePbSvNQ/1jx49GitWrMCIESOgUCgwffp0pKamIjIyEtu2bdMbTGCsuLg4xMXFITs7G3K5nAMHHsIHY63jwVhDyjEebaNNxqNx7RQUq/DRlvP48cgtAECovysWDY9AM9/SEZsjfH0xtOOjOHo9A+k5RfBxdUS7II8aWfaJ8Wgd8ciBA+UYPnw4unXrhgkTJqCgoABRUVG4fv06BEHA6tWrMWTIEKOOFx0dDUGo/Jb0hAkTLP715sM4cEAfH4y1jgdjDSnHeLSNNhmPhu0/ezsLE1efxDVFHgDg5a7BeCu2ORztZDrlpFLg8RDvqk6rWjAerSMeOXDgIfv27UPXrl0BAOvXr4cgCMjMzMSSJUvw4YcfGns4IiIiAKUT1S7dexVPf/EXriny4OvmiFXjOuC9/mF6CRpRXWB0kpaVlaWdF23btm0YMmQIXFxc0L9/f4uN6iQiorrlTmYBnvvmMD7ZegElKgGxLX2xbVI3dGnmJXbXiERj9NedAQEBOHToEDw8PLBt2zasXr0aAPDgwQOjvme1NhzdqYujl6xj9JIh5RiPttEm47Hi/VvO3MX/rT+L7EIlnO1lmD4gFMPbNoZEIrHKuGU8Wkc8cnRnOSZPnoznnnsO9evXR2BgIKKjowGUfg1a0Sz/1oijOyvH0UvWMXrJkHKMR9tok/Gov79AKWD+7lvYcv4+ACDM1wUz+gajSQNHKBQKg85RDIxH64hHju4sx/jx49G+fXvcunULvXv31p7YI488go8++sjYw4mGozsrx9FL1jF6yZByjEfbaJPxqLv/VoE93vrlNG5mFEAqAV7v3hQTe4XAXmb9Mcp4tI545OjOCkRFRSEqKkpnW5s2bbBq1So8/vjjphxSdBzdqY+jl6xj9JIh5RiPttEm4xFQqtT49u+7+PZIKlRqAY3cnbFwRKTNrRDAeLSOeKztozuNTtLGjh1b7vYbN27gyJEjeOedd4w9JBER1QE37+djypqTOH4zEwDwVGRDzHoqHHJne3E7RmSljE7SHjx4oPNapVLh2rVrSE5OxpdffmmxjhERUe0gCAJ+PXEbCb+fQ26REvUcpJg9KByDHwsQu2tEVs3oJG39+vXlbv/oo4+wYcMGvPrqq2Z3ioiIaoes/BK8t+EMNp2+CwBoG9gA7/VshMhmjUTuGZH1M3tZKI1nnnnGpiez5RQcujjE3DqGmBtSjvFoG23WxXj8+9p9xK89jbtZhZBJJZjUMwSvdA3Cg4z7jEeR26wN8cgpOIxw6tQptGnTxlKHq3acgqNyHGJuHUPMDSnHeLSNNutSPN7LyMQXB25j1fE0CAAayx0xo28wwv3rIeP+PcajFbRZG+KRU3CUQ7MAellpaWn47bff0L9/f539CxYsMPbwNYZTcFSOQ8ytY4i5IeUYj7bRZl2JxytpOXjnz4u4kJ4PABge1Rgf9A9FPUc7k8/D2jAerSMeOQVHOU6ePFnu9nbt2iE9PR3p6ekASoe12hJOwaGPQ8ytY4i5IeUYj7bRZm2OR0EQsProLczamIyCEhXkzvb4ZHAr9Gvlb5HzsDaMR+uIR07B8ZDdu3cbW4WIiGqxjLxivLvuNLYnpwEAogJcsfjZtmjUoJ7IPSOybRZ7Jo2IiOqefZcUeHPtKShyimAvk+DtPs3x5KMu8JM7i901Iptn0D23vn374vDhw1WWy8nJwaefforExESzO0ZERNarqESFWRuTMerbI1DkFCHEpz42xHXGS12DIbWxx12IrJVBd9KGDRuGIUOGQC6XY8CAAYiKikLDhg3h5OSEBw8eIDk5GQcOHMCWLVvQv39/zJs3r7r7bXGcgkMXh5hbxxBzQ8oxHm2jzdoUj1cU+Zj900VcTMsFADzfoQmm9WsBZwcZ49FG2qwN8cgpOP41btw4PP/881i7di3WrFmDZcuWISsrC0DpQ3dhYWGIjY3F0aNHERoaalqvaxin4Kgch5hbxxBzQ8oxHm2jzdoQj4IgYM3JNCQeuIMStYAGznZ4r3cgujzijpzM+8gx4DiMR+toszbEI6fgKMPR0RHPP/88nn/+eQBAVlYWCgoK4OnpCXt721t3jVNwVI5DzK1jiLkh5RiPttGmrcejIqcI7647jb2X7gEAoh/1wqdDWsPb1dGo4zAeraNNW49HS5SplVNwaMjlcsjlclOrWx1OwaGPQ8ytY4i5IeUYj7bRpq3G4/bkNLy77jQy8orhaCfFG10b4fWYlpDJZCa1yXi0jjZtNR4tWabWTcFBRER1Q0GxCh9uTsYPf98EALTwc8WiERFwlxTY3FyYRLaISRoREek5ezsLE1efxDVFHgDg5a7BeCu2OeylEqSnF4jcO6K6gUkaERFpqdUClu2/ivl/XkSJSoCPqyMWDI9El2Ze/+633VGZRLaGSRoRUR2kUgs4kpKB9JxC+Lg6ISrQHWk5xZjy+xEcupYBAIht6YtPBrdGg3oOIveWqG4yOkkbPXo0xo0bh27dulVHf4iIqJptO3sXMzcm425WoXabu7M9CktUKFSq4Wwvw4yBYRgeFcBnz4hEZPRwhqysLMTExKBZs2b4+OOPcfv27eroFxERVYNtZ+/i9VUndBI0AMgsKEGhUo1ATxdsmdQVI9o1YYJGJDKj76Rt2LABCoUCK1euxHfffYeEhATExMRg3LhxeOqpp2xyzjSAKw48jDNqW8eM2oaUYzzaRpvWEI8qtYAZvydDqKROUYkKjd2dqi3eGI/W0aY1xKO5ZbniQAW8vb0RHx+P+Ph4nDhxAsuXL8cLL7yA+vXr4/nnn8f48ePRrFkzUw5dY7jiQOU4o7Z1zKhtSDnGo220aQ3xePxWDlKzCyutk5pdhD9PXkPbAFeT+sF4tI02rSEezS3LFQeqcPfuXWzfvh3bt2+HTCbDE088gTNnziAsLAxz587FlClTzDl8teKKA5XjjNrWMaO2IeUYj7bRpjXEY8ldpUH1Suyc4ePjY1I/GI+20aY1xKO5ZbniQDlKSkrw+++/Y/ny5fjzzz/RunVrTJ48Gc8++6w2yVm/fj3Gjh1r1Unaw7jigD7OqG0dM2obUo7xaBttihmPEokEl/5dEL0qvm7OZs3yzni0jTZrw/WRKw48xN/fH2q1Gs888wyOHDmCyMhIvTI9evSAu7u7sYcmIqJqkF1Qgg9+T8bGU3eqLOsvd0L7YI8a6BURVcXoJG3hwoUYNmxYpbfr3N3dkZKSYlbHiIjIfCf+ycHs7edwN6sQMqkE/cP9sPH0XQDQGUAg+ff1B/1DIZNyVCeRNTA6SRs4cCDy8/P1krSMjAzY2dlpv/IkIiLxFCvVWLD9Ir7aew0CgEBPFywaEYk2TRrgidb686T5yZ0wsWtD9A33E6/TRKTD6CRt5MiRGDBgAMaPH6+z/eeff8bvv/+OLVu2WKxzRERkvKuKXExenYQzt7MAAMPaNsaMgS1Rz7H0kt833B+9w/z0Vhy4f08hZreJ6CFGJ2l///03FixYoLc9Ojoa7733nkU6RURExhMEAauP3sKsjckoKFFB7myPd3sGYGTn5noPK8ukEnRq6ql9bcvzlhHVVkYnaUVFRVAq9Ydxl5SUoKCgwCKdIiIi42TkFePddaexPTkNAPB4U0/MG9oKsiLD52QiIuti9JjT9u3bY9myZXrbly5dirZt21qkU0REZLh9lxSIXbQP25PTYC+T4L0nQrFqXAf4y53F7hoRmcHoO2kffvghYmJicOrUKfTq1QsAsHPnThw9ehR//vmnxTtYU7gslC4ue2Idy54YUo7xaBttVkc8FpWoMPfPS1j+13UAQIh3PSwaEYmwhm4ABMZjNWE8Wsf1kctClaNz5844dOgQ5s2bh59//hnOzs5o3bo1/vvf/1r9UlBlcVmoynHZE+tY9sSQcoxH22jT0vF49V4Bpm9LwdV7pY+ZDGntjTe6NoaTXSHS0w27lpnSP8Yj49HU+lwWqlS1LwsVGRmJH374wZSqVoPLQlWOy55Yx7InhpRjPNpGm5aKRy8vL6z8+xY+2XYRxUo1POs54JMhrdCrhf4yTozH6sF4tI7rI5eFqoBarcaVK1eQnp6ud9uuW7duphxSdFwWSh+XPbGOZU8MKcd4tI02zT1eRr4SU1eewN5L9wAAPZp7Y+7QCHi7OlqkTcaj4RiP1nF95LJQDzl8+DCeffZZ3LhxA4Ig6OyTSCTarw+JiMhydpxPwzu/JCOzQAlHOyne6x+KFzoGQiLh6gBEtZXRSdprr72GqKgobN68Gf7+/rxAEBFVo4JiFT7cnIwf/r4JAGjh54olz7TBo76uIveMiKqb0Una5cuX8csvvyAkJKQ6+kNERP86ezsLE1efxDVFHgDg2cd88cGgCDg72IvcMyKqCUYnaR06dMCVK1eYpBERVROVWsDX+69h/p8XUaIS4OvmiHlDW+NRNzUc7WRid4+IaojRSdobb7yBN998E6mpqWjVqhXs7XX/omvdurXFOkdEVNfcySxA/M9JOHwtAwAQ29IXnwxuDbmzHdLT00XuHRHVJKOTtCFDhgAAxo4dq90mkUggCAIHDhARmWHz6buY9utpZBcq4Wwvw4yBYRgeFQCJRGLTk78SkWmMTtJSUlKqox9ERHVWbpESCb+dw7oT/wAAIhrLsWhkGwR71RO5Z0QkJqOTtMDAwOroBxFRnXT8xgNMWZOEmxn5kEqA8dEhmBTTDPYy251DjIgsw6SrwMqVK9G5c2c0bNgQN27cAAAsWrQIv/32m0U7R0RUWylVaizecRnDvzqEmxn5aOTujNWvdMJbsc2ZoBERABOStC+//BLx8fF44oknkJmZqX0Gzd3dHYsWLbJ0/4iIap2b9/MxYtlhLNxxCSq1gKciG2LLpK5oH+whdteIyIoY/XXn559/jq+//hqDBg3CJ598ot0eFRWFt956y6Kdq0lqtVr7YK5arYYgCBU+qGvuflsgxjlYuk1zj2dKfWPqGFqW8Vh74lGtVmPd8VuYuek8cotUqO9oh1lPhWFQZCNtGUv2h/FYPWpLPNr69dHcMmLFojHtmTRwoE2bNnrbHR0dkZeXZ+zhRJOYmIjExETtnUCFQoHCwkIApW9gVlYWBEGocIFgc/bbAjHOwdJtmns8U+obU8fQsozH2hGPmfnF+PjPa9h3vfQ6GdGwPhJig9BQbm/Q1BqMR+tRG+KxNlwfzS0jVizm5OQYXNboJC04OBhJSUl6Awi2bduG0NBQYw8nmri4OMTFxSE7OxtyuRze3t5wc3MDUPrBSSQSeHt7V3iRMWe/LRDjHCzdprnHM6W+MXUMLct4tP14/PvafcSvvYy7WYWQSSWY3CsEr3VvCpnU8GX1GI/Ww9bj0RLHs4Z4NLeMWLHo5ORkcFmjk7T4+HjExcWhsLAQgiDgyJEj+OmnnzBnzhx88803xh7OakilUp0PSSKR6G0ry9z9tkCMc7B0m+Yez5T6xtQxtCzj0TbjsVipxsIdl7B071UIAtBY7oglzz6GxwJNe/aM8Wg9bDEeLX08a4hHc8uI8Tka05bRSdpLL70EZ2dnvP/++8jPz8ezzz6Lhg0bYvHixRg5cqSxhyMiqpWuKnIxeXUSztzOAgAMj2qM19p7Iaixu7gdIyKbYXSSBgDPPfccnnvuOeTn5yM3Nxc+Pj6W7hcRkU0SBAE/HbmF2ZuSUVCigtzZHp8MboXYlr5c1omIjGL0/b2ePXsiMzMTAODi4qJN0LKzs9GzZ0+Ldo6IyJZk5BXjlZXH8X/rz6CgRIXOIZ74Y3I39GvlL3bXiMgGGX0nbc+ePSguLtbbXlhYiP3791ukU0RE1kylFnAkJQPpOYXwcXVC+2AP/HXlHt5cewqKnCLYyyR4J7YFxnUJhtSIwQFERGUZnKSdPn1a+//JyclITU3VvlapVNi2bRsaNWpk2d4REVmZbWfvYubGZNzNKtRuc3GQIb+4dDqfEJ/6WDwyEi0bysXqIhHVEgYnaZGRkZBIJJBIJOV+rens7IzPP//cop0jIrIm287exeurTkB4aLsmQev+qDeWPt8Wzg6ymu8cEdU6BidpKSkpEAQBjzzyCI4cOQJvb2/tPgcHB/j4+EAm44WJiGonlVrAzI3JeglaWZfScuBgZ7vTShCRdTE4SdNMXmvLS3kQEZnqSEqGzlec5bmbVYgjKRno1NSzhnpFRLWZSVNwAKXPpd28eVNvEMHAgQPN7hQRkbVJz6k8QTO2HBFRVYxO0q5du4ann34aZ86cgUQigSCU3vyXSEpHMGnWwiQiqi0KilXYeOqOQWV9XA1f8oWIqDJGPzwxadIkBAcHIz09HS4uLjh37hz27duHqKgo7Nmzpxq6SEQknrO3s9D/8/3Ycb7yiWglAPzlpdNxEBFZgtFJ2qFDhzBr1ix4eXlp17vq0qUL5syZg4kTJ1ZHH4mIapxKLWDlsVQMWXoI1xR58HVzxKRezSBBaUJWluZ1woAwoxZNJyKqjNFfd6pUKri6ugIAvLy8cOfOHTRv3hyBgYG4ePGixTtIRFTT7mQWIH5NEg6nZAAA+rb0w5zBrdCgngNC/V315knzkzshYUAY+oZzZQEishyjk7Tw8HCcOnUKwcHB6NChA+bOnQsHBwcsW7YMjzzySHX0kYioxmw+fRfTfj2N7EIlnO2lSBgQhhHtmmifu+0b7o/eYX56Kw7wDhoRWZrRSdr777+PvLw8AMCsWbPw5JNPomvXrvD09MSaNWss3kEiopqQW6REwm/nsO7EPwCA1o3leL9XY0Q1D9AmaBoyqYTTbBBRtTM6SYuNjdX+f0hICC5cuICMjAw0aNBA70JGRGQLjt94gClrknAzIx9SCTA+OgRv9GyKB/fvid01IqrDTJ4nrSwPD45mIiLbo1Spkbj7KpbsugyVWkAjd2csHBGJ9sEenLibiERnUJI2ePBggw/466+/mtwZManVau1FWa1WQxCECi/S5u63BWKcg6XbNPd4ptQ3po6hZRmP1XMONzPy8ebPp3D8ZiYAYGCEP2YNbAk3Z3vt9YDxaHw5xqNttFkb4tHcMmLFojHtGZSkyeVykztjrRITE5GYmKidfFehUKCwsHS0llqtRlZWFgRBgFSqP0uJufttgRjnYOk2zT2eKfWNqWNoWcajZc9BEARsPZ+Bz/bcRH6xGvUcpHi7ZxP0beGJwpwHKMyxfJuWOB7j0Xrw+mgd8WhuGbFiMScnx+CyBiVpy5cvN7kz1iouLg5xcXHIzs6GXC6Ht7c33NzcAJR+cBKJBN7e3hVeZMzZbwvEOAdLt2nu8Uypb0wdQ8syHi13DlkFJfhgwzlsOnMXABAV2AALhrdG4wYu1dampY7HeLQevD5aRzyaW0asWHRyMnxVEos8k1YbaCbm1ZBIJHrbyjJ3vy0Q4xws3aa5xzOlvjF1DC3LeDT/HA5fu4/4NUm4k1UImVSCKTHN8Hp0SKVTZzAeTSvHeLSNNmtDPJpbRozP0Zi2mKQRUa1WrFRj4Y5LWLr3KgQBCPR0waIRkWjTpIHYXSMiqhSTNCKqta4qcjF5dRLO3M4CAIyICsD0AWGo58hLHxFZP16piKjWEQQBq4/ewqyNySgoUUHubI9PBrdCv1ZctomIbAeTNCKqVTLyivHuutPYnpwGAHi8qSfmD4+Av9xZ5J4RERnHoCRtyZIlBh9w4sSJJneGiMgc+y4p8ObaU1DkFMFeJsE7sS0wrkswpFxXk4hskEFJ2sKFCw06mEQiYZJGRDWusESFudsu4tu/UgAAIT71sXhkJFo2rH1zPBJR3WFQkpaSklLd/SAiMsnF1BxMWn0SF1JLJ4gc1SkQ0/qFwtlBJnLPiIjMw2fSiMgmCYKAFQevY87WCyhWquFZzwFzh7ZGr1BfsbtGRGQRBiVp8fHxBh9wwYIFJneGiMgQ6TmFeHvtaey9pAAA9GjujblDI+Dt6ihyz4iILMegJO3kyZMGHUwi4cO5RFS9tien4d11p5GRVwxHOyne6x+KFzoG8vpDRLWOQUna7t27q7sfRESVKixR4/0NZ/HjkVsAgBZ+rljyTBs86usqcs+IiKoHn0kjIqt39nYW3vgxGTceFAEAXu4ajLdim8PRjoMDiKj2MilJO3bsGH7++WfcvHkTxcXFOvt+/fVXi3SMiEilFvD1/muY/+dFlKgE+Lo5Yv6wSHRp5iV214iIqp3Ry76vXr0ajz/+OM6fP4/169ejpKQE586dw65duyCXc04iIrKMO5kFeO6bw/hk6wWUqAR0b+qOLRO7MEEjojrD6DtpH3/8MRYuXIi4uDi4urpi8eLFCA4Oxquvvgp/f66LR0Tm23z6Lqb9ehrZhUo428swfUAoogMc0MDFQeyuERHVGKPvpF29ehX9+/cHADg4OCAvLw8SiQRTpkzBsmXLLN5BIqo7couUePPnU4j78QSyC5WIaCzHlkldMSIqgKM3iajOMfpOWoMGDZCTUzqzd6NGjXD27Fm0atUKmZmZyM/Pt3gHiahuOHHzASavTsLNjHxIJEBcdAgmxTSDvUwKtVotdveIiGqc0Ulat27dsH37drRq1QrDhg3DpEmTsGvXLmzfvh29evWqjj4SUS2iUgs4kpKB9JxC+Lg64bEm7li69xqW7LoMlVpAI3dnLBwRifbBHmJ3lYhIVEYnaf/5z39QWFgIAHjvvfdgb2+PgwcPYsiQIXj//fct3kEiqj22nb2LmRuTcTerULvNXiZBiUoAADwV2RCzngqH3NlerC4SEVkNo5M0D4///XUrlUoxdepUi3aIiGqnbWfv4vVVJyA8tF2ToI3pHISEAS1rvmNERFbK6IEDZfXv3x937961VF+IqJZSqQXM3Jisl6CVte1sKlTqykoQEdUtZiVp+/btQ0FBgaX6QkS11JGUDJ2vOMtzN6sQR1IyaqhHRETWz6wkjYjIEHczDftjLj2n8kSOiKguMStJCwwMhL09H/AloopdVeTi891XDCrr4+pUzb0hIrIdZi2wfvbsWUv1g4hqGUEQsProLczamIyCEhUkQIXPpEkA+MmdOO0GEVEZJidpx48fx/nz5wEAYWFheOyxxyzWKSKybRl5xXh33WlsT04DADze1BMDIvzxf7+W/mFXNlnTrCOQMCAMMilXFSAi0jA6SUtPT8fIkSOxZ88euLu7AwAyMzPRo0cPrF69Gt7e3pbuY5Wefvpp7NmzB7169cIvv/xS4+0T0f/su6TAm2tPQZFTBHuZBO/EtsC4LsGQSiVo4OKgN0+an9wJCQPC0Deca/8SEZVldJL2xhtvICcnB+fOnUNoaCgAIDk5GaNHj8bEiRPx008/WbyTVZk0aRLGjh2L7777rsbbJqJShSUqzN12Ed/+lQIACPGpj8UjI9GyoVxbpm+4P3qH+emsONA+2IN30IiIymF0krZt2zbs2LFDm6ABpV93JiYmok+fPhbtnKGio6OxZ88eUdomIuBiag4mrT6JC6ml6/q+0DEQ//dEKJwdZHplZVIJOjX1rOkuEhHZHKNHd6rV6nJHdNrb25u0CPK+ffswYMAANGzYEBKJBBs2bNArk5iYiKCgIDg5OaFDhw44cuSI0e0QkeUJgoDlf6VgwH8O4EJqDjzrOeC/o6Mwe1B4uQkaEREZzugkrWfPnpg0aRLu3Lmj3Xb79m1MmTLFpAXW8/LyEBERgcTExHL3r1mzBvHx8UhISMCJEycQERGB2NhYpKenG90WEVlOek4hXlx+FDM3JqNYqUaP5t7YNrkbeoX6it01IqJawaQF1gcOHIigoCAEBAQAAG7duoXw8HCsWrXK6A7069cP/fr1q3D/ggUL8PLLL2PMmDEAgKVLl2Lz5s349ttvTVo3tKioCEVFRdrX2dnZAErvEGruBKrVagiCUOGdQXP32wIxzsHSbZp7PFPqG1PH0LLWGI87zqdh6rozyMgvgaOdFNP6tcALHZtAIpFUS8wwHhmP1oTxaB3xaG4ZsWLRmPaMTtICAgJw4sQJ7NixAxcuXAAAhIaGIiYmxthDVam4uBjHjx/HtGnTtNukUiliYmJw6NAhk445Z84czJw5U2+7QqFAYWHpiDO1Wo2srCwIggCpVP9mo7n7bYEY52DpNs09nin1jaljaFlrisfCEjUW77uF9WfuAQBCvJwxq18wHvF0hkKhMOvYlWE8Mh6tCePROuLR3DJixWJOTo7BZY1O0r7//nuMGDECvXv3Ru/evbXbi4uLsXr1aowaNcrYQ1bo3r17UKlU8PXV/frE19dXmyACQExMDE6dOoW8vDw0btwYa9euRadOnco95rRp0xAfH699nZ2djYCAAHh7e8PNzQ1A6QcnkUjg7e1d4UXGnP22QIxzsHSb5h7PlPrG1DG0rLXE49nbWZi85hSu3csDALzUJRhv9mkGR7vqf/aM8ch4tCaMR+uIR3PLiBWLTk6Gr6xidJI2ZswY9O3bFz4+Pjrbc3JyMGbMGIsmaYbasWOHwWUdHR3h6Oiot10qlep8SBKJRG9bWebutwVinIOl2zT3eKbUN6aOoWXFjEeVWsDX+69h/p8XUaIS4OvmiPnDItGlmZfRxzIH45HxaE0Yj9YRj+aWEeNzNKYto5M0QRAgkejPafTPP/9ALpeXU8N0Xl5ekMlkSEtL09melpYGPz8/i7ZFRPruZBYg/uckHL6WAQCIbemLTwa3RoN6DiL3jIio9jM4SWvTpg0kEgkkEgl69eoFO7v/VVWpVEhJSUHfvn0t2jkHBwe0bdsWO3fuxKBBgwCU3p7cuXMnJkyYYNG2OHBAFx+MtY4HYw0pV13xuOXMXfzf+rPILlTC2V6G6QNCMbxt42obHFAZxiPj0ZowHq0jHjlwoAxNkpSUlITY2FjUr19fu8/BwQFBQUEYMmSI4b38V25uLq5cuaJ9nZKSgqSkJHh4eKBJkyaIj4/H6NGjERUVhfbt22PRokXIy8vTjvY0VWJiIhITE6FSqQBw4MDD+GCsdTwYa0g5S8djXrEK83ffwpbz9wEAYb4umNE3GE0aOFbr4IDKMB7rbjxaI8ajdcQjBw6UkZCQAAAICgrCiBEjjHrwrTLHjh1Djx49tK81D/WPHj0aK1aswIgRI6BQKDB9+nSkpqYiMjIS27Zt0xtMYKy4uDjExcUhOzsbcrmcAwcewgdjrePBWEPKWTIeT9x8gPifk3EzowBSCfB696aY2CsE9jJx45jxWDfj0VoxHq0jHjlwoByjR4/W/v/48eMxa9YseHmZ/gBxdHQ0BEGotMyECRMs/vXmwzhwQB8fjLWOB2MNKWfufqVKjcTdV7Fk12Wo1AIauTtj4YhItA/2qPIcagrjse7Eoy1gPFpHPNb2gQNm9WrVqlXayWCJyDbdvJ+PEcsOY+GOS1CpBTwV2RBbJnW1qgSNiKguMvpOWllV3QEjIuslCAJ+PXEbCb+fQ26REq6Odpg9KByD2jQSu2tERAQzk7TahKM7dXH0knWMXjKknCn7swpK8MGGc9h05i4AICqwARYMb43GDVysMm4Zj7U7Hm0N49E64pGjO6tgzAgFa8PRnZXj6CXrGL1kSDlj95/4Jwcz/0hBWk4JZBLgpU4NMSrKD7KSXKSn5xp0njWN8Vh749EWMR6tIx45uvNfxjx3phkhae04urNyHL1kHaOXDCln6H55A08s2X0VX+27BkEAAj1csHBEBCID3A06NzExHmtfPPL6KG6btSEeObrzX+7u7uWuMlCWZiUCzZ0pW8PRnfo4esk6Ri8ZUq6q/TcfFGH2L3/j7O3SP7hGRAVg+oAw1HO0nSceGI+1Jx55fbSONmtDPNb20Z0GXaF3795tcmeISDyCIOCnIzcxe9N5FCrVkDvb45PBrdCvlb/YXSMioioYlKR17969uvtBRBaWkVeMd9edxvbk0rVvH2/qiQXDI+Ent8xE1EREVL0MStJOnz5t8AFbt25tcmeIyDL2XVLgzbWnoMgpgr1Mgtceb4RJseGws5OJ3TUiIjKQQUlaZGQkJBJJlfOi2fIzaZyCQxeHmFvHEHNDypXdX1Siwtw/L2H5X9cBACHe9bBgeGt42xUBsN2YZDzaZjya2ydrxXi0jnjkFBz/SklJMbkz1opTcFSOQ8ytY4i5IeU0+68o8jHjzxu4eq8AADA0whsTujSGg6wQmZmMR7HbrGvxyOujdbdZG+KRU3D8KzAw0OTOWCtOwVE5DjG3jiHmhpRTqVRYe0qBxL+uo1iphmc9B3w6pBV6tvAx+TysDePRduKR10fbaLM2xCOn4KhEcnIybt68ieLiYp3tAwcONPWQouIUHPo4xNw6hphrygmQ4O+UB0jPKYSPqxPaB3vgfl4R3l57Cnsv3QMA9GjujblDI+Dt6mj2eVgbxqN1xSOvj4xHa4hHTsHxkGvXruHpp5/GmTNndJ5T08yjZqvPpBFZs91XHmDxt+eQml2o3dbAxR4lKgG5RUo4yiSY9kQoRj8eVOWchkREZBuMTh0nTZqE4OBgpKenw8XFBefOncO+ffsQFRWFPXv2VEMXieq2bWdTMW3TNZ0EDQAe5Jcgt0iJRu5OWPFsKEZ1CmSCRkRUixh9J+3QoUPYtWsXvLy8tLcIu3Tpgjlz5mDixIk4efJkdfSTqE5SqQXM2nS+yjJNGnDuMyKi2sboO2kqlQqurq4AAC8vL9y5cwdA6eCCixcvWrZ3RHXckZQMvTtoD0vNLkLSbetcGJ2IiExn9J208PBwnDp1CsHBwejQoQPmzp0LBwcHLFu2DI888kh19LFGcJ40XZwHyDrmAUrLLjCo3Xt5xYxHK2+zNsSjIeV4fbSNNmtDPHKetHK8//77yMvLAwDMmjULTz75JLp27QpPT0+sWbPG2MOJhvOkVY7zAFnHPEBnrqcb1K6TUIT09HTGoxW3WRvi0ZByvD7aRpu1IR45T1o5YmNjtf8fEhKCCxcuICMjAw0aNLCph5Y5T1rlOA+QuPMA5RYpMXNjMtaduFvpMSQA/ORO6PyoH3x8fBiPVtymLcejMeV4fbSNNmtDPHKeNAN5eHhY4jCi4jxp+jgPkDjzAB2/8QBT1iThZkY+pBKgT5gvtp1LgwRA2YXZNH8SfdA/FHYyKePRBtq0xXg0pRyvj7bRZm2IR86TRkQ1QqlS4z+7r+DzXVegUgto5O6MhSMiERXojjUHL2Lxvjs6gwj85E5IGBCGPmG+SE837GtRIiKyHUzSiKzAzfv5mLzmJE7czAQAPBXZELMHhcPNyR5qtRo9QhpgaMdHcexGps6KAzKpxKYfwCYioooxSSMSkSAI+PXEbczYmIzcIiVcHe0we1A4BrVppFdWJpWgU1NPEXpJRERiYJJGJJKsghJ8sDUFOy49AAC0C2qABcMjEeDhInLPiIjIGjBJIxLB4Wv3MWVNEu5mFUImlWBKTDO8Hh0CmdR2RkgTEVH1YpL2L05mq4uTNVbPZI3FSjUW7byMr/ZdgyAAjeSOWDwyEo8FegAQoFYLenUYj4xHU+tzMtvqwXi0jnjkZLa1GCezrRwna7T8ZI03MgqRsC0FF9LzAQADwjwxJtIVfo7FlY7OZDwyHk2tz8lsqwfj0TrikZPZ1mKczLZynKzRcpM1CoKA1Udv4cPNF1BQooLc2R4fPx2O2DAfKBQKTh5qAMajdUweakg5xqNttFkb4pGT2dYhnMxWHydrNH+yxoy8Yry77jS2J6cBADqHeGL+sEj4yZ20FwhOHmoYxqN1TB5qSDnGo220WRvikZPZEpFJ9l1S4M21p6DIKYK9TIJ3YltgXJdgSDk4gIiIDMAkjcjCipRqzN58Hsv/ug4ACPGpj8UjI9GyoVzcjhERkU1hkkZkQRdTczBh9QVcvVcAABjVKRDT+oXC2UEmcs+IiMjWMEkjsgBBELDi4HXM2XoBxUo1POs5YO7Q1ugV6it214iIyEYxSSMyU3pOId5eexp7LykAAI8HuWHhM1HwlTuL3DMiIrJlTNKIzLA9OQ3vrjuNjLxiONpJMa1fC8Q+4gRvV0exu0ZERDaOSRqRCQqKVfhwczJ++PsmACDU3w1LRkaiqXe9SiemJSIiMhSTtH9xWShdXPak4vpnb2dh8ppTuHYvDwDwUpdgvNmnGRztZFyGp5owHq1jGR5DyjEebaPN2hCPXBaqFuOyUJXjsif69VVqAT+eSMNXB+9AqRbgXc8eH8QGoX0TN2Rl3De6TS7DYzjGo3Usw2NIOcajbbRZG+KRy0LVYlwWqnJc9kS3fmp2Ed5aexqHUzIAAH3CfDFncDgauDiY3CaX4TEc49E6luExpBzj0TbarA3xyGWh6hAuC6WPy56U1t92Lg3/t/4ssguVcLaXYcbAMAyPCoBEUv7KAVyGp3owHq1jGR5DyjEebaPN2hCPXBaKqI7KLVJi1h/XseV86VeZEY3lWDSyDYK96oncMyIiqguYpBGV4/iNB5iy5iRuZhRAKgHGR4dgUkwz2Mts9y9/IiKyLUzSiMpQqtRI3H0VS3ZdhkotwM/VAYtGtkHHpl5id42IiOoYJmlE/7p5Px9Tfk7C8RsPAAADI/zxRicfNG3iIXLPiIioLmKSRnWeIAj49cRtJPx+DrlFSrg62mH2oHAMjPDnxLRERCQaJmlUp2Xll+C9DWew6fRdAEBUYAMsHBGJAA8Xm55sk4iIbB+TNKoTVGoBR1IykJ5TCB9XJ7QP9sDR6xmIX5OEO1mFkEklmNyrGV6Pbgo7Dg4gIiIrwCSNar1tZ1Mxe/N53M0q1G6r5yhDXlHpahOBni5YNCISbZo0EKuLREREepikUa22+8oD/N+maxAe2q5J0Do39cSyUVGo58gfBSIisi78XodqLZVawMI9t/QStLKu3cuDk72sxvpERERkKCZpVGsdvZ6B9NySSsvczSrEkX/X4yQiIrIm/I7nX2q1WjuaT61WQxCECkf3mbvfFohxDpZuM63MM2iVlssuKLdNU/pjTB1DyzIea0c8mns8xqP1YDxaRzyaW0asWDSmvTqbpCUmJiIxMREqVemzSQqFAoWFpb/U1Wo1srKyIAhCuQuhmrvfFohxDpZss0ipxu8nbhhU1l5ZUO58aKb0x5g6hpZlPNp+PFrieIxH68F4tI54NLeMWLGYk5NjcNk6m6TFxcUhLi4O2dnZkMvl8Pb2hpubG4DSD04ikcDb27vCi4w5+22BGOdgqTYvpuZg8i+ncDG18h8ECQA/uRP6tHkEMqnEIv0xpo6hZRmPth2Pljoe49F6MB6tIx7NLSNWLDo5ORlcts4maQ+TSqU6H5JEItHbVpa5+22BGOdgTpuCIGDFweuYs/UCipVqeNZzwJOhDfD9sbTS/WXb+fe/CQPCYG9X8cABU/pjTB1DyzIebS8eq+N4jEfrwXi0jng0t4wYn6MxbTFJo1ohPacQb689jb2XFACAHs298cngVhAKstDh0YZ686T5yZ2QMCAMfcP9xeoyERFRpZikkc3bnpyGd9edRkZeMRztpHivfyhe6BgIQRCQXgD0DfdDbLi/3ooD5X3FSUREZC2YpJHNKihW4cPNyfjh75sAgBZ+rljyTBs86usKoPTrTw2ZVIJOTT1F6ScREZEpmKSRTTp7OwsTV5/ENUUeAODlrsF4K7Y5HCt5voyIiMiWMEkjm6JWC1i2/xrm/3kRJSoBvm6OmD8sEl2aeYndNSIiIotikkY2405mAd78+RQOXbsPAIht6YtPBrdGg3oOIveMiIjI8pikkU3YfPoupv16GtmFSjjbyzBjYBiGRwVAIuHD/0REVDsxSSOrllukRMJv57DuxD8AgIjGciwa2QbBXvVE7hkREVH1YpJGVuvEzQeYvDoJNzPyIZUA46NDMCmmGexltjsBJhERkaGYpJHVUarU+HLXVSzZdRkqtYBG7s5YOCIS7YM9xO4aERFRjWGSRlbldlYR4n79G8dvZgIAnopsiFlPhUPubC9ux4iIiGoYkzSyCoIgYP3J2/jgt2TkF6vh6miH2YPCMahNI7G7RkREJAomaSS6rIISvL/hLDaeugMAiApsgIUjIhHg4SJyz4iIiMTDJI1EdfjafcSvScKdrELIpBK81NEfbz3RGvZcOYCIiOo4JmkkimKlGgt3XMLSvVchCECgpwsWDo9AQ8diLnxOREQEJmkkgquKXExenYQzt7MAACOiAjB9QBic7aVIT08XuXdERETWgUnav9RqNdRqtfb/BUHQvi6vrDn7bUF1nIMgCFhz7B/M3nQeBSUqyJ3t8fHT4egX7lctbZp7PFPqG1PH0LKMR3HOgfFoWjnGo220WRvi0dwyYsWiMe3V2SQtMTERiYmJUKlUAACFQoHCwkIApW9gVlYWBEGAVKo/caq5+22Bpc8hs0CJj7dfx75rpXfPogJcMT02CD71/3f3zNJtmns8U+obU8fQsoxHcc6B8WhaOcajbbRZG+LR3DJixWJOTo7BZetskhYXF4e4uDhkZ2dDLpfD29sbbm5uAEo/OIlEAm9v7wovMubstwWWPIf9lxV465cLUOQUwV4mwdt9mmNs5yBIH3r2zNLvm7nHM6W+MXUMLct4FOccGI+mlWM82kabtSEezS0jViw6OTkZXLbOJmkPk0qlOh+SRCLR21aWufttgbnnUFiiwtxtF/HtXykAgBCf+lg8MhItG8qrrU1LH8+U+sbUMbQs41Gcc2A8mlaO8WgbbdaGeDS3jBifozFtMUmjanExNQeTVp/EhdTS27qjOgViWr9QODtwag0iIiJDMEkjixIEASsOXsecrRdQrFTDq74D5g5tjZ4tfMXuGhERkU1hkkYWk55TiLfXnsbeSwoAQI/m3pg7NALero4i94yIiMj2MEkji9iRnIZ31p1GRl4xHO2keK9/KF7oGAiJhBPTEhERmYJJGpmloFiFDzcn44e/bwIAQv3dsGRkJJr5uorcMyIiItvGJI1MdvZ2FiauPolrijwAwMtdg/FWbHM4ct1NIiIiszFJoyqp1AKOpGQgPacQPq5OiApsgP/+lYL5f15EiUqAr5sj5g+LRJdmXmJ3lYiIqNZgkkaV2nY2FbM3n8fdrELtNgeZFMWq0mUt+rb0w5zBrdCgnoNYXSQiIqqVmKRRhXZfeYD/23QNwkPbNQnaCx0DMeuplhwcQEREVA1sd7pnqlYqtYCFe27pJWhl7TifBnVlBYiIiMhkTNKoXEevZyA9t6TSMnezCnEkJaOGekRERFS3MEmjcqWWeQatMuk5hpUjIiIi4zBJIz23MvKxdO9Vg8r6uDpVc2+IiIjqJg4cIC1BELD+5G1M/+0ccouUkAAVPpMmAeAnd0L7YI8a7CEREVHdwSSNAABZBSV4f8NZbDx1BwAQFdgAPR+pj3m7bwHQTdY0YzkTBoRBJuXITiIiourAJI1w+Np9xK9Jwp2sQsikEkyJaYZXuz2C+/cUCPL30psnzU/uhIQBYegb7i9ir4mIiGo3Jml1WLFSjUU7LuHLvVchCECgpwsWj2yDyAB3qNX/TlYb7ofYcH+dFQfaB3vwDhoREVE1Y5JWR11V5GLy6iScuZ0FABgRFYDpA8JQz1E/JGRSCTo19azpLhIREdVpTNLqGEEQsProLczamIyCEhXkzvb4ZHAr9GvFry6JiIisCZO0OiQjrxjvrjuN7clpAIDOIZ6YPywSfnJOo0FERGRtmKTVEfsuKfDm2lNQ5BTBXibBO7EtMK5LMKR8toyIiMgqMUmr5QpLVJi77SK+/SsFABDiUx+LR0aiZUO5yD0jIiKiyjBJq8UupuZg0uqTuJCaAwAY1SkQ0/qFwtlBJnLPiIiIqCpM0mohQRCw4uB1zNl6AcVKNbzqO2Du0Nbo2cJX7K4RERGRgZik1TLpOYV4e+1p7L2kAAD0aO6NuUMj4O3qKHLPiIiIyBhM0mqRHclpeGfdaWTkFcPRTor3+ofihY6BkEg4OICIiMjWSMXugCVs2rQJzZs3R7NmzfDNN9+I3Z0aV1Cswnvrz+Cl748hI68Yof5u2PRGF4zqFMQEjYiIyEbZ/J00pVKJ+Ph47N69G3K5HG3btsXTTz8NT8+6MUP+2dtZmLj6JK4p8gAAr3R7BG/2eRSOdhwcQEREZMts/k7akSNH0LJlSzRq1Aj169dHv3798Oeff4rdrWqnVgtYuvcqnv7iL1xT5MHXzRGrxnXA/z0RygSNiIioFhA9Sdu3bx8GDBiAhg0bQiKRYMOGDXplEhMTERQUBCcnJ3To0AFHjhzR7rtz5w4aNWqkfd2oUSPcvn27JroumjuZBXjum7/xydYLKFEJ6NvSD9smdUOXZl5id42IiIgsRPQkLS8vDxEREUhMTCx3/5o1axAfH4+EhAScOHECERERiI2NRXp6eg331DpsPn0XfRftw6Fr9+HiIMPcIa3x5fOPoUE9B7G7RkRERBYk+jNp/fr1Q79+/Srcv2DBArz88ssYM2YMAGDp0qXYvHkzvv32W0ydOhUNGzbUuXN2+/ZttG/fvsLjFRUVoaioSPs6OzsbAKBWq6FWq7X/LwiC9vXDzN1vitwiJWZuTMa6E6XnGtFYjgXDIxDsVQ+CIEAQBIu1BVTPOdR0m+Yez5T6xtQxtKw1xmNNYzwyHq0J49E64tHcMmLFojHtiZ6kVaa4uBjHjx/HtGnTtNukUiliYmJw6NAhAED79u1x9uxZ3L59G3K5HFu3bsUHH3xQ4THnzJmDmTNn6m1XKBQoLCwEUPoGZmVlQRAESKX6NxvN3W+ss3dzkbAtBbeziiGVAKPb+WFch4awU+chPT3P7OOXx9LnIEab5h7PlPrG1DG0rLXFoxgYj4xHa8J4tI54NLeMWLGYk5NjcFmrTtLu3bsHlUoFX1/dmfJ9fX1x4cIFAICdnR3mz5+PHj16QK1W45133ql0ZOe0adMQHx+vfZ2dnY2AgAB4e3vDzc0NQOkHJ5FI4O3tXeFFxpz9hlKq1Phiz1V8vvsqVGoBjdydMX9Ya7QP9jD5mIay1DmI2aa5xzOlvjF1DC1rLfEoJsYj49GaMB6tIx7NLSNWLDo5ORlc1qqTNEMNHDgQAwcONKiso6MjHB31Z9+XSqU6H5JEItHbVpa5+6ty834+pvychOM3HgAABkU2xKxB4XBzsjfpeKYw9xysoU1zj2dKfWPqGFpW7Hi0BoxHxqM1YTxaRzyaW0aMz9GYtqw6SfPy8oJMJkNaWprO9rS0NPj5+YnUq+olCAJ+PXEbCb+fQ26REq6Odvjw6XA8Fdmo6spERERUa1h1kubg4IC2bdti586dGDRoEIDS25M7d+7EhAkTLNqWNQwcyCoowQcbzmHTmbsAgKjABlgwvDUaN3AR5cFGPhgr/oOxhpTjg9q20Sbj0fg+WSvGo3XEIwcO1IDc3FxcuXJF+zolJQVJSUnw8PBAkyZNEB8fj9GjRyMqKgrt27fHokWLkJeXpx3taarExEQkJiZCpVIBEH/gwIl/cjDzjxSk5ZRAJgFe6tQQo6L8ICvJRXp6rhlnaho+GGsdD8YaUo4PattGm4xH08/D2jAerSMeOXCgBhw7dgw9evTQvtY81D969GisWLECI0aMgEKhwPTp05GamorIyEhs27ZNbzCBseLi4hAXF4fs7GzI5XLRBg4UK9VYtPMyvtp3DYIABHq6YNHwCEQEuJt1fubig7HW8WCsIeX4oLZttMl4NP08rA3j0TrikQMHakB0dHSVc3xNmDDB4l9vPkyMgQNXFbmYvDoJZ25nAQBGRAVg+oAw1HMU/WMBwAdjTa3PB7WrB+OR8WhNGI/WEY8cOEAWJwgCVh+9hVkbk1FQooK7iz0+GdwKfcP9xe4aERERWQkmadVIpRbwd8p9pOcUwsfVCe2DPZBVUIJ3153G9uTSEaudQzwxf1gk/OSG3/4kIiKi2o9J2r8sPbpz9+UMLP72LFKz/7cEVQMXe6jUArILlXCQSfBWn+YY2zkIUqnE6kY6cfSSdYxeMqQcR9PZRpuMR+P7ZK0Yj9YRjxzdWYtV5+jOXZcy8H9bUvS2P8gvAQB417fH/IEheNTHBffuKSx5WhbD0UvWMXrJkHIcTWcbbTIeTT8Pa8N4tI545OjOWqy6Rneq1AKWfHu20rZlUik6hQVCJpVY7oQsjKOXrGP0kiHlOJrONtpkPJp+HtaG8Wgd8cjRnXWIpUZ3/p1yX+crzvKkZhfh2I1MdGpa8Rqj1oCjl6xj9JIh5TiazjbaZDwa3ydrxXi0jnis7aM7bfcnxEql5xRatBwRERHVTUzSLMzH1bDbmIaWIyIiorqJSZqFtQ/2gJ9bxQmYBIC/vHQ6DiIiIqKK8Jm0f1lqCg4JgA/6t0DcT0mQABAe2gcAH/QPhQQC1OrKV1oQE4eYW8cQc0PKccoD22iT8Wh8n6wV49E64pFTcNRi1TkFRxtvCd6P9sWyYxlIzy3Rbvepb4/J0QF4zEeK9PT0ajozy+AQc+sYYm5IOU55YBttMh5NPw9rw3i0jnjkFBy1WHUvsN4/QoIXekbg+M1MpOcUwcfVEe2CPKx62o2yOMTcOoaYG1KOUx7YRpuMR9PPw9owHq0jHjkFRx1iqSk4yu63t5Ph8RDvaulvTeAQc+sYYm5IOU55YBttMh6N75O1YjxaRzxyCg4iIiIiqnFM0oiIiIisEJM0IiIiIivEJI2IiIjICnHgwL8sNU+aIfttAecBso55gAwpx3i0jTYZj8b3yVoxHq0jHjlPWi1WnfOkcR4g62izNswDZEg5xqNttMl4NP08rA3j0TrikfOk1WLVPU8a5wESv83aMA+QIeUYj7bRJuPR9POwNoxH64hHzpNWh1THPGmcB0j8NmvDPECGlGM82kabjEfj+2StGI/WEY+cJ42IiIiIalydv5MmCKWLnGdnZ2u3qdVq5OTkwMnJqcLb9ebstwVinIOl2zT3eKbUN6aOoWUZj4xHU+szHqsH49E64tHcMmLFoibf0OQflanzSZrmAb6AgACRe0JERER1RU5ODuRyeaVlJIIhqVwtplarcefOHbi6ukIi+d/i5+3atcPRo0crrFfZ/uzsbAQEBODWrVvawQi2qKr3wBbaNPd4ptQ3po6hZRmPjEdT6zMeqwfj0Tri0ZwyYsWiIAjIyclBw4YNq7yDV+fvpEmlUjRu3Fhvu0wmq/RDq2o/ALi5udn0RciQc7T2Ns09nin1jaljaFnGI+PR1PqMx+rBeLSOeLREGTFisao7aBq2+UBADYiLizNrf20gxjlauk1zj2dKfWPqGFqW8ch4NLU+47F6MB6tIx4tVcZa1fmvO6uDZu61rKwsm/5LkWoHxiNZE8YjWQtbiEXeSasGjo6OSEhIgKOjo9hdIWI8klVhPJK1sIVY5J00IiIiIivEO2lEREREVohJGhEREZEVYpJGREREZIWYpBERERFZISZpRERERFaISVoNS0lJQY8ePRAWFoZWrVohLy9P7C5RHXXx4kVERkZq/zk7O2PDhg1id4vqsIULF6Jly5YICwvDxIkTDVqAmqi6fPbZZ2jZsiXCw8OxatUqUfrAKThqWPfu3fHhhx+ia9euyMjIgJubG+zs6vzqXCSy3NxcBAUF4caNG6hXr57Y3aE6SKFQoGPHjjh37hzs7e3RrVs3fPbZZ+jUqZPYXaM66MyZMxg9ejQOHjwIQRDQo0cPbNu2De7u7jXaD95Jq0Gai0/Xrl0BAB4eHkzQyCr8/vvv6NWrFxM0EpVSqURhYSFKSkpQUlICHx8fsbtEddT58+fRqVMnODk5wdnZGREREdi2bVuN94NJmhH27duHAQMGoGHDhpBIJOV+NZSYmIigoCA4OTmhQ4cOOHLkiHbf5cuXUb9+fQwYMACPPfYYPv744xrsPdU25sZjWT///DNGjBhRzT2m2szcePT29sZbb72FJk2aoGHDhoiJiUHTpk1r8AyoNjE3HsPDw7Fnzx5kZmbiwYMH2LNnD27fvl2DZ1CKSZoR8vLyEBERgcTExHL3r1mzBvHx8UhISMCJEycQERGB2NhYpKenAyj9K3H//v344osvcOjQIWzfvh3bt2+vyVOgWsTceNTIzs7GwYMH8cQTT9REt6mWMjceHzx4gE2bNuH69eu4ffs2Dh48iH379tXkKVAtYm48ap6L7NmzJwYPHoyOHTtCJpPV5CmUEsgkAIT169frbGvfvr0QFxenfa1SqYSGDRsKc+bMEQRBEA4ePCj06dNHu3/u3LnC3Llza6S/VLuZEo8a33//vfDcc8/VRDepjjAlHn/++Wdh/Pjx2v1z584VPv300xrpL9Vu5lwfNcaNGyds2rSpOrtZLt5Js5Di4mIcP34cMTEx2m1SqRQxMTE4dOgQAKBdu3ZIT0/HgwcPoFarsW/fPoSGhorVZarFDIlHDX7VSdXNkHgMCAjAwYMHUVhYCJVKhT179qB58+ZidZlqMUOvj5q7ahcvXsSRI0cQGxtb433lU+sWcu/ePahUKvj6+ups9/X1xYULFwAAdnZ2+Pjjj9GtWzcIgoA+ffrgySefFKO7VMsZEo8AkJWVhSNHjmDdunU13UWqQwyJx44dO+KJJ55AmzZtIJVK0atXLwwcOFCM7lItZ+j18amnnkJWVhbq1auH5cuXizLQj0laDevXrx/69esndjeIAAByuRxpaWlid4MIAPDRRx/ho48+ErsbRACg962DGPh1p4V4eXlBJpPp/cJLS0uDn5+fSL2iuorxSNaE8UjWxJbikUmahTg4OKBt27bYuXOndptarcbOnTs5GSPVOMYjWRPGI1kTW4pHft1phNzcXFy5ckX7OiUlBUlJSfDw8ECTJk0QHx+P0aNHIyoqCu3bt8eiRYuQl5eHMWPGiNhrqq0Yj2RNGI9kTWpNPNb4eFIbtnv3bgGA3r/Ro0dry3z++edCkyZNBAcHB6F9+/bC4cOHxesw1WqMR7ImjEeyJrUlHrl2JxEREZEV4jNpRERERFaISRoRERGRFWKSRkRERGSFmKQRERERWSEmaURERERWiEkaERERkRVikkZERERkhZikEREREVkhJmlEZJIXX3wRgwYNqrbj79mzBxKJBJmZmQCAFStWwN3dvdra0wgKCsKiRYuqvZ2KREdHY/LkydXaRnFxMUJCQnDw4EEAwPXr1yGRSJCUlFSt7ZqiY8eOWLdundjdIBIFkzSiWsLUX+41kRRYo4qSvqNHj+KVV16p+Q7969dff8Xs2bOrtY2lS5ciODgYjz/+uEWPK5FIsGHDBose8/3338fUqVOhVqstelwiW8AkjYioDG9vb7i4uIjWvoeHB1xdXavt+IIg4D//+Q/GjRtXbW1YUr9+/ZCTk4OtW7eK3RWiGsckjagWePHFF7F3714sXrwYEokEEokE169fBwDs3bsX7du3h6OjI/z9/TF16lQolcpK66lUKowbNw7BwcFwdnZG8+bNsXjxYqP79ddffyE6OhouLi5o0KABYmNj8eDBAwBAUVERJk6cCB8fHzg5OaFLly44evSoUcf/7bff8Nhjj8HJyQmPPPIIZs6cqT03AMjMzMSrr74KX19fODk5ITw8HJs2bcKePXswZswYZGVlac97xowZAHS/7nz22WcxYsQInTZLSkrg5eWF77//HgCgVqsxZ84c7XsVERGBX375pdJ+f/HFF2jWrBmcnJzg6+uLoUOHaveVvbOp+cr34X8vvviiwe/Bw44fP46rV6+if//+evsuXLiAxx9/XPte7d27F0BpYhcSEoLPPvtMp3xSUhIkEgmuXLmCoKAgAMDTTz8NiUSifV1VHwVBwIwZM9CkSRM4OjqiYcOGmDhxorauTCbDE088gdWrV1f6nhLVSqIu705EFpGZmSl06tRJePnll4W7d+8Kd+/eFZRKpfDPP/8ILi4uwvjx44Xz588L69evF7y8vISEhIRK6xUXFwvTp0////buPybq+o8D+JMfd8cBd4Bcvwg5NBAkwx8rBK7BSklmsVZ0OmsOS/+Q89dqltofgpu21Am0xmhuCfPXsnJOIULkl7Jz4S3DRO1cJyY0G2msuoDjunt+/3B+8uLwaMuvjF6PjT94v3m936/P+267196fe3+gzWbjlStXuH//foaHh/PQoUPKnEVFRXzxxRdHzembb76hRqNhcXExOzs72dXVxQ8//JA///wzSXLt2rWMi4tjfX09L1y4wKKiIsbExPDmzZskydbWVgJgf38/SbK6uppRUVHK+KdOnaJer2dNTQ0dDgcbGxuZmJjI0tJSkqTH42FmZiYff/xxNjY20uFwsLa2lvX19XS5XKyoqKBer1eu+/fffydJGo1GlpeXkyTr6uqo1WqVPpKsra2lVqvlb7/9RpLcunUrU1NT2dDQQIfDwerqamo0Gra1tfldF5vNxpCQEB48eJBXr17l2bNn+cEHHyj9ubm5XLduHUnS5XIp+V2/fp0tLS0MCwvjxx9/PKY18KesrIypqak+bd3d3QTA+Ph4fv7557x48SJXrFhBnU7HGzdukCS3bdvGtLQ0n7i1a9cyJyeHJNnX10cArK6u5vXr19nX1zemHD/77DPq9XrW19fzhx9+YEdHB3fv3u0zT1VVFY1G46jXJMREJUWaEBPEnR/ut7377rtMSUmh1+tV2iorKxkZGUmPxzNqnD+rVq1iYWGh8nugIm3JkiU0mUx++5xOJ1UqFQ8cOKC0DQ8PMy4ujjt27CAZuEibN28e33vvPZ9x9+3bx0ceeYQkefz4cQYHB9Nut/vN4e/j3XZnkeZ2u2kwGLh3716f61q8eDFJcmhoiOHh4Tx9+rTPGMuXL+eSJUv8znv48GHq9XqlyPu70V6PGzducOrUqbRYLEpboDXwZ926dXz22Wd92m4Xae+//77S5na7GR8fz+3bt5Mkf/zxR4aEhLCjo4PkrdfLYDCwpqZGiQHAI0eO+IwdKMddu3Zx2rRpHB4eHjXno0ePMjg4WHnPCvFfEXo/d/GEEPfWpUuXkJWVhaCgIKXNZDLB6XSit7cXCQkJo8ZWVlZiz549uHbtGgYHBzE8PIxZs2aNee7Ozk6YzWa/fQ6HA263GyaTSWlTqVTIyMjApUuXxjT+uXPnYLVasW3bNqXN4/FgaGgIAwMD6OzsRHx8PKZNmzbmnP8uNDQUixYtwoEDB7B06VL88ccfOHr0qHLr7fvvv8fAwADy8vJ84oaHhzF79my/Y+bl5cFoNGLq1KnIz89Hfn4+Xnrppbt+D87tdqOwsBBGo9HntnOgNfA35uDgIMLCwvzOk5WV5XPtTz75pPJ6xMXF4fnnn8eePXuQkZGB2tpauFyuUV/jseZoNptRUVGhrMfChQtRUFCA0NC/Pp60Wi28Xi9cLhe0Wu1d5xNiIpEiTQgxwieffIL169dj165dyMrKgk6nw86dO9HR0THmMe71h6nT6cSWLVvw8ssvj+gLCwv71+Z/7bXXkJubi76+Ppw4cQJarRb5+flKDgDwxRdf4NFHH/WJ02g0fsfT6XQ4e/Ys2tra0NjYiM2bN6O0tBQ2m23UR4wUFxejp6cHZ86c8SleAq2BPwaDAefPnw943f6sWLECS5cuRXl5Oaqrq7F48eKAhywC5Th58mTY7XY0NTXhxIkTsFgs2LlzJ06ePAmVSgUA+OWXXxARESEFmvjPkSJNiAlCrVbD4/H4tE2fPh2HDx8GSWU3zWq1QqfTIT4+ftQ4q9WK7OxsWCwWpc3hcPyjfNLT09Hc3IwtW7aM6HvsscegVqthtVphNBoB3NotstlsY34cyJw5c2C325GUlDTq/L29vbh8+bLf3TR/1+1PdnY2Jk+ejEOHDuHLL7+E2WxWioe0tDRoNBpcu3YNubm5Y8obuLVLNX/+fMyfPx8lJSWIjo5GS0uL30KmrKwMn376KU6fPo3Y2FifvkBr4M/s2bNRVVXl85647auvvkJOTg4A4M8//8TXX3+N1atXK/0LFy5EREQEqqqq0NDQgFOnTvnEq1SqEWs6lhy1Wi0KCgpQUFCAVatWITU1FefPn8ecOXMAAF1dXaPuTAoxkUmRJsQEkZiYiI6ODly9ehWRkZGYNGkSLBYLKioqsGbNGqxevRp2ux0lJSV46623EBwcPGpccnIy9u7di+PHj2PKlCnYt28fbDYbpkyZMuZ8Nm3ahCeeeAIWiwUrV66EWq1Ga2srzGYzDAYDiouL8fbbb2PSpElISEjAjh07MDAwMOZHQ2zevBkvvPACEhIS8MorryA4OBjnzp1DV1cXtm7ditzcXOTk5KCwsBBlZWVISkrCd999h6CgIOTn5yMxMRFOpxPNzc2YOXMmwsPDR90VevXVV/HRRx/h8uXLaG1tVdp1Oh3Wr1+PN998E16vF08//TR+/fVXWK1W6PV6FBUVjRirrq4OV65cQU5ODmJiYlBfXw+v14uUlJQRf9vU1IR33nkHlZWVMBgM+OmnnwDcKmqioqICroE/zzzzDJxOJy5cuIAZM2b49FVWViI5ORnTp09HeXk5+vv78cYbbyj9ISEhWLZsGTZt2oTk5GSf26PArfdSc3MzTCYTNBoNYmJiAuZYU1MDj8eDuXPnIjw8HPv374dWq1WKdwBob2/Hc8895/d6hJjQ7veX4oQQ/w673c7MzExqtVoCYHd3N0myra2NTz31FNVqNR9++GFu2LCBbrf7rnFDQ0NctmwZo6KiGB0dzeLiYm7cuJEzZ85U4gIdHLg9d3Z2NjUaDaOjo7lgwQLlIMDg4CDXrFlDg8FAjUZDk8nEM2fOKLGBDg6QZENDA7Ozs6nVaqnX65mRkeFzMvDmzZt8/fXXGRsby7CwMM6YMYN1dXVK/8qVKxkbG0sAyonXOw8O3Hbx4kUCoNFo9DmEQZJer5cVFRVMSUmhSqXiAw88wAULFvDkyZN+16S9vZ25ubmMiYmhVqtlenq6z6nZOw8OlJSUEMCIn6KiojGvgT+LFi3ixo0bld9vHxw4ePAgMzIyqFarmZaWxpaWlhGxDoeDAJQDHnc6duwYk5KSGBoa6nMa8245HjlyhHPnzqVer2dERAQzMzPZ1NSkxPb29lKlUrGnp+eu1yTERBREkvenPBRCCHE/fPvtt8jLy4PD4UBkZOQ/im1vb8e8efPQ09ODhx566B5l+JcNGzagv78fu3fvvudzCTHeyMNshRDiPyY9PR3bt29Hd3f3mGNcLhd6e3tRWloKs9n8fynQAODBBx+85/8mS4jxSnbShBBCBFRTU4Ply5dj1qxZOHbs2IjTrEKIf58UaUIIIYQQ45Dc7hRCCCGEGIekSBNCCCGEGIekSBNCCCGEGIekSBNCCCGEGIekSBNCCCGEGIekSBNCCCGEGIekSBNCCCGEGIekSBNCCCGEGIekSBNCCCGEGIf+B2YYuW17ct0CAAAAAElFTkSuQmCC",
+ "text/plain": [
+ "
"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ }
+ ],
+ "source": [
+ "fig, ax = plt.subplots(figsize=(7, 5))\n",
+ "ax.loglog(sweep[\"collective_size_bytes\"], sweep[\"model_latency_s\"] * 1e6, \"o-\",\n",
+ " label=f\"ISL model ({LINK_BW_GBPS:g} GB/s links)\")\n",
+ "ax.set_xlabel(\"total collective size (bytes)\")\n",
+ "ax.set_ylabel(\"all-to-all latency (µs)\")\n",
+ "ax.set_title(f\"{n}-GPU fully-connected NVLink all-to-all (ISL network model)\")\n",
+ "ax.grid(True, which=\"both\", alpha=0.3)\n",
+ "ax.legend() # ASTRA-sim 2.0 and EC2 nccl-tests series overlay here later\n",
+ "plt.show()"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "449a7991",
+ "source": "## Torus leg: 2×2×2 logical torus (3-cube) all-to-all\n\nGPU $r$ sits at binary coordinates $(g_0, g_1, g_2) \\in \\{0, 1\\}^3$ of a 3-dimensional `noc` space, one coordinate per torus dimension. Rank is recovered row-major with the **last** dimension fastest-varying, $r = 4g_0 + 2g_1 + g_2$ — matching the `torus_bench` benchmark's rank$\\leftrightarrow$coordinate convention.\n\nEvery dimension here has extent 2, so the $\\pm 1 \\pmod 2$ torus wraparound edge coincides with the only other value that coordinate can take: wraparound degenerates exactly to the hypercube edge, and `HypercubeMulticastModel` applies directly. Its precondition — a Manhattan, translation-invariant `dist_fn` — is satisfied by Hamming distance on the binary coordinates (ISL has no `!=`/`<>` operator, so per-axis inequality is written as the disjunction `(a < b or a > b)`).\n\nCost per $(s, d)$ pair at Hamming distance $k$ is $2^k - 1$ (the bounding-box multicast-tree cost `HypercubeMulticastModel` computes), summing to **152** hops over all 64 `data[s, d]` chunks. A dimension-ordered *minimal unicast* routing — one packet per hop instead of a shared multicast tree, which is what the empirical torus benchmark actually implements — would instead traverse $\\sum_{s,d} \\mathrm{hamming}(s, d) = $ **96** hops. The two series bracket real routing behavior; the gap between them is itself a correlation observable, quantifying how far a multicast-aware model diverges from point-to-point unicast.",
+ "metadata": {}
+ },
+ {
+ "cell_type": "code",
+ "id": "d166f6c4",
+ "source": "def binary_bounds(prefix: str, k: int) -> str:\n \"\"\"Bit constraints over dims ``{prefix}0..{prefix}{k-1}``.\"\"\"\n return \" and \".join(f\"0 <= {prefix}{i} <= 1\" for i in range(k))\n\n\ndef rank_expr(prefix: str, k: int) -> str:\n \"\"\"Affine recovery of the rank from binary torus coordinates.\n\n Row-major with the LAST coordinate fastest-varying:\n ``rank = sum_i 2**(k-1-i) * g_i``, matching the ``torus_bench``\n benchmark's rank<->coordinate convention (see the markdown above).\n \"\"\"\n return \" + \".join(f\"{2 ** (k - 1 - i)}*{prefix}{i}\" for i in range(k))\n\n\ndef torus_all_to_all_maps(\n dims: tuple[int, ...] = (2, 2, 2),\n) -> tuple[isl.Map, isl.Map, isl.Map]:\n \"\"\"Build (occupancy, fill, dist_fn) for a logical-torus all-to-all.\n\n Mirrors ``all_to_all_maps`` above but keys GPUs by binary torus\n coordinates instead of a one-hot id. ``data[s, d]`` is the chunk sent\n by GPU ``s`` to GPU ``d``; each GPU holds the chunks it sends (occ)\n and requests the chunks addressed to it (fill).\n\n Parameters\n ----------\n dims : tuple of int, default (2, 2, 2)\n Per-dimension extents of the logical torus. Every entry must\n equal 2 (see Raises); ``K = len(dims)`` sets the dimensionality\n of the ``noc`` space and the number of ``SpatialTag`` axes the\n caller must build to match.\n\n Returns\n -------\n occ : isl.Map\n ``noc[g0..g_{K-1}] -> data[s, d]``: the GPU at ``(g0..g_{K-1})``\n holds the chunk it is the source of, for every destination ``d``.\n fill : isl.Map\n ``noc[g0..g_{K-1}] -> data[s, d]``: the GPU at ``(g0..g_{K-1})``\n requests the chunk addressed to it, for every source ``s``.\n dist_fn : isl.Map\n Piecewise Hamming distance between two ``noc`` points, encoded as\n a disjunction over the ``2**K`` equal/differ patterns per axis.\n\n Raises\n ------\n ValueError\n If any entry of ``dims`` is not 2. General (non-power-of-two)\n torus extents need a genuine wraparound-aware Manhattan distance\n function, not the hypercube-degeneracy shortcut used here — see\n the markdown cell above for why extent 2 is special.\n \"\"\"\n if any(d != 2 for d in dims):\n raise ValueError(\n f\"torus_all_to_all_maps only supports extent-2 dims (hypercube \"\n f\"wraparound degeneracy), got dims={dims!r}; general torus \"\n f\"extents need a wraparound-aware Manhattan dist_fn\"\n )\n k = len(dims)\n n = 2 ** k\n\n gs = \", \".join(f\"gs{i}\" for i in range(k))\n gd = \", \".join(f\"gd{i}\" for i in range(k))\n occ = isl.Map.read_from_str(\n CTX,\n f\"{{ noc[{gs}] -> data[s, d] : {binary_bounds('gs', k)} \"\n f\"and s = {rank_expr('gs', k)} and 0 <= d < {n} }}\",\n )\n fill = isl.Map.read_from_str(\n CTX,\n f\"{{ noc[{gd}] -> data[s, d] : {binary_bounds('gd', k)} \"\n f\"and d = {rank_expr('gd', k)} and 0 <= s < {n} }}\",\n )\n\n # Piecewise Hamming distance: one disjunct per subset of axes that\n # differ (2**K sign patterns), cost = popcount of the subset. ISL has\n # no `!=`, so per-axis inequality is the disjunction (a < b or a > b).\n pieces = []\n for mask in range(2 ** k):\n dist = bin(mask).count(\"1\")\n conds = [\n f\"(gd{i} < gs{i} or gd{i} > gs{i})\" if mask & (1 << i) else f\"gd{i} = gs{i}\"\n for i in range(k)\n ]\n pieces.append(\n f\"[noc[{gd}] -> noc[{gs}]] -> hops[{dist}] : \" + \" and \".join(conds)\n )\n dist_fn = isl.Map.read_from_str(CTX, \"{ \" + \"; \".join(pieces) + \" }\")\n\n return occ, fill, dist_fn",
+ "metadata": {},
+ "execution_count": null,
+ "outputs": []
+ },
+ {
+ "cell_type": "code",
+ "id": "256ea183",
+ "source": "dims_torus = (2, 2, 2)\nk = len(dims_torus)\nn_t = 2 ** k # = 8 GPUs, same physical count as the FC leg above\n\nocc_t, fill_t, dist_fn_t = torus_all_to_all_maps(dims_torus)\ntags_t = [SpatialTag(i, 0) for i in range(k)]\nmodel_t = HypercubeMulticastModel(dist_fn_t)\n\n\ndef torus_pair_hops(src: int, dst: int) -> int:\n \"\"\"Hops for a single (src, dst) chunk under the torus (Hamming) dist_fn.\"\"\"\n chunk = isl.Set.read_from_str(CTX, f\"{{ data[{src}, {dst}] }}\")\n info = model_t.apply(\n 0,\n Fill(tags_t, fill_t.intersect_range(chunk)),\n Occupancy(tags_t, occ_t.intersect_range(chunk)),\n )\n return eval_total(info.hops)\n\n\ndef hamming(a: int, b: int) -> int:\n \"\"\"Hamming distance between the binary representations of two ints.\"\"\"\n return bin(a ^ b).count(\"1\")\n\n\n# --- Run the tool over the full all-to-all and validate against the\n# verified reference numbers (see torus_hops_check.py pre-verification). ---\ninfo_t = model_t.apply(0, Fill(tags_t, fill_t), Occupancy(tags_t, occ_t))\ntorus_hops = eval_total(info_t.hops)\nassert torus_hops == 152, f\"model returned {torus_hops} total hops, expected 152\"\n\nprobes = {(0, 1): 1, (0, 3): 3, (0, 7): 7, (5, 5): 0}\nfor (src, dst), expected in probes.items():\n got = torus_pair_hops(src, dst)\n assert got == expected, f\"pair ({src}->{dst}) cost {got} != {expected}\"\n\n# Dimension-ordered minimal unicast routing (what the empirical torus\n# benchmark implements): one hop per differing bit, summed over every\n# ordered (s, d) pair. Computed programmatically -- not hardcoded -- so\n# it stays correct if dims_torus ever changes.\nminroute_hops = sum(hamming(s, d) for s in range(n_t) for d in range(n_t))\nassert minroute_hops == 96, f\"minroute_hops {minroute_hops} != 96\"\n\nprint(f\"ISL network model: {n_t}-GPU 2x2x2 logical torus all-to-all (binary coords)\")\nprint(f\" bounding-box multicast hops: {torus_hops} (Sigma 2^hamming(s,d) - 1)\")\nprint(f\" min-route unicast hops : {minroute_hops} (Sigma hamming(s,d))\")\nprint(f\" routing gap : {torus_hops - minroute_hops} hops \"\n f\"({torus_hops / minroute_hops:.2f}x)\")",
+ "metadata": {},
+ "execution_count": null,
+ "outputs": []
+ },
+ {
+ "cell_type": "code",
+ "id": "8f3f4946",
+ "source": "# 8 nodes * 3 distinct neighbors/node. Extent-2 wraparound means the +1\n# and -1 neighbor along a dim coincide (mod 2), so each node has exactly\n# one neighbor per torus dimension, not two -- hence 3, not 6.\nDIRECTED_LINKS_TORUS = n_t * k\n\nrows_t = []\nmib = MIN_MIB\nwhile mib <= MAX_MIB:\n size = mib * (1 << 20)\n total = size * n_t if PER_RANK else size\n per_rank = total / n_t\n chunk_bytes = total / (n_t * n_t)\n # By symmetry, all-to-all traffic loads all 24 directed torus links\n # uniformly: total hop-bytes (chunk_bytes * hops, summed over the 64\n # chunks) divides evenly across the links. Two series bracket\n # routing: `torus_hops` (bounding-box multicast) vs. `minroute_hops`\n # (dimension-ordered unicast) -- see the markdown cell above.\n per_link_bytes_model = (torus_hops / DIRECTED_LINKS_TORUS) * chunk_bytes\n per_link_bytes_minroute = (minroute_hops / DIRECTED_LINKS_TORUS) * chunk_bytes\n model_latency_s = ALPHA_S + per_link_bytes_model / bw\n minroute_latency_s = ALPHA_S + per_link_bytes_minroute / bw\n # Closed-form cross-check: 152/(64*24) = 19/192, 96/(64*24) = 1/16.\n assert math.isclose(model_latency_s, ALPHA_S + total * 19 / (192 * bw), rel_tol=1e-12)\n assert math.isclose(minroute_latency_s, ALPHA_S + total / (16 * bw), rel_tol=1e-12)\n rows_t.append((n_t, int(total), int(per_rank), per_link_bytes_model,\n model_latency_s, minroute_latency_s))\n mib *= 2\n\nsweep_torus = pd.DataFrame(rows_t, columns=[\n \"nodes\", \"collective_size_bytes\", \"per_rank_bytes\", \"per_link_bytes_model\",\n \"model_latency_s\", \"minroute_latency_s\",\n])\nsweep_torus",
+ "metadata": {},
+ "execution_count": null,
+ "outputs": []
+ },
+ {
+ "cell_type": "markdown",
+ "id": "5e209b22",
+ "source": "## Empirical overlay: EC2 nccl-tests + torus_bench\n\nCSV files under `correlation/data/` (e.g. `correlation/data//csv/*.csv`) are fetched by `correlation/orchestrate.py` (sibling infrastructure, not part of this notebook) and share a unified schema: `source, topology, dims, collective, size_bytes, count, dtype, time_us, algbw_GBps, busbw_GBps, wrong`.\n\nSize normalization to the total collective size $S$ used by `sweep` / `sweep_torus` above differs by source: for `source == \"torus_bench\"` rows, `size_bytes` is **already** the total $S$. For `source == \"nccl-tests\"` rows with `collective == \"alltoall\"`, `size_bytes` is the *per-rank* message size, so $S = N \\times \\texttt{size\\_bytes}$ — verify this against the nccl-tests version actually in use, since the size semantics nccl-tests reports differ per collective. Only `alltoall` is normalized and correlated against the model here; other collectives are inventoried below for later work.",
+ "metadata": {}
+ },
+ {
+ "cell_type": "code",
+ "id": "4715980a",
+ "source": "from pathlib import Path\n\n# Design: prefer the path relative to the notebook's own directory (the\n# normal case when the notebook is run in place), but fall back to the\n# repo-root-relative path so `jupyter execute` from a different cwd (e.g.\n# CI running from the repo root) still finds real data when it exists,\n# instead of silently taking the empty-data path below.\nDATA_DIR = (\n Path(\"correlation/data\") if Path(\"correlation\").is_dir()\n else Path(\"notebooks/astrasim2_correlation/correlation/data\")\n)\n\nUNIFIED_CSV_COLUMNS = [\n \"source\", \"topology\", \"dims\", \"collective\", \"size_bytes\", \"count\",\n \"dtype\", \"time_us\", \"algbw_GBps\", \"busbw_GBps\", \"wrong\",\n]\n_NUMERIC_COLUMNS = [\"size_bytes\", \"count\", \"time_us\", \"algbw_GBps\", \"busbw_GBps\"]\n\nfiles = sorted(DATA_DIR.glob(\"**/*.csv\"))\nif files:\n emp = pd.concat(\n [pd.read_csv(f, dtype={\"wrong\": str}) for f in files],\n ignore_index=True,\n )\n for col in _NUMERIC_COLUMNS:\n emp[col] = pd.to_numeric(emp[col], errors=\"coerce\")\nelse:\n # Design: keep the unified schema even when no files were found, so\n # downstream cells can reference emp[...] columns inside their\n # `if HAVE_EMPIRICAL:` guards without ever hitting a KeyError.\n emp = pd.DataFrame(columns=UNIFIED_CSV_COLUMNS)\n\nHAVE_EMPIRICAL = len(files) > 0 and not emp.empty\n\nif not HAVE_EMPIRICAL:\n print(\"no empirical data yet -- run correlation/orchestrate.py; \"\n \"overlay cells below will no-op\")\nelse:\n inventory = emp.groupby([\"topology\", \"collective\"]).agg(\n rows=(\"size_bytes\", \"count\"),\n min_size=(\"size_bytes\", \"min\"),\n max_size=(\"size_bytes\", \"max\"),\n )\n display(inventory)",
+ "metadata": {},
+ "execution_count": null,
+ "outputs": []
+ },
+ {
+ "cell_type": "code",
+ "id": "a2e080ce",
+ "source": "fig, ax = plt.subplots(figsize=(7, 5))\n\nax.loglog(sweep[\"collective_size_bytes\"], sweep[\"model_latency_s\"] * 1e6, \"o-\",\n label=f\"FC model ({LINK_BW_GBPS:g} GB/s links)\")\nax.loglog(sweep_torus[\"collective_size_bytes\"], sweep_torus[\"model_latency_s\"] * 1e6, \"s-\",\n label=\"torus model, 152 hops (bounding-box multicast)\")\nax.loglog(sweep_torus[\"collective_size_bytes\"], sweep_torus[\"minroute_latency_s\"] * 1e6, \"s--\",\n label=\"torus min-route, 96 hops (dimension-ordered unicast)\")\n\nif HAVE_EMPIRICAL:\n alltoall = emp[emp[\"collective\"] == \"alltoall\"]\n fc_rows = alltoall[alltoall[\"topology\"] == \"fc\"]\n torus_rows = alltoall[alltoall[\"topology\"] == \"torus\"]\n if not fc_rows.empty:\n # nccl-tests alltoall size_bytes is the PER-RANK message size (see\n # markdown above); scale to total S to match `sweep`'s convention.\n ax.scatter(fc_rows[\"size_bytes\"] * NODES, fc_rows[\"time_us\"],\n marker=\"x\", color=\"tab:red\", label=\"EC2 nccl-tests alltoall (fc)\")\n if not torus_rows.empty:\n # torus_bench size_bytes is ALREADY total collective size S.\n ax.scatter(torus_rows[\"size_bytes\"], torus_rows[\"time_us\"],\n marker=\"+\", color=\"tab:green\", label=\"EC2 torus_bench alltoall (torus)\")\n\nax.set_xlabel(\"total collective size (bytes)\")\nax.set_ylabel(\"all-to-all latency (µs)\")\nax.set_title(\"Fully-connected vs. 2x2x2 torus all-to-all: model vs. empirical\")\nax.grid(True, which=\"both\", alpha=0.3)\nax.legend()\nplt.show()",
+ "metadata": {},
+ "execution_count": null,
+ "outputs": []
+ },
+ {
+ "cell_type": "code",
+ "id": "a8ec5d63",
+ "source": "import numpy as np\n\n\ndef linear_fit(S: np.ndarray, t: np.ndarray) -> tuple[float, float, float]:\n \"\"\"Ordinary-least-squares fit of ``t = intercept + slope * S``.\n\n Parameters\n ----------\n S : np.ndarray\n Total collective size in bytes for each empirical sample.\n t : np.ndarray\n Measured latency in seconds for each empirical sample. Must be\n the same length as `S`.\n\n Returns\n -------\n intercept : float\n Fitted per-operation latency overhead (the model's alpha), seconds.\n slope : float\n Fitted dt/dS, seconds/byte.\n r_squared : float\n Coefficient of determination of the linear fit; NaN if `t` has\n zero total variance (R^2 is undefined in that degenerate case).\n \"\"\"\n slope, intercept = np.polyfit(S, t, 1)\n t_pred = slope * S + intercept\n ss_res = float(np.sum((t - t_pred) ** 2))\n ss_tot = float(np.sum((t - np.mean(t)) ** 2))\n r_squared = 1.0 - ss_res / ss_tot if ss_tot > 0 else float(\"nan\")\n return float(intercept), float(slope), r_squared\n\n\ndef bw_from_slope(slope: float, slope_factor: float) -> float:\n \"\"\"Translate a fitted t-vs-S slope into an effective bandwidth, bytes/s.\n\n Given the model form ``t = alpha + S * slope_factor / BW``, the\n fitted slope is ``dt/dS = slope_factor / BW``, so\n ``BW = slope_factor / slope``.\n \"\"\"\n return slope_factor / slope\n\n\nif HAVE_EMPIRICAL and (emp[\"collective\"] == \"alltoall\").any():\n alltoall = emp[emp[\"collective\"] == \"alltoall\"].dropna(\n subset=[\"size_bytes\", \"time_us\"]\n )\n\n fc_rows = alltoall[alltoall[\"topology\"] == \"fc\"]\n if len(fc_rows) >= 2:\n # fc slope_factor = (N-1)/N^2 = 7/64 (N=8), from the closed form\n # in the \"Size and latency conventions\" cell above.\n S_fc = fc_rows[\"size_bytes\"].to_numpy(dtype=float) * NODES\n t_fc = fc_rows[\"time_us\"].to_numpy(dtype=float) * 1e-6\n alpha_fc, slope_fc, r2_fc = linear_fit(S_fc, t_fc)\n bw_fc = bw_from_slope(slope_fc, 7 / 64)\n print(f\"FC calibration : alpha = {alpha_fc * 1e6:.3f} us, \"\n f\"BW_eff = {bw_fc / 1e9:.2f} GB/s, R^2 = {r2_fc:.4f}\")\n else:\n print(\"FC calibration : fewer than 2 alltoall rows, skipping fit\")\n\n torus_rows = alltoall[alltoall[\"topology\"] == \"torus\"]\n if len(torus_rows) >= 2:\n # torus slope_factor = hops/(64*24) = hops/1536; report BW under\n # both routing hypotheses (h=152 bounding-box, h=96 min-route)\n # from a single fit, since alpha/R^2 don't depend on slope_factor\n # -- only the slope-to-BW translation does.\n S_t = torus_rows[\"size_bytes\"].to_numpy(dtype=float)\n t_t = torus_rows[\"time_us\"].to_numpy(dtype=float) * 1e-6\n alpha_t, slope_t, r2_t = linear_fit(S_t, t_t)\n bw_t_152 = bw_from_slope(slope_t, 152 / 1536)\n bw_t_96 = bw_from_slope(slope_t, 96 / 1536)\n print(f\"torus calibration : alpha = {alpha_t * 1e6:.3f} us, R^2 = {r2_t:.4f}\")\n print(f\" BW_eff @ h=152 (bounding-box model): {bw_t_152 / 1e9:.2f} GB/s\")\n print(f\" BW_eff @ h=96 (dimension-ordered) : {bw_t_96 / 1e9:.2f} GB/s\")\n else:\n print(\"torus calibration : fewer than 2 alltoall rows, skipping fit\")\nelse:\n print(\"calibration: no empirical alltoall rows yet\")",
+ "metadata": {},
+ "execution_count": null,
+ "outputs": []
+ }
+ ],
+ "metadata": {
+ "kernelspec": {
+ "display_name": "accelforge",
+ "language": "python",
+ "name": "python3"
+ },
+ "language_info": {
+ "codemirror_mode": {
+ "name": "ipython",
+ "version": 3
+ },
+ "file_extension": ".py",
+ "mimetype": "text/x-python",
+ "name": "python",
+ "nbconvert_exporter": "python",
+ "pygments_lexer": "ipython3",
+ "version": "3.12.3"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/notebooks/astrasim2_correlation/correlation/.gitignore b/notebooks/astrasim2_correlation/correlation/.gitignore
new file mode 100644
index 00000000..6bddd09b
--- /dev/null
+++ b/notebooks/astrasim2_correlation/correlation/.gitignore
@@ -0,0 +1,4 @@
+keys/
+.state/
+*.pem
+logs/
diff --git a/notebooks/astrasim2_correlation/correlation/README.md b/notebooks/astrasim2_correlation/correlation/README.md
new file mode 100644
index 00000000..4eb33486
--- /dev/null
+++ b/notebooks/astrasim2_correlation/correlation/README.md
@@ -0,0 +1,188 @@
+# correlation/ -- AWS provisioning for NCCL profiling
+
+This directory is the **empirical leg** of the ISL-model correlation
+study: it profiles real NCCL collective-communication performance on one
+AWS `p5.48xlarge` instance (8x H100, NVSwitch), both in the GPUs' native
+fully-connected (FC) topology and under a logical torus overlay, so those
+measurements can be compared against the model's and ASTRA-sim's
+predictions.
+
+## Prerequisites
+
+- **AWS credentials** available to boto3 via the standard mechanisms (env
+ vars `AWS_ACCESS_KEY_ID`/`AWS_SECRET_ACCESS_KEY`/`AWS_SESSION_TOKEN`, a
+ named profile via `AWS_PROFILE`, or an attached IAM role). This
+ directory never manages or stores credentials itself.
+- **boto3** installed in whatever Python environment you run these
+ scripts from: `pip install boto3`. It is deliberately *not* a dependency
+ of the `accelforge` package -- only this notebook's provisioning corner
+ needs it.
+- **Service Quotas.** `p5.48xlarge` requires 192 vCPUs of quota. Before
+ your first run, check the Service Quotas console for your target region
+ and confirm both:
+ - the on-demand P-instance quota ("Running On-Demand P instances" or
+ equivalent, depending on current AWS naming), and
+ - the corresponding Spot P-instance vCPU quota (if you intend to use
+ `--purchasing spot` or the default `spot-then-ondemand`),
+
+ are each **>= 192 vCPUs**. Quota codes and exact names change over time
+ and vary by account/region presentation -- look them up in the console
+ rather than trusting a hardcoded code here. A quota that's too low
+ surfaces as an `InsufficientInstanceCapacity`-adjacent or limit-exceeded
+ error at launch time.
+
+## Cost warning
+
+`p5.48xlarge` on-demand pricing is roughly **$30-55/hr**, depending on
+region and current AWS pricing -- **verify current pricing** at
+ before running anything.
+A full profiling sweep (FC + torus, all collectives, the full message-size
+range) is expected to take well under an hour, but see the dead-man-timer
+note below: the default 120-minute timer is comfortably above that
+estimate, but a large custom `--min-mib`/`--max-mib` range or collective
+list can push a full sweep close to or past it. Spot pricing is
+substantially cheaper when capacity is available, which is why
+`spot-then-ondemand` is the default purchasing mode -- see
+`provision.launch_instance`'s docstring for the exact fallback conditions.
+Both `orchestrate.py` and `provision.py` print this same warning and
+require an interactive `yes` confirmation before touching AWS (skippable
+with `--yes`; not required for `--dry-run`).
+
+## Quickstart
+
+**`orchestrate.py` is the end-to-end path.** It provisions its own
+instance, pushes the profiling scripts, runs the full FC and/or torus
+sweep, fetches the results locally, and tears the instance down again --
+all in one command:
+
+```bash
+# 1. Sanity-check permissions and request shape without launching anything
+# (still creates a real SSH key pair + security group -- see the
+# --dry-run note below):
+python orchestrate.py --topology both --dry-run
+
+# 2. The real run: provision, profile both legs, fetch results, tear down.
+python orchestrate.py --topology both
+```
+
+**Do not run `provision.py` before `orchestrate.py`.** `orchestrate.py`
+provisions its *own* instance internally; running `provision.py` first
+would launch a *second*, independent `p5.48xlarge` instance that nothing
+in the `orchestrate.py` run above knows about or tears down, silently
+doubling your bill. `provision.py`/`teardown.py` are for advanced, manual
+control only -- see the section below.
+
+Every `Config` field (region, purchasing mode, topology, message-size
+range, collectives, ...) is a CLI flag; run `python orchestrate.py --help`
+for the full list, or pass `--config some.yaml` to load a batch of
+overrides from YAML (CLI flags still win over anything in the YAML file).
+Useful flags:
+
+- `--keep-alive`: skip `orchestrate.py`'s own teardown at the end, leaving
+ the instance running (its SSH command is printed) for manual
+ inspection. The on-instance dead-man timer still applies regardless.
+- `--dead-man-minutes `: extend the dead-man timer past its 120-minute
+ default -- see the note below.
+- `--ssh-cidr /32`: skip auto-detecting your IP for the SSH
+ security-group rule; required if IP auto-detection fails (see
+ `provision.caller_ip`'s docstring).
+
+### Fetched data layout
+
+Each leg's results land under `data///`:
+
+```
+data///
+├── csv/ # parsed, unified-schema CSVs consumed by the correlation notebook
+├── raw/ # raw nccl-tests/torus_bench stdout logs, one per collective
+└── metadata.txt # machine/software provenance (nvidia-smi topo, driver/NCCL versions, git rev)
+```
+
+`` is `fc` or `torus`. `data/` is the one subdirectory of this
+project *not* gitignored -- fetched results are committed intentionally.
+
+## Advanced / manual control: `provision.py` + `teardown.py`
+
+`provision.py` and `teardown.py` are the individual "up" and "down" halves
+`orchestrate.py` composes internally. Use them directly only if you need
+manual control between provisioning and profiling (e.g. debugging the
+instance by hand, or running a custom workload instead of
+`run_profile.sh`) -- most users should use `orchestrate.py` above instead.
+
+```bash
+# Provision one instance and leave it running (see the warning below):
+python provision.py
+
+# ... do whatever manual work you need on the instance ...
+
+# Tear down that one run:
+python teardown.py --run-id
+
+# Tear down everything this study has tagged, across every region this
+# study's local state knows about:
+python teardown.py --all
+
+# Also delete the SSH key pair (AWS-side) and local PEM:
+python teardown.py --run-id --delete-key
+
+# Audit: confirm nothing is left running, without tearing anything down.
+# Exits 0 ("no running instances") when clean, 1 with a table otherwise --
+# safe to use as a post-teardown check or a periodic cron/CI safety net.
+python teardown.py --verify
+```
+
+**Warning: `provision.py` leaves the instance running with NO dead-man
+timer until `setup_node.sh` is run on it.** The dead-man timer is armed
+*by* `setup_node.sh` (a step `orchestrate.py` always runs for you, but
+`provision.py` alone does not reach) -- so an instance provisioned via
+`provision.py` and never followed up with `setup_node.sh` (or
+`teardown.py`) will run, and bill, indefinitely with no automatic
+backstop. If you provision manually, either run `setup_node.sh` on the
+instance promptly or tear it down yourself as soon as you're done.
+
+`--dry-run` on **both** `provision.py` and `orchestrate.py` performs
+`DryRun=True` authorization checks only and launches no instance -- but a
+**real** SSH key pair and security group ARE still created in AWS either
+way (there is no dry-run equivalent for those two calls). Clean them up
+with:
+
+```bash
+python teardown.py --run-id --delete-key
+```
+
+`teardown.py` discovers resources two ways and reconciles them: local
+`.state/*.json` files written by `provision.py`/`orchestrate.py`, and a
+live `describe_instances` search by `Project`/`RunId` tags. The tag search
+is authoritative, so teardown still works even if a state file was lost or
+a run was started from a different machine. `--region` defaults to
+*resolving per run* rather than to a fixed region: an explicit `--region`
+flag always wins, otherwise each run's own state file (if any) supplies
+its region, otherwise `us-east-1` (`Config`'s default) is used -- so
+`--all`/`--verify` correctly span every region this study's local state
+knows about in one invocation, not just one hardcoded region.
+
+## Safety guardrails
+
+- **Dead-man timer -- armed by `setup_node.sh`, not at instance launch.**
+ Once `setup_node.sh` has run on the instance (always true for an
+ `orchestrate.py` run; not automatic for a manual `provision.py` one --
+ see the warning above), an on-instance timer force-shuts-down the box
+ after `--dead-man-minutes` (default 120) regardless of whether
+ `teardown.py` was ever run -- a backstop against a forgotten or failed
+ teardown. Re-running `setup_node.sh` (e.g. via a second `orchestrate.py`
+ leg) pushes the deadline back rather than erroring or stacking. Note
+ that the 120-minute default can be tight for a long custom sweep --
+ raise it with `--dead-man-minutes` if you expect to run past it.
+- **`InstanceInitiatedShutdownBehavior=terminate`.** The instance is
+ launched so that an in-instance `shutdown` (including the dead-man
+ timer firing) *terminates* it rather than merely stopping it, so it
+ cannot be left billing in a stopped state.
+- **Teardown-in-`finally`.** `orchestrate.py` calls teardown from a
+ `finally` block around the provisioning-through-profiling sequence, so a
+ mid-sweep crash or a failed leg still tears the instance down (unless
+ `--keep-alive` was passed).
+- **PEM files are gitignored.** `keys/`, `.state/`, `*.pem`, and `logs/`
+ are all excluded (see `.gitignore`) -- private key material and
+ per-run state never get committed. `data/` (fetched profiling CSVs) is
+ the one subdirectory *not* ignored; those results are committed
+ intentionally.
diff --git a/notebooks/astrasim2_correlation/correlation/config.py b/notebooks/astrasim2_correlation/correlation/config.py
new file mode 100644
index 00000000..9fb675f0
--- /dev/null
+++ b/notebooks/astrasim2_correlation/correlation/config.py
@@ -0,0 +1,666 @@
+"""Shared configuration for the NCCL profiling correlation-study AWS scripts.
+
+This module defines :class:`Config`, the single source of truth for every
+tunable knob used by ``provision.py``, ``teardown.py``, and (per the plan)
+the sibling ``orchestrate.py`` that a later work package will add. Keeping
+the configuration in one frozen dataclass -- rather than threading loose
+kwargs through each script -- means every script agrees on defaults,
+validation, and CLI flag names without duplicating logic.
+
+Construction paths
+-------------------
+``Config`` instances can be built three ways, all of which funnel through
+the same validation in :meth:`Config.__post_init__`:
+
+1. Directly, e.g. ``Config(region="us-west-2")`` -- handy for tests and for
+ any future caller that wants a config without touching argparse at all.
+2. Via :meth:`Config.from_args`, which owns its own
+ :class:`argparse.ArgumentParser` end to end.
+3. Via :meth:`Config.add_args` + :meth:`Config.from_parsed`, which lets a
+ *caller* (``provision.py``, ``teardown.py``) build one shared parser,
+ add its own script-specific flags (e.g. ``--dry-run``), and only then
+ hand the resulting namespace back to ``Config`` to extract just the
+ fields that belong to it. This is the path ``provision.py`` and
+ ``teardown.py`` actually use.
+
+Design: frozen dataclass
+-------------------------
+``Config`` is declared ``frozen=True`` so that once built it can be passed
+into ``provision.py`` functions (``launch_instance``, ``write_state``, ...)
+without any risk of one function's edits leaking into another's view of the
+same run. Frozen dataclasses cannot assign to ``self.`` in the usual
+way, so any post-construction normalization (parsing "2x2x2" into a tuple,
+generating a run id, coercing str paths to :class:`pathlib.Path`) goes
+through ``object.__setattr__`` inside ``__post_init__`` -- this is the
+standard, documented escape hatch for "derive a field after validation" on
+a frozen dataclass.
+"""
+
+from __future__ import annotations
+
+import argparse
+import dataclasses
+import datetime
+import math
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+
+# Design: anchor key_dir/state_dir to this file's directory (not the
+# process cwd) so that `python provision.py` behaves identically no matter
+# where the user's shell happens to be sitting when they invoke it.
+_THIS_DIR = Path(__file__).resolve().parent
+
+_ALLOWED_PURCHASING = frozenset({"spot", "ondemand", "spot-then-ondemand"})
+_ALLOWED_TOPOLOGY = frozenset({"fc", "torus", "both"})
+
+# Default sweep of NCCL collectives profiled by the (sibling-work-package)
+# profiling scripts. Kept here, not in profiling code, so a single
+# `--collectives` override on the CLI is the one place a user needs to
+# change to alter the sweep.
+_DEFAULT_COLLECTIVES: Tuple[str, ...] = (
+ "all_reduce",
+ "all_gather",
+ "reduce_scatter",
+ "alltoall",
+ "broadcast",
+ "sendrecv",
+)
+
+# The study profiles one 8x H100 node; a torus overlay must therefore have
+# axis dimensions whose product is exactly 8 (one "logical GPU slot" per
+# axis position), regardless of how many axes are used.
+_TORUS_GPU_COUNT = 8
+
+_RUN_ID_TIME_FORMAT = "%Y%m%d-%H%M%S"
+
+
+def _default_run_id() -> str:
+ """Generate a fresh, sortable run identifier.
+
+ Returns
+ -------
+ str
+ ``"correl-" + UTC timestamp`` formatted as
+ ``%Y%m%d-%H%M%S`` (e.g. ``"correl-20260716-161503"``). Lexicographic
+ sort order matches chronological order, which is convenient when
+ listing ``.state/*.json`` files or AMI/SG names in a shell.
+
+ Notes
+ -----
+ Uses UTC (not local time) so run ids are unambiguous and comparable
+ regardless of which machine or timezone invokes the script.
+ """
+ return "correl-" + datetime.datetime.now(datetime.timezone.utc).strftime(
+ _RUN_ID_TIME_FORMAT
+ )
+
+
+def _is_power_of_two(n: int) -> bool:
+ """Return whether a positive integer is an exact power of two.
+
+ Parameters
+ ----------
+ n : int
+ Value to test.
+
+ Returns
+ -------
+ bool
+ ``True`` if ``n > 0`` and ``n & (n - 1) == 0``, ``False`` otherwise
+ (including for ``n <= 0``).
+
+ Examples
+ --------
+ >>> _is_power_of_two(1024)
+ True
+ >>> _is_power_of_two(0)
+ False
+ >>> _is_power_of_two(3)
+ False
+ """
+ return n > 0 and (n & (n - 1)) == 0
+
+
+def _parse_torus_dims(value: Any) -> Tuple[int, ...]:
+ """Normalize a torus-dimensions spec into a tuple of ints.
+
+ Accepts the three shapes this value can arrive in depending on
+ construction path: an already-correct ``tuple[int, ...]`` (direct
+ ``Config(...)`` construction), a ``list``/``tuple`` of ints or numeric
+ strings (from a parsed YAML ``--config`` file), or a delimited string
+ like ``"2x2x2"`` (from the CLI, where argparse hands raw strings to
+ ``type=`` callables).
+
+ Parameters
+ ----------
+ value : Any
+ Torus dimensions in any of the accepted shapes described above.
+
+ Returns
+ -------
+ tuple[int, ...]
+ The parsed per-axis dimensions, in the order given.
+
+ Raises
+ ------
+ ValueError
+ If ``value`` is an empty string, or any component cannot be parsed
+ as an integer.
+
+ Examples
+ --------
+ >>> _parse_torus_dims("2x2x2")
+ (2, 2, 2)
+ >>> _parse_torus_dims([2, 4])
+ (2, 4)
+ >>> _parse_torus_dims((8,))
+ (8,)
+ """
+ if isinstance(value, (list, tuple)):
+ try:
+ return tuple(int(v) for v in value)
+ except (TypeError, ValueError) as exc:
+ raise ValueError(
+ f"torus_dims entries must all be integers, got {value!r}"
+ ) from exc
+
+ text = str(value).strip().lower()
+ if not text:
+ raise ValueError("torus_dims must not be empty")
+ try:
+ return tuple(int(part) for part in text.split("x"))
+ except ValueError as exc:
+ raise ValueError(
+ f"torus_dims={value!r} is not a valid dims string; "
+ "expected a form like '2x2x2'"
+ ) from exc
+
+
+def _parse_collectives(value: Any) -> Tuple[str, ...]:
+ """Normalize a collectives spec into a tuple of collective names.
+
+ Parameters
+ ----------
+ value : Any
+ Either a ``list``/``tuple`` of collective-name strings (from a
+ parsed YAML ``--config`` file or direct construction), or a
+ comma-separated string (from the CLI).
+
+ Returns
+ -------
+ tuple[str, ...]
+ The parsed collective names, in the order given, with surrounding
+ whitespace stripped and empty entries dropped (so a trailing comma
+ like ``"all_reduce,"`` does not produce a spurious ``""`` entry).
+
+ Examples
+ --------
+ >>> _parse_collectives("all_reduce, broadcast")
+ ('all_reduce', 'broadcast')
+ >>> _parse_collectives(["all_reduce", "broadcast"])
+ ('all_reduce', 'broadcast')
+ """
+ if isinstance(value, (list, tuple)):
+ return tuple(str(v) for v in value)
+ return tuple(part.strip() for part in str(value).split(",") if part.strip())
+
+
+@dataclasses.dataclass(frozen=True)
+class Config:
+ """Immutable configuration for one correlation-study AWS provisioning run.
+
+ All fields have defaults, so ``Config()`` alone yields a fully valid
+ configuration (one on-demand-or-spot ``p5.48xlarge`` in ``us-east-1``,
+ fully-connected topology, a freshly generated ``run_id``). Every
+ construction path (direct, :meth:`from_args`, :meth:`from_parsed`) runs
+ the same validation in :meth:`__post_init__`, so an invalid ``Config``
+ can never be observed by downstream code.
+
+ Parameters
+ ----------
+ region : str, default "us-east-1"
+ AWS region to provision in. Also used for the SSM AMI lookup and
+ must have the requested ``instance_type`` available.
+ availability_zone : str or None, default None
+ Specific AZ within ``region`` to pin the instance to. ``None``
+ lets AWS/the spot fleet choose.
+ instance_type : str, default "p5.48xlarge"
+ EC2 instance type. The 8x H100 NVSwitch topology this study
+ profiles is specific to ``p5.48xlarge``; other types are accepted
+ without validation but are not what the rest of this work package
+ was designed against.
+ purchasing : str, default "spot-then-ondemand"
+ One of ``"spot"``, ``"ondemand"``, ``"spot-then-ondemand"``. See
+ ``provision.launch_instance`` for the fallback semantics of the
+ combined mode.
+ topology : str, default "fc"
+ One of ``"fc"`` (fully connected, the NVSwitch-native topology),
+ ``"torus"`` (a logical torus overlay profiled on top of the same
+ physical node), or ``"both"``.
+ torus_dims : tuple[int, ...], default (2, 2, 2)
+ Per-axis dimensions of the logical torus. Precondition: the
+ product of all dimensions must equal 8 (one axis slot per GPU on
+ the node). Accepts a ``"2x2x2"``-style string on construction and
+ normalizes it to a tuple of ints.
+ collectives : tuple[str, ...], default (all_reduce, all_gather,
+ reduce_scatter, alltoall, broadcast, sendrecv)
+ NCCL collectives the (sibling) profiling scripts should sweep.
+ Accepts a comma-separated string on construction.
+ min_mib : int, default 1
+ Smallest message size, in MiB, in the profiling sweep.
+ Precondition: must be a power of two and ``<= max_mib``.
+ max_mib : int, default 1024
+ Largest message size, in MiB, in the profiling sweep.
+ Precondition: must be a power of two and ``>= min_mib``.
+ run_id : str or None, default None
+ Identifier used to name/tag every AWS resource created for this
+ run (instance, security group, key pair, state file). If left as
+ ``None``, a fresh id is generated in :meth:`__post_init__` as
+ ``"correl-" + UTC timestamp``. See the NOTE in
+ :meth:`__post_init__` for why generation happens there rather than
+ only in :meth:`from_args`.
+ tag_project : str, default "accelforge-correlation"
+ Value written to the ``Project`` tag on every AWS resource this
+ run creates; also what ``teardown.py`` filters
+ ``describe_instances`` on to discover a run's resources.
+ ssh_user : str, default "ubuntu"
+ Login user baked into the deep-learning AMI, used when printing
+ the ``ssh`` command at the end of provisioning.
+ ami_ssm_parameter : str, default
+ "/aws/service/deeplearning/ami/x86_64/base-oss-nvidia-driver-gpu-ubuntu-22.04/latest/ami-id"
+ SSM public parameter name that resolves to the latest matching
+ Deep Learning AMI id for ``region``.
+ root_volume_gb : int, default 200
+ Size, in GiB, of the root ``gp3`` EBS volume attached at
+ ``/dev/sda1``.
+ dead_man_minutes : int, default 120
+ Minutes of wall-clock time after which the on-instance dead-man
+ timer (armed by setup scripts owned by a sibling work package)
+ force-shuts-down the instance. This is a safety net independent of
+ ``teardown.py`` actually being run; see the README's guardrails
+ section.
+ key_dir : pathlib.Path, default "/keys"
+ Directory where generated SSH private keys are written.
+ state_dir : pathlib.Path, default "/.state"
+ Directory where per-run provisioning state JSON is written.
+
+ Raises
+ ------
+ ValueError
+ Raised by :meth:`__post_init__` if ``purchasing`` or ``topology``
+ is not one of the allowed values, if ``torus_dims`` does not
+ multiply out to 8 or contains a non-positive entry, if
+ ``min_mib > max_mib``, or if either ``min_mib`` or ``max_mib`` is
+ not a power of two.
+
+ Examples
+ --------
+ >>> cfg = Config(region="us-west-2", purchasing="ondemand")
+ >>> cfg.region, cfg.purchasing
+ ('us-west-2', 'ondemand')
+ >>> cfg.run_id is not None
+ True
+ """
+
+ region: str = "us-east-1"
+ availability_zone: Optional[str] = None
+ instance_type: str = "p5.48xlarge"
+ purchasing: str = "spot-then-ondemand"
+ topology: str = "fc"
+ torus_dims: Tuple[int, ...] = (2, 2, 2)
+ collectives: Tuple[str, ...] = _DEFAULT_COLLECTIVES
+ min_mib: int = 1
+ max_mib: int = 1024
+ run_id: Optional[str] = None
+ tag_project: str = "accelforge-correlation"
+ ssh_user: str = "ubuntu"
+ ami_ssm_parameter: str = (
+ "/aws/service/deeplearning/ami/x86_64/"
+ "base-oss-nvidia-driver-gpu-ubuntu-22.04/latest/ami-id"
+ )
+ root_volume_gb: int = 200
+ dead_man_minutes: int = 120
+ key_dir: Path = _THIS_DIR / "keys"
+ state_dir: Path = _THIS_DIR / ".state"
+
+ def __post_init__(self) -> None:
+ """Normalize field representations and validate all preconditions.
+
+ Raises
+ ------
+ ValueError
+ See the class docstring's ``Raises`` section; this method is
+ where every one of those checks is actually enforced.
+
+ Notes
+ -----
+ Runs for *every* construction path (``Config(...)`` directly,
+ :meth:`from_args`, :meth:`from_parsed`), because dataclasses always
+ call ``__post_init__`` after ``__init__``. This is a deliberate
+ choice over validating only inside :meth:`from_args`: it means a
+ test (or a future caller) that builds ``Config(purchasing="bogus")``
+ directly fails loudly at construction time instead of silently
+ producing an invalid config that only misbehaves once it reaches
+ AWS calls.
+ """
+ # Frozen dataclasses disallow `self.field = ...`; object.__setattr__
+ # is the standard, documented way to set fields from __post_init__.
+ object.__setattr__(self, "torus_dims", _parse_torus_dims(self.torus_dims))
+ object.__setattr__(self, "collectives", _parse_collectives(self.collectives))
+ object.__setattr__(self, "key_dir", Path(self.key_dir))
+ object.__setattr__(self, "state_dir", Path(self.state_dir))
+
+ if self.run_id is None:
+ # NOTE: the work-package spec describes run_id's default as
+ # "generated ... at parse time", which most literally refers to
+ # Config.from_args. We instead generate it here, in
+ # __post_init__, so every construction path gets a valid run_id
+ # -- see the docstring Notes above for why. This is the more
+ # conservative reading: it can never produce a Config with
+ # run_id=None reaching AWS tag values, which the "at parse
+ # time" phrasing on its own does not guarantee for direct
+ # `Config(...)` construction.
+ object.__setattr__(self, "run_id", _default_run_id())
+
+ if self.purchasing not in _ALLOWED_PURCHASING:
+ raise ValueError(
+ f"purchasing={self.purchasing!r} is not one of "
+ f"{sorted(_ALLOWED_PURCHASING)}"
+ )
+ if self.topology not in _ALLOWED_TOPOLOGY:
+ raise ValueError(
+ f"topology={self.topology!r} is not one of {sorted(_ALLOWED_TOPOLOGY)}"
+ )
+
+ # A negative-dimension axis (e.g. (-2, -2, 2)) could still multiply
+ # out to 8, silently passing a bare product check; reject it
+ # explicitly since it is never physically meaningful for a torus.
+ if any(d <= 0 for d in self.torus_dims):
+ raise ValueError(
+ f"torus_dims={self.torus_dims!r} must all be positive integers"
+ )
+ product = math.prod(self.torus_dims)
+ if product != _TORUS_GPU_COUNT:
+ raise ValueError(
+ f"torus_dims={self.torus_dims!r} has product {product}, "
+ f"expected {_TORUS_GPU_COUNT} (one p5.48xlarge node = "
+ f"{_TORUS_GPU_COUNT} GPUs)"
+ )
+
+ if self.min_mib > self.max_mib:
+ raise ValueError(
+ f"min_mib={self.min_mib} must be <= max_mib={self.max_mib}"
+ )
+ if not _is_power_of_two(self.min_mib):
+ raise ValueError(f"min_mib={self.min_mib} must be a power of two")
+ if not _is_power_of_two(self.max_mib):
+ raise ValueError(f"max_mib={self.max_mib} must be a power of two")
+
+ @classmethod
+ def add_args(cls, parser: argparse.ArgumentParser) -> None:
+ """Register every ``Config`` field as a kebab-case CLI flag.
+
+ Intended to be called by a script's own parser setup (see
+ ``provision.py``/``teardown.py``) *before* that script adds its
+ own extra flags (e.g. ``--dry-run``), so the two flag sets share
+ one ``argparse.ArgumentParser`` and one ``--help`` output.
+
+ Parameters
+ ----------
+ parser : argparse.ArgumentParser
+ Parser to add arguments to, mutated in place.
+
+ Notes
+ -----
+ Design: every flag added here uses ``default=argparse.SUPPRESS``
+ instead of the field's real default. This means an unset flag is
+ simply *absent* from the parsed namespace, which is exactly what
+ :meth:`from_parsed` needs to implement "CLI flags override
+ ``--config`` YAML values override dataclass defaults" -- if every
+ flag instead defaulted to its real value, :meth:`from_parsed` could
+ not distinguish "user explicitly passed the default value" from
+ "user didn't pass this flag at all", and CLI flags could never be
+ overridden by anything.
+ """
+ # Only used to render human-readable defaults into --help text;
+ # never used for the actual default values (see Notes above).
+ defaults = cls()
+
+ parser.add_argument(
+ "--config",
+ type=str,
+ default=None,
+ help=(
+ "Path to a YAML file of Config field overrides, applied "
+ "before CLI flags (CLI flags always win over --config)."
+ ),
+ )
+ parser.add_argument(
+ "--region",
+ type=str,
+ default=argparse.SUPPRESS,
+ help=f"AWS region (default: {defaults.region!r}).",
+ )
+ parser.add_argument(
+ "--availability-zone",
+ type=str,
+ default=argparse.SUPPRESS,
+ help="AWS availability zone, e.g. us-east-1a (default: let AWS choose).",
+ )
+ parser.add_argument(
+ "--instance-type",
+ type=str,
+ default=argparse.SUPPRESS,
+ help=f"EC2 instance type (default: {defaults.instance_type!r}).",
+ )
+ parser.add_argument(
+ "--purchasing",
+ type=str,
+ choices=sorted(_ALLOWED_PURCHASING),
+ default=argparse.SUPPRESS,
+ help=f"Purchasing strategy (default: {defaults.purchasing!r}).",
+ )
+ parser.add_argument(
+ "--topology",
+ type=str,
+ choices=sorted(_ALLOWED_TOPOLOGY),
+ default=argparse.SUPPRESS,
+ help=f"NCCL topology to profile (default: {defaults.topology!r}).",
+ )
+ parser.add_argument(
+ "--torus-dims",
+ type=_parse_torus_dims,
+ default=argparse.SUPPRESS,
+ help=(
+ "Logical torus dims as e.g. '2x2x2'; product must be 8 "
+ f"(default: {'x'.join(str(d) for d in defaults.torus_dims)!r})."
+ ),
+ )
+ parser.add_argument(
+ "--collectives",
+ type=_parse_collectives,
+ default=argparse.SUPPRESS,
+ help=(
+ "Comma-separated NCCL collectives to sweep "
+ f"(default: {','.join(defaults.collectives)!r})."
+ ),
+ )
+ parser.add_argument(
+ "--min-mib",
+ type=int,
+ default=argparse.SUPPRESS,
+ help=(
+ "Smallest message size in MiB, must be a power of two "
+ f"(default: {defaults.min_mib})."
+ ),
+ )
+ parser.add_argument(
+ "--max-mib",
+ type=int,
+ default=argparse.SUPPRESS,
+ help=(
+ "Largest message size in MiB, must be a power of two "
+ f"(default: {defaults.max_mib})."
+ ),
+ )
+ parser.add_argument(
+ "--run-id",
+ type=str,
+ default=argparse.SUPPRESS,
+ help=(
+ "Run identifier used to tag/name all resources "
+ "(default: generated as 'correl-')."
+ ),
+ )
+ parser.add_argument(
+ "--tag-project",
+ type=str,
+ default=argparse.SUPPRESS,
+ help=f"Value for the 'Project' tag on all resources (default: {defaults.tag_project!r}).",
+ )
+ parser.add_argument(
+ "--ssh-user",
+ type=str,
+ default=argparse.SUPPRESS,
+ help=f"SSH login user for the AMI (default: {defaults.ssh_user!r}).",
+ )
+ parser.add_argument(
+ "--ami-ssm-parameter",
+ type=str,
+ default=argparse.SUPPRESS,
+ help="SSM parameter name to resolve the AMI id from.",
+ )
+ parser.add_argument(
+ "--root-volume-gb",
+ type=int,
+ default=argparse.SUPPRESS,
+ help=f"Root EBS volume size in GiB (default: {defaults.root_volume_gb}).",
+ )
+ parser.add_argument(
+ "--dead-man-minutes",
+ type=int,
+ default=argparse.SUPPRESS,
+ help=(
+ "Minutes before the on-instance dead-man timer force-shuts-down "
+ f"the instance (default: {defaults.dead_man_minutes})."
+ ),
+ )
+ parser.add_argument(
+ "--key-dir",
+ type=Path,
+ default=argparse.SUPPRESS,
+ help=f"Directory to store the generated SSH private key (default: {defaults.key_dir}).",
+ )
+ parser.add_argument(
+ "--state-dir",
+ type=Path,
+ default=argparse.SUPPRESS,
+ help=f"Directory to store per-run provisioning state JSON (default: {defaults.state_dir}).",
+ )
+
+ @classmethod
+ def from_parsed(cls, namespace: argparse.Namespace) -> "Config":
+ """Build a :class:`Config` from an already-parsed argparse namespace.
+
+ Meant to be used together with :meth:`add_args`: a caller builds
+ one ``ArgumentParser``, calls ``Config.add_args(parser)``, adds its
+ own extra flags, calls ``parser.parse_args(argv)``, and passes the
+ resulting namespace here. Any namespace attributes that are not
+ ``Config`` field names (e.g. a caller's own ``--dry-run``) are
+ ignored, so the same namespace can safely be shared with
+ script-specific flags.
+
+ Parameters
+ ----------
+ namespace : argparse.Namespace
+ Parsed CLI arguments, as produced by
+ ``parser.parse_args(...)`` on a parser that included
+ :meth:`add_args`'s flags. If it has a ``config`` attribute
+ (from the ``--config`` flag) that is truthy, that path is
+ loaded as a YAML mapping of field overrides.
+
+ Returns
+ -------
+ Config
+ A validated ``Config`` built from, in increasing priority:
+ dataclass defaults, then ``--config`` YAML values, then
+ explicitly-passed CLI flags.
+
+ Raises
+ ------
+ ValueError
+ If ``--config`` points at a YAML document whose top level is
+ not a mapping, or if any field fails :meth:`__post_init__`
+ validation.
+ OSError
+ If ``--config`` points at a path that cannot be opened.
+ """
+ field_names = {f.name for f in dataclasses.fields(cls)}
+ values: Dict[str, Any] = {}
+
+ config_path = getattr(namespace, "config", None)
+ if config_path:
+ # Design: import PyYAML lazily, inside this branch, rather than
+ # at module top. Only the --config path needs it; a plain
+ # `provision.py --help` (or any run that never passes
+ # --config) must keep working even in an environment that
+ # only has boto3 installed and not PyYAML.
+ import yaml
+
+ with open(config_path, "r") as fh:
+ yaml_values = yaml.safe_load(fh)
+ if yaml_values is None:
+ yaml_values = {}
+ if not isinstance(yaml_values, dict):
+ raise ValueError(
+ f"--config file {config_path!r} must contain a top-level "
+ f"YAML mapping, got {type(yaml_values).__name__}"
+ )
+ # Silently drop unknown keys rather than raising: this lets a
+ # single shared YAML file carry keys meant for other tools
+ # (e.g. a future orchestrate.py section) without every
+ # consumer needing to know about every other consumer's keys.
+ values.update({k: v for k, v in yaml_values.items() if k in field_names})
+
+ # CLI flags win over --config values. Because add_args() gives
+ # every flag default=argparse.SUPPRESS, `namespace` only carries a
+ # key for a field the user actually typed on the command line, so
+ # this unconditional overwrite is exactly "CLI beats YAML beats
+ # dataclass default".
+ for key, value in vars(namespace).items():
+ if key in field_names:
+ values[key] = value
+
+ return cls(**values)
+
+ @classmethod
+ def from_args(cls, argv: Optional[List[str]] = None) -> "Config":
+ """Parse ``argv`` with a fresh, ``Config``-only parser.
+
+ Convenience wrapper around :meth:`add_args` + :meth:`from_parsed`
+ for callers that only need ``Config``'s own flags and do not have
+ any script-specific flags of their own to add.
+
+ Parameters
+ ----------
+ argv : list[str] or None, default None
+ Argument list to parse, as passed to
+ ``argparse.ArgumentParser.parse_args``. ``None`` means "read
+ from ``sys.argv[1:]``", argparse's own default behavior.
+
+ Returns
+ -------
+ Config
+ The parsed, validated configuration.
+
+ Examples
+ --------
+ >>> Config.from_args(["--region", "us-west-2", "--purchasing", "ondemand"]).region
+ 'us-west-2'
+ """
+ parser = argparse.ArgumentParser(
+ description="accelforge NCCL correlation-study provisioning config."
+ )
+ cls.add_args(parser)
+ namespace = parser.parse_args(argv)
+ return cls.from_parsed(namespace)
diff --git a/tests/not_working/distribuffers/__init__.py b/notebooks/astrasim2_correlation/correlation/data/.gitkeep
similarity index 100%
rename from tests/not_working/distribuffers/__init__.py
rename to notebooks/astrasim2_correlation/correlation/data/.gitkeep
diff --git a/notebooks/astrasim2_correlation/correlation/orchestrate.py b/notebooks/astrasim2_correlation/correlation/orchestrate.py
new file mode 100644
index 00000000..00d5dddd
--- /dev/null
+++ b/notebooks/astrasim2_correlation/correlation/orchestrate.py
@@ -0,0 +1,935 @@
+"""End-to-end CLI for the NCCL correlation study's empirical leg.
+
+This is the "run everything" work package: it ties together the
+provisioning infrastructure (``config.py``, ``provision.py``,
+``teardown.py``) and the profiling infrastructure (``setup_node.sh``,
+``run_profile.sh``, ``parse_nccl.py``, ``torus_bench/``) -- all owned by
+sibling work packages and imported/invoked here, never modified -- into one
+command an operator runs to go from "nothing provisioned" to "CSVs sitting
+in ``data///csv/`` and the instance torn down".
+
+Pipeline
+--------
+1. Parse CLI args into a :class:`config.Config` plus this script's own
+ ``--yes``/``--keep-alive``/``--dry-run``/``--ssh-cidr`` flags.
+2. Print the run header, cost warning, a wall-time estimate, and the leg
+ plan (via :func:`legs_for`); ask for interactive confirmation unless
+ ``--yes`` (skipped entirely for ``--dry-run``, which spends no money).
+3. Resolve the AMI, create the SSH key pair and security group, and call
+ ``provision.launch_instance`` -- exactly the sequence ``provision.py``
+ itself runs, reusing its functions directly rather than reimplementing
+ any of them. If ``--dry-run``, stop here (the key pair and security
+ group above were still created for real; see the ``--dry-run`` flag's
+ help text).
+4. Write the provisioning state file immediately, before waiting for SSH
+ (see the design comment at that call site for why).
+5. Wait for the instance to become SSH-reachable, then update the state
+ file with the now-known public IP.
+6. scp the profiling scripts and ``torus_bench/`` onto the instance.
+7. ssh in to run ``setup_node.sh`` (arms the dead-man timer, builds
+ nccl-tests/torus_bench).
+8. For each leg selected by ``--topology`` (see :func:`legs_for`), ssh in
+ to run ``run_profile.sh`` for the full collective sweep, then scp the
+ results back to ``data///``.
+9. In a ``finally`` block around steps 5-8: tear the instance down (unless
+ ``--keep-alive``), so a crash or a failed profiling leg never leaves an
+ (expensive, 8x H100) instance running unattended. See the design
+ comment on :func:`_teardown_and_cleanup` for how a teardown failure
+ itself is handled without masking whatever exception was already
+ propagating.
+
+Design: no new AWS/SSH logic here
+-----------------------------------
+Every AWS API call in this module goes through a ``provision.py`` or
+``teardown.py`` function that already exists, is already tested, and is
+already documented as part of this work package's contract (see those
+modules' docstrings). This module's own responsibility is narrower:
+sequencing those calls correctly, building ``ssh``/``scp`` argv lists
+(:func:`build_ssh_cmd`, :func:`build_scp_cmd`), and running them as
+subprocesses. Keeping that boundary sharp is also what makes this module
+testable without any real AWS/SSH/SCP access -- every seam it introduces
+(the two ``build_*_cmd`` functions, plus the imported provisioning
+functions) is a plain function that a test can monkeypatch or inspect the
+return value of, per this work package's "no AWS, no network, no ssh in
+tests" constraint.
+"""
+
+from __future__ import annotations
+
+import argparse
+import subprocess
+import sys
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+
+from config import Config
+from provision import (
+ _COST_WARNING,
+ _prompt_yes_no,
+ _require_boto3,
+ caller_ip,
+ ensure_key_pair,
+ ensure_security_group,
+ launch_instance,
+ resolve_ami,
+ wait_for_instance,
+ write_state,
+)
+from teardown import teardown_run
+
+# Design: guarded import, matching provision.py/teardown.py's own
+# convention (see provision.py's module docstring for the full rationale)
+# -- so `python orchestrate.py --help` keeps working even in a Python
+# environment that lacks boto3, since argparse's own --help handling exits
+# before main() ever reaches _require_boto3(). Importing `boto3` here as a
+# module-level name of orchestrate.py's own (rather than reaching into
+# `provision.boto3`) keeps this module's boto3.client(...) calls readable
+# without poking at another module's internals; _require_boto3() (reused
+# from provision.py, not redefined) is still what actually validates
+# availability before any client is constructed, since both imports
+# resolve to the same cached sys.modules entry (or both to None) in any
+# given interpreter.
+try:
+ import boto3
+except ImportError: # pragma: no cover - exercised only when boto3 truly absent
+ boto3 = None
+
+# ---------------------------------------------------------------------------
+# Module-level path constants
+# ---------------------------------------------------------------------------
+
+# Anchor every sibling-file path to this file's own directory (not the
+# process cwd), matching config.py's identical _THIS_DIR convention -- so
+# `python orchestrate.py` behaves the same regardless of the caller's shell
+# cwd.
+_THIS_DIR = Path(__file__).resolve().parent
+
+_SETUP_NODE_SH = _THIS_DIR / "setup_node.sh"
+_RUN_PROFILE_SH = _THIS_DIR / "run_profile.sh"
+_PARSE_NCCL_PY = _THIS_DIR / "parse_nccl.py"
+_TORUS_BENCH_DIR = _THIS_DIR / "torus_bench"
+
+# Design: expose the fetched-results root as its own module-level constant
+# (rather than inlining `_THIS_DIR / "data"` at the one call site) purely
+# so tests can monkeypatch `orchestrate._DATA_DIR` to a tmp_path and
+# guarantee the end-to-end test never writes into the real repository's
+# data/ directory -- every other AWS/ssh/scp side effect in a test is
+# already monkeypatched away, and this is the one remaining plain
+# filesystem write main() would otherwise perform unconditionally.
+_DATA_DIR = _THIS_DIR / "data"
+
+# run_profile.sh's argument for the FC leg: the profiling scripts
+# treat "8" as "all 8 GPUs, fully connected" -- there is no logical torus
+# shape to describe for that leg (see run_profile.sh's `-g 8` on the FC
+# path), unlike the torus leg where is a "DxDx..." shape string.
+_FC_DIMS = "8"
+
+_BYTES_PER_MIB = 2**20
+
+
+# ---------------------------------------------------------------------------
+# Pure helpers (no I/O, no AWS, no subprocess) -- kept separate from main()
+# specifically so they are trivially unit-testable per this work package's
+# "no AWS, no network, no ssh in tests" constraint.
+# ---------------------------------------------------------------------------
+
+
+def legs_for(topology: str) -> List[str]:
+ """Expand a ``Config.topology`` value into the ordered list of legs to run.
+
+ Parameters
+ ----------
+ topology : str
+ Typically ``cfg.topology``, one of ``"fc"``, ``"torus"``, or
+ ``"both"`` (``Config.__post_init__`` already validates this, so
+ this function does not re-validate it -- see Notes).
+
+ Returns
+ -------
+ list[str]
+ ``["fc", "torus"]`` if ``topology == "both"``; otherwise the
+ single-element list ``[topology]``.
+
+ Notes
+ -----
+ Deliberately permissive for any value other than ``"both"``: it simply
+ echoes that value back as a one-element list rather than checking it
+ against the allowed set. Re-validating here would duplicate
+ ``Config.__post_init__``'s already-authoritative check for no benefit,
+ since every caller in this module only ever passes an already-validated
+ ``cfg.topology``.
+
+ Examples
+ --------
+ >>> legs_for("both")
+ ['fc', 'torus']
+ >>> legs_for("fc")
+ ['fc']
+ >>> legs_for("torus")
+ ['torus']
+ """
+ if topology == "both":
+ return ["fc", "torus"]
+ return [topology]
+
+
+def _dims_for_leg(leg: str, torus_dims: Tuple[int, ...]) -> str:
+ """Compute run_profile.sh's ```` argument for one leg.
+
+ Parameters
+ ----------
+ leg : str
+ ``"fc"`` or ``"torus"``.
+ torus_dims : tuple[int, ...]
+ Per-axis torus dimensions, e.g. ``(2, 2, 2)`` (only consulted when
+ ``leg == "torus"``).
+
+ Returns
+ -------
+ str
+ :data:`_FC_DIMS` (``"8"``) for the FC leg; otherwise
+ ``torus_dims`` joined with ``"x"``, e.g. ``"2x2x2"``.
+
+ Examples
+ --------
+ >>> _dims_for_leg("fc", (2, 2, 2))
+ '8'
+ >>> _dims_for_leg("torus", (2, 2, 2))
+ '2x2x2'
+ """
+ if leg == "fc":
+ return _FC_DIMS
+ return "x".join(str(d) for d in torus_dims)
+
+
+def _estimate_wall_time_message(legs: List[str]) -> str:
+ """Build the human-readable wall-time estimate printed before confirmation.
+
+ Parameters
+ ----------
+ legs : list[str]
+ The leg plan, as returned by :func:`legs_for`.
+
+ Returns
+ -------
+ str
+ A one-line estimate: ~30-60 minutes when both legs are selected
+ (the work-package spec's own figure for a full FC+torus sweep),
+ or roughly half that for a single leg.
+
+ Notes
+ -----
+ This is a coarse, documented *estimate* for setting operator
+ expectations before a real-money confirmation prompt, not a measured
+ or SLA'd figure -- actual time depends on instance boot time, spot
+ availability, and the exact collective/message-size sweep configured.
+ """
+ if len(legs) >= 2:
+ return (
+ "Estimated wall time: ~30-60 minutes for both legs "
+ "(excludes instance boot/provisioning time)."
+ )
+ return (
+ "Estimated wall time: ~15-30 minutes for a single leg "
+ "(excludes instance boot/provisioning time)."
+ )
+
+
+def build_ssh_cmd(
+ key_path: Path,
+ user: str,
+ ip: str,
+ remote_cmd: str,
+ known_hosts_path: Path,
+) -> List[str]:
+ """Build an ``ssh`` argv list to run one command on the provisioned instance.
+
+ Parameters
+ ----------
+ key_path : pathlib.Path
+ Path to the local PEM private key (as returned by
+ ``provision.ensure_key_pair``).
+ user : str
+ SSH login user, typically ``cfg.ssh_user``.
+ ip : str
+ Target host's public IP address.
+ remote_cmd : str
+ The full remote command line to execute, e.g.
+ ``"DEADMAN_MINUTES=120 bash ~/setup_node.sh"``. Passed to ``ssh``
+ as a single trailing argv element; ``ssh`` hands it to the remote
+ login shell for interpretation, so ordinary shell syntax (env var
+ prefixes, multiple space-separated arguments) works as expected
+ without any extra quoting from this function.
+ known_hosts_path : pathlib.Path
+ Path to a per-run known-hosts file (typically
+ ``cfg.state_dir / "known_hosts"``). Kept as an explicit parameter
+ (rather than a hardcoded/global path) so this function stays pure
+ and independently testable -- see the module docstring's "Design:
+ no new AWS/SSH logic here" section.
+
+ .. note::
+ NOTE ON SPEC DEVIATION: the work-package spec lists this
+ function's signature as ``build_ssh_cmd(key_path, user, ip,
+ remote_cmd)`` -- four parameters -- but also requires the
+ ``-o UserKnownHostsFile=/known_hosts`` option, which
+ cannot be constructed without knowing ``state_dir`` from
+ *somewhere*. Reaching into a module-global ``Config`` from
+ inside this function would break the "pure helper, easy to
+ unit-test" property the spec explicitly asks for. Adding
+ ``known_hosts_path`` as an explicit fifth parameter is the
+ minimal change that preserves purity/testability; every
+ call site in this module passes ``cfg.state_dir / "known_hosts"``
+ for it, matching the spec's intent.
+
+ Returns
+ -------
+ list[str]
+ argv list of the form ``["ssh", "-i", , "-o",
+ "StrictHostKeyChecking=accept-new", "-o", "UserKnownHostsFile=<...>",
+ "-o", "ConnectTimeout=30", "@", ]``, ready to
+ pass to ``subprocess.run``.
+
+ Examples
+ --------
+ >>> build_ssh_cmd(Path("/k/id.pem"), "ubuntu", "1.2.3.4", "echo hi", Path("/s/known_hosts"))
+ ... # doctest: +NORMALIZE_WHITESPACE
+ ['ssh', '-i', '/k/id.pem', '-o', 'StrictHostKeyChecking=accept-new',
+ '-o', 'UserKnownHostsFile=/s/known_hosts', '-o', 'ConnectTimeout=30',
+ 'ubuntu@1.2.3.4', 'echo hi']
+ """
+ return [
+ "ssh",
+ "-i",
+ str(key_path),
+ "-o",
+ "StrictHostKeyChecking=accept-new",
+ "-o",
+ f"UserKnownHostsFile={known_hosts_path}",
+ "-o",
+ "ConnectTimeout=30",
+ f"{user}@{ip}",
+ remote_cmd,
+ ]
+
+
+def build_scp_cmd(
+ key_path: Path,
+ sources: List[str],
+ dest: str,
+ known_hosts_path: Path,
+ recursive: bool = False,
+) -> List[str]:
+ """Build an ``scp`` argv list to copy one or more paths to/from the instance.
+
+ Parameters
+ ----------
+ key_path : pathlib.Path
+ Path to the local PEM private key.
+ sources : list[str]
+ Source path(s) to copy, in whatever form ``scp`` accepts: plain
+ local paths for a push, or a single ``"user@ip:remote/path"``
+ string for a fetch. Must be non-empty.
+ dest : str
+ Destination, again in whatever form ``scp`` accepts (a local
+ directory for a fetch, or ``"user@ip:remote/path"`` for a push).
+ known_hosts_path : pathlib.Path
+ Path to a per-run known-hosts file. See :func:`build_ssh_cmd`'s
+ docstring for why this is an explicit parameter rather than an
+ implicit global (the same rationale applies here).
+ recursive : bool, default False
+ If ``True``, prepend ``-r`` (needed for copying a directory, e.g.
+ ``torus_bench/`` or a remote ``results_/`` directory).
+
+ Returns
+ -------
+ list[str]
+ argv list of the form ``["scp", ["-r"], "-i", , "-o",
+ "StrictHostKeyChecking=accept-new", "-o",
+ "UserKnownHostsFile=<...>", "-o", "ConnectTimeout=30", *sources,
+ dest]``, ready to pass to ``subprocess.run``. The ``-r`` flag (when
+ present) is placed immediately after ``"scp"``, before every other
+ option, so its position is fixed and independently assertable in
+ tests regardless of how many sources are given.
+
+ Raises
+ ------
+ ValueError
+ If ``sources`` is empty -- an ``scp`` invocation with no source
+ path is never meaningful and would otherwise fail cryptically at
+ the OS level instead of at this argv-building step.
+
+ Examples
+ --------
+ >>> build_scp_cmd(Path("/k/id.pem"), ["a.sh", "b.sh"], "ubuntu@1.2.3.4:~/", Path("/s/known_hosts"))
+ ... # doctest: +NORMALIZE_WHITESPACE
+ ['scp', '-i', '/k/id.pem', '-o', 'StrictHostKeyChecking=accept-new',
+ '-o', 'UserKnownHostsFile=/s/known_hosts', '-o', 'ConnectTimeout=30',
+ 'a.sh', 'b.sh', 'ubuntu@1.2.3.4:~/']
+ >>> build_scp_cmd(Path("/k/id.pem"), ["dir"], "ubuntu@1.2.3.4:~/dir", Path("/s/known_hosts"), recursive=True)[:2]
+ ['scp', '-r']
+ """
+ if not sources:
+ raise ValueError("build_scp_cmd requires at least one source path")
+
+ cmd: List[str] = ["scp"]
+ if recursive:
+ cmd.append("-r")
+ cmd += [
+ "-i",
+ str(key_path),
+ "-o",
+ "StrictHostKeyChecking=accept-new",
+ "-o",
+ f"UserKnownHostsFile={known_hosts_path}",
+ "-o",
+ "ConnectTimeout=30",
+ ]
+ cmd += list(sources)
+ cmd.append(dest)
+ return cmd
+
+
+# ---------------------------------------------------------------------------
+# Subprocess execution
+# ---------------------------------------------------------------------------
+
+
+def _run_streaming(cmd: List[str]) -> None:
+ """Run an ``ssh``/``scp`` argv list, streaming its output live.
+
+ Parameters
+ ----------
+ cmd : list[str]
+ argv list, typically from :func:`build_ssh_cmd` or
+ :func:`build_scp_cmd`.
+
+ Raises
+ ------
+ subprocess.CalledProcessError
+ If the subprocess exits with a non-zero status (``check=True``).
+
+ Notes
+ -----
+ Deliberately does not pass ``capture_output``/``stdout``/``stderr`` --
+ the child process's output streams straight through to this process's
+ own stdout/stderr, so an operator watching a multi-minute profiling
+ sweep sees live progress rather than a silent hang followed by a wall
+ of buffered text at the end. Prints the command line first (to stdout)
+ so the corresponding output block is identifiable in a long combined
+ log.
+ """
+ print("+ " + " ".join(cmd))
+ subprocess.run(cmd, check=True)
+
+
+# ---------------------------------------------------------------------------
+# Pipeline stages -- each a thin, single-responsibility wrapper around a
+# handful of _run_streaming calls, factored out of main() so that function
+# stays a readable top-to-bottom sequence rather than one long body.
+# ---------------------------------------------------------------------------
+
+
+def _push_files(cfg: Config, key_path: Path, public_ip: str, known_hosts_path: Path) -> None:
+ """scp the profiling scripts and torus_bench/ onto the instance.
+
+ Parameters
+ ----------
+ cfg : Config
+ Run configuration; only ``cfg.ssh_user`` is consulted directly
+ (the rest flows through ``key_path``/``public_ip``/
+ ``known_hosts_path``).
+ key_path : pathlib.Path
+ Local PEM private key path.
+ public_ip : str
+ Instance's public IP, as returned by ``provision.wait_for_instance``.
+ known_hosts_path : pathlib.Path
+ Per-run known-hosts file path.
+
+ Notes
+ -----
+ Two separate ``scp`` invocations, matching the two different transfer
+ shapes: (1) three individual files pushed non-recursively to the
+ remote home directory, and (2) the ``torus_bench/`` directory pushed
+ recursively to ``~/torus_bench`` so ``setup_node.sh`` can build it.
+ Side effect: two subprocess invocations over the network to the
+ instance.
+ """
+ remote_home = f"{cfg.ssh_user}@{public_ip}:~/"
+ print("=== Pushing profiling scripts (setup_node.sh, run_profile.sh, parse_nccl.py) ===")
+ _run_streaming(
+ build_scp_cmd(
+ key_path,
+ [str(_SETUP_NODE_SH), str(_RUN_PROFILE_SH), str(_PARSE_NCCL_PY)],
+ remote_home,
+ known_hosts_path,
+ )
+ )
+
+ remote_torus_dir = f"{cfg.ssh_user}@{public_ip}:~/torus_bench"
+ print("=== Pushing torus_bench/ ===")
+ _run_streaming(
+ build_scp_cmd(
+ key_path,
+ [str(_TORUS_BENCH_DIR)],
+ remote_torus_dir,
+ known_hosts_path,
+ recursive=True,
+ )
+ )
+
+
+def _run_setup(cfg: Config, key_path: Path, public_ip: str, known_hosts_path: Path) -> None:
+ """ssh in and run ``setup_node.sh`` (arms dead-man timer, builds binaries).
+
+ Parameters
+ ----------
+ cfg : Config
+ Run configuration; ``cfg.dead_man_minutes`` is forwarded to
+ ``setup_node.sh`` as the ``DEADMAN_MINUTES`` environment variable.
+ key_path : pathlib.Path
+ Local PEM private key path.
+ public_ip : str
+ Instance's public IP.
+ known_hosts_path : pathlib.Path
+ Per-run known-hosts file path.
+
+ Notes
+ -----
+ Must run after :func:`_push_files` (``setup_node.sh`` and
+ ``torus_bench/`` must already be on the instance) and before any
+ profiling leg (``setup_node.sh`` is what builds the nccl-tests and
+ torus_bench binaries those legs depend on, and arms the on-instance
+ dead-man safety timer). Side effect: one subprocess invocation over
+ the network to the instance.
+ """
+ remote_cmd = f"DEADMAN_MINUTES={cfg.dead_man_minutes} bash ~/setup_node.sh"
+ print("=== Running setup_node.sh ===")
+ _run_streaming(build_ssh_cmd(key_path, cfg.ssh_user, public_ip, remote_cmd, known_hosts_path))
+
+
+def _run_leg(
+ cfg: Config,
+ leg: str,
+ key_path: Path,
+ public_ip: str,
+ known_hosts_path: Path,
+ run_data_dir: Path,
+) -> None:
+ """Run one profiling leg on the instance and fetch its results locally.
+
+ Parameters
+ ----------
+ cfg : Config
+ Run configuration (``torus_dims``, ``min_mib``/``max_mib``,
+ ``collectives``, ``ssh_user``).
+ leg : str
+ ``"fc"`` or ``"torus"``.
+ key_path : pathlib.Path
+ Local PEM private key path.
+ public_ip : str
+ Instance's public IP.
+ known_hosts_path : pathlib.Path
+ Per-run known-hosts file path.
+ run_data_dir : Path
+ Local directory this run's results should land under (typically
+ ``_DATA_DIR / cfg.run_id``); this leg's results land specifically
+ at ``run_data_dir / leg``.
+
+ Notes
+ -----
+ Side effects: one ``ssh`` subprocess invocation that runs the full
+ collective sweep for this leg on the instance (this is the
+ long-running step -- see :func:`_estimate_wall_time_message`), a
+ ``mkdir`` of ``run_data_dir`` (the PARENT of this leg's fetch
+ destination -- see the inline comment above that call for why the leaf
+ directory itself is deliberately left uncreated), and one ``scp -r``
+ subprocess invocation that fetches the results back, creating the new
+ local directory tree at ``run_data_dir / leg`` itself.
+
+ Prints the list of fetched ``*.csv`` files at the end, so an operator
+ watching the run can immediately confirm data landed without a
+ separate ``ls``.
+
+ Raises
+ ------
+ RuntimeError
+ If ``run_data_dir / leg`` already exists before the fetch -- see
+ the inline comment above the check for why this is treated as an
+ error rather than silently proceeding.
+ """
+ dims = _dims_for_leg(leg, cfg.torus_dims)
+ min_bytes = cfg.min_mib * _BYTES_PER_MIB
+ max_bytes = cfg.max_mib * _BYTES_PER_MIB
+ remote_results_dir = f"~/results_{leg}"
+ remote_cmd = (
+ f"bash ~/run_profile.sh {remote_results_dir} {leg} {min_bytes} {max_bytes} {dims} "
+ f"{' '.join(cfg.collectives)}"
+ )
+
+ print(f"=== Profiling leg: {leg} (dims={dims}, {cfg.min_mib}-{cfg.max_mib} MiB) ===")
+ _run_streaming(build_ssh_cmd(key_path, cfg.ssh_user, public_ip, remote_cmd, known_hosts_path))
+
+ # Design: fetch into a leaf directory (run_data_dir/leg) that must NOT
+ # already exist when scp runs. `scp -r user@host:~/results_fc
+ # ` renames the copied directory to when
+ # doesn't yet exist, landing its contents (csv/, raw/,
+ # metadata.txt) directly inside it -- exactly the data///
+ # layout this work package's spec requires. Pre-creating that leaf
+ # directory first would instead make scp nest an extra results_fc/
+ # level inside it, since scp's "copy into vs. rename to" behavior
+ # depends on whether the destination path already exists.
+ local_leg_dir = run_data_dir / leg
+ if local_leg_dir.exists():
+ # A pre-existing leaf directory means a previous run already fetched
+ # results for this exact run_id+leg combination (or something else
+ # created the path). Silently proceeding would make scp nest an
+ # extra results_/ level inside the existing directory instead
+ # of landing csv/raw/metadata.txt directly in it (see the comment
+ # above), corrupting the data/// layout without any
+ # error -- raise instead so a re-run with a colliding run_id+leg
+ # fails loudly here rather than silently mis-nesting fetched data.
+ raise RuntimeError(
+ f"{local_leg_dir} already exists; a previous fetch for run_id="
+ f"{cfg.run_id!r} leg={leg!r} already landed results there. "
+ "Refusing to scp into it again (that would nest an extra "
+ f"results_{leg}/ level inside the existing directory instead of "
+ "csv/raw/metadata.txt landing directly in it). Remove or rename "
+ "the existing directory, or re-run with a different --run-id."
+ )
+ # Only the PARENT (run_data_dir) is created here -- local_leg_dir itself
+ # must be left absent (see the comment above) for scp's rename-vs-nest
+ # behavior to produce the right layout. Without this mkdir, the very
+ # first leg fetched under a fresh run_id would fail outright: scp cannot
+ # write to / if itself doesn't exist
+ # yet (parents=True/exist_ok=True mirrors provision.write_state's own
+ # "create the directory, then write into it" idiom).
+ run_data_dir.mkdir(parents=True, exist_ok=True)
+ fetch_source = f"{cfg.ssh_user}@{public_ip}:{remote_results_dir}"
+ print(f"=== Fetching {leg} leg results ===")
+ _run_streaming(
+ build_scp_cmd(
+ key_path,
+ [fetch_source],
+ str(local_leg_dir),
+ known_hosts_path,
+ recursive=True,
+ )
+ )
+
+ csv_files = sorted((local_leg_dir / "csv").glob("*.csv"))
+ print(f"Fetched {leg} leg results to {local_leg_dir}:")
+ for csv_path in csv_files:
+ print(f" {csv_path}")
+
+
+def _print_keep_alive_notice(cfg: Config, key_path: Path, state: Dict[str, Any]) -> None:
+ """Print the SSH command and a loud cost reminder for a ``--keep-alive`` run.
+
+ Parameters
+ ----------
+ cfg : Config
+ Run configuration (``ssh_user``, ``run_id``, ``dead_man_minutes``).
+ key_path : pathlib.Path
+ Local PEM private key path, included in the printed SSH command.
+ state : dict
+ This run's state dict; ``state["public_ip"]`` is used if present.
+
+ Notes
+ -----
+ Called from :func:`main`'s ``finally`` block in place of
+ :func:`_teardown_and_cleanup` when ``--keep-alive`` was passed. Does
+ not delete the state file (unlike the teardown path), since
+ ``teardown.py`` needs it later to find and tear down this
+ still-running instance.
+ """
+ ip_display = state.get("public_ip") or ""
+ print("=" * 70)
+ print("--keep-alive set: instance left RUNNING. YOU ARE STILL BEING BILLED.")
+ print(f"SSH command: ssh -i {key_path} {cfg.ssh_user}@{ip_display}")
+ print(
+ f"Dead-man shutdown deadline: ~{cfg.dead_man_minutes} minutes after "
+ "setup_node.sh armed it (this is a backstop independent of teardown.py)."
+ )
+ print(f"Tear down manually with: python teardown.py --run-id {cfg.run_id}")
+ print("=" * 70)
+
+
+def _teardown_and_cleanup(
+ ec2_client, state: Dict[str, Any], state_path: Path, cfg: Config
+) -> None:
+ """Tear down this run's AWS resources and remove its local state file.
+
+ Parameters
+ ----------
+ ec2_client : botocore.client.BaseClient
+ A boto3 ``ec2`` client (or a stub thereof).
+ state : dict
+ This run's state dict, as passed to ``teardown.teardown_run``.
+ state_path : pathlib.Path
+ Path to this run's ``.state/.json`` file, removed only if
+ teardown succeeds.
+ cfg : Config
+ Run configuration, used only to print ``cfg.run_id`` into the
+ follow-up messages below.
+
+ Notes
+ -----
+ Called from :func:`main`'s ``finally`` block, which may itself be
+ running while an exception from the profiling steps (e.g. a failed
+ ``run_profile.sh`` invocation) is already propagating out of the
+ surrounding ``try``. This is exactly the scenario the work-package
+ spec calls out: a ``finally`` block that itself raises would, under
+ Python's ordinary exception semantics, cause *that new* exception to
+ be what the caller sees instead of the original one -- effectively
+ masking the original failure behind a teardown failure, even though
+ both are independently worth surfacing. This function therefore
+ catches any exception ``teardown_run`` raises, prints it (and does not
+ delete the state file, since a failed teardown may have left resources
+ behind that ``teardown.py --run-id `` will still need it to
+ find), and simply returns -- letting whatever exception was already
+ active in the caller's ``try`` block continue propagating undisturbed.
+ """
+ print(f"Tearing down run {cfg.run_id}...")
+ try:
+ teardown_run(ec2_client, state, delete_key=False)
+ except Exception as teardown_exc: # noqa: BLE001
+ # Design/WHY deliberately broad and deliberately NOT re-raised:
+ # see this function's docstring Notes above. Printing to stderr
+ # (rather than raising) is what prevents this teardown failure
+ # from masking an in-flight exception from the try block this
+ # finally belongs to.
+ print(f"ERROR: teardown failed: {teardown_exc}", file=sys.stderr)
+ print(
+ "Manual cleanup required -- the instance may still be running. "
+ f"Try: python teardown.py --run-id {cfg.run_id}",
+ file=sys.stderr,
+ )
+ return
+
+ if state_path.exists():
+ state_path.unlink()
+ print(f"Removed state file {state_path}.")
+ print("Recommended audit: python teardown.py --verify")
+
+
+# ---------------------------------------------------------------------------
+# CLI entry point
+# ---------------------------------------------------------------------------
+
+
+def main(argv: Optional[List[str]] = None) -> int:
+ """Parse CLI args and run the full provision -> profile -> teardown pipeline.
+
+ Parameters
+ ----------
+ argv : list[str] or None, default None
+ Argument list, as passed to ``argparse``'s ``parse_args``. ``None``
+ reads from ``sys.argv[1:]``.
+
+ Returns
+ -------
+ int
+ ``0`` on success (including a completed ``--dry-run``), ``1`` if
+ the user declined the interactive confirmation prompt.
+
+ Raises
+ ------
+ Exception
+ Any exception raised by provisioning, SSH/SCP subprocess failures
+ (``subprocess.CalledProcessError``), or waiting for the instance
+ propagates out of this function uncaught -- per the work-package
+ spec, a profiling/provisioning failure is a real failure and
+ should surface as one (a nonzero process exit code via
+ ``raise SystemExit(main())`` in ``__main__``), not be silently
+ downgraded to a return code. The ``finally`` block described below
+ still runs before that propagation completes.
+
+ Notes
+ -----
+ Side effects (skipped when ``--dry-run`` is passed, past the point
+ where the SSH key pair and security group are created -- see the
+ ``--dry-run`` flag's help text): creates an SSH key pair and security
+ group, launches an instance, writes/updates a state JSON file, scp's
+ profiling scripts and results to/from the instance, ssh's in to run
+ setup and profiling commands, and (unless ``--keep-alive``) tears the
+ instance down again.
+
+ Design/WHY the state file is written immediately after
+ ``launch_instance`` returns, before ``wait_for_instance`` is even
+ called: ``wait_for_instance`` can block for several minutes (instance
+ boot, then polling for SSH) and can itself raise (``TimeoutError``,
+ ``WaiterError``). If this controller process crashes or is killed
+ during that wait, an instance is still running and being billed with
+ no local record of it unless the state file was already written
+ beforehand. Writing state right after launch -- with ``public_ip`` as
+ an explicit placeholder, filled in by a second ``write_state`` call
+ once ``wait_for_instance`` returns -- means ``teardown.py`` (which
+ also cross-checks live AWS tags, not just local state, per its own
+ docstring) can find and tear down this run from local state alone even
+ in that crash scenario, without waiting on SSH reachability first.
+ """
+ parser = argparse.ArgumentParser(
+ prog="orchestrate.py",
+ description=(
+ "End-to-end NCCL correlation-study orchestration: provision one "
+ "p5.48xlarge instance, push profiling scripts, run the FC "
+ "and/or torus profiling sweep, fetch results into data//, "
+ "and tear the instance down."
+ ),
+ )
+ Config.add_args(parser)
+ parser.add_argument(
+ "--dry-run",
+ action="store_true",
+ help=(
+ "Perform DryRun=True authorization checks only; launch no "
+ "instance and run no profiling. NOTE: a real SSH key pair and "
+ "security group ARE still created in AWS even with --dry-run, "
+ "mirroring provision.py's documented --dry-run semantics -- "
+ "only launch_instance's actual launch and everything after it "
+ "(waiting for SSH, pushing files, profiling, teardown) is "
+ "skipped."
+ ),
+ )
+ parser.add_argument(
+ "--keep-alive",
+ action="store_true",
+ help=(
+ "Skip automatic teardown after profiling completes (or fails); "
+ "leave the instance running and print its SSH command instead "
+ "of tearing it down. The on-instance dead-man timer "
+ "(--dead-man-minutes) still applies regardless of this flag -- "
+ "it only disables orchestrate.py's own teardown call, not that "
+ "backstop."
+ ),
+ )
+ parser.add_argument(
+ "--ssh-cidr",
+ type=str,
+ default=None,
+ help=(
+ "CIDR block allowed SSH access, e.g. 1.2.3.4/32. Overrides "
+ "auto-detected caller IP (see provision.caller_ip())."
+ ),
+ )
+ parser.add_argument(
+ "--yes",
+ action="store_true",
+ help="Skip the interactive cost-confirmation prompt (not required for --dry-run).",
+ )
+ args = parser.parse_args(argv)
+ cfg = Config.from_parsed(args)
+
+ # Fail fast on a missing boto3 before printing anything else, matching
+ # provision.py's own main() -- no point asking the operator to confirm
+ # a cost warning for a run that cannot possibly proceed.
+ _require_boto3()
+
+ legs = legs_for(cfg.topology)
+ print(f"=== accelforge correlation-study orchestration: run_id={cfg.run_id} ===")
+ print(_COST_WARNING)
+ print(_estimate_wall_time_message(legs))
+ print(f"Legs to profile: {', '.join(legs)} (topology={cfg.topology!r})")
+ print(
+ f"Purchasing mode: {cfg.purchasing}. Instance type: {cfg.instance_type}. "
+ f"Region: {cfg.region}."
+ )
+
+ if not args.dry_run and not args.yes:
+ if not _prompt_yes_no("Proceed with provisioning and profiling? [yes/N]: "):
+ print("Aborted by user.")
+ return 1
+
+ ec2_client = boto3.client("ec2", region_name=cfg.region)
+ ssm_client = boto3.client("ssm", region_name=cfg.region)
+
+ ssh_cidr = args.ssh_cidr
+ if not ssh_cidr:
+ ip = caller_ip()
+ ssh_cidr = f"{ip}/32"
+ print(f"SSH will be allowed from: {ssh_cidr}")
+
+ ami_id = resolve_ami(ssm_client, cfg.ami_ssm_parameter)
+ print(f"Resolved AMI: {ami_id}")
+
+ key_name = f"{cfg.tag_project}-{cfg.run_id}"
+ key_path = ensure_key_pair(ec2_client, key_name, cfg.key_dir)
+ print(f"Key pair ready: {key_name} -> {key_path}")
+
+ sg_name = f"{cfg.tag_project}-{cfg.run_id}-sg"
+ sg_id = ensure_security_group(ec2_client, sg_name, ssh_cidr, cfg.tag_project, cfg.run_id)
+ print(f"Security group ready: {sg_id}")
+
+ # Mirrors provision.py's main(): the dry_run flag flows all the way
+ # into launch_instance (which makes a real, DryRun=True API call) so a
+ # dry run exercises the exact same request-building code path a real
+ # launch would, rather than short-circuiting before this call.
+ launch_result = launch_instance(
+ ec2_client, cfg, ami_id, sg_id, key_name, dry_run=args.dry_run
+ )
+
+ if args.dry_run:
+ print(
+ f"Dry run complete (purchasing checked: {launch_result['purchasing_used']}); "
+ "no instance was launched, no SSH was attempted, and no profiling ran. "
+ "The key pair and security group above WERE created for real -- "
+ f"clean them up with: python teardown.py --run-id {cfg.run_id} --delete-key"
+ )
+ return 0
+
+ instance_id = launch_result["instance_id"]
+ purchasing_used = launch_result["purchasing_used"]
+ print(f"Launched instance {instance_id} ({purchasing_used}).")
+
+ # See this function's "Design/WHY" docstring Notes above for the full
+ # rationale: write state now, with public_ip left as an explicit
+ # placeholder, rather than waiting until wait_for_instance (which can
+ # block for minutes and can itself fail) returns.
+ state: Dict[str, Any] = {
+ "run_id": cfg.run_id,
+ "region": cfg.region,
+ "instance_id": instance_id,
+ "sg_id": sg_id,
+ "key_name": key_name,
+ "key_path": str(key_path),
+ "public_ip": None,
+ "purchasing_used": purchasing_used,
+ "ami_id": ami_id,
+ }
+ state_path = write_state(cfg.state_dir, cfg.run_id, state)
+ print(f"State written to: {state_path} (public_ip pending SSH reachability).")
+
+ known_hosts_path = cfg.state_dir / "known_hosts"
+ run_data_dir = _DATA_DIR / cfg.run_id
+
+ try:
+ public_ip = wait_for_instance(ec2_client, instance_id)
+ state["public_ip"] = public_ip
+ write_state(cfg.state_dir, cfg.run_id, state)
+ print(f"Instance is running and SSH-reachable at {public_ip}")
+ print(f"SSH command: ssh -i {key_path} {cfg.ssh_user}@{public_ip}")
+
+ _push_files(cfg, key_path, public_ip, known_hosts_path)
+ _run_setup(cfg, key_path, public_ip, known_hosts_path)
+
+ for leg in legs:
+ _run_leg(cfg, leg, key_path, public_ip, known_hosts_path, run_data_dir)
+ finally:
+ # Teardown-in-finally: this block runs whether the try body above
+ # succeeded, raised (e.g. a failed run_profile.sh -> a propagating
+ # subprocess.CalledProcessError), or was interrupted -- so an
+ # (expensive, 8x H100) instance is never left running just because
+ # one profiling leg failed partway through. See
+ # _teardown_and_cleanup's docstring for how a *second* failure
+ # (teardown itself failing) is handled without masking whichever
+ # exception was already propagating out of the try body.
+ if args.keep_alive:
+ _print_keep_alive_notice(cfg, key_path, state)
+ else:
+ _teardown_and_cleanup(ec2_client, state, state_path, cfg)
+
+ return 0
+
+
+if __name__ == "__main__":
+ raise SystemExit(main())
diff --git a/notebooks/astrasim2_correlation/correlation/parse_nccl.py b/notebooks/astrasim2_correlation/correlation/parse_nccl.py
new file mode 100644
index 00000000..8a4c6d41
--- /dev/null
+++ b/notebooks/astrasim2_correlation/correlation/parse_nccl.py
@@ -0,0 +1,483 @@
+"""Parse raw NCCL collective-communication profiling logs into tidy CSVs.
+
+This module turns the stdout of two profiling tools into a single, unified
+CSV schema that the correlation notebook (``correlation.ipynb``) consumes:
+
+1. **nccl-tests** (upstream NVIDIA binaries: ``all_reduce_perf``,
+ ``all_gather_perf``, ``reduce_scatter_perf``, ``alltoall_perf``,
+ ``broadcast_perf``, ``sendrecv_perf``) run against the fully-connected
+ (FC) NVSwitch fabric.
+2. ``torus_bench`` (a custom binary built by a sibling work package) run
+ against a torus-topology emulation on the same physical fabric.
+
+Both tools are parsed independently (:func:`parse_nccl_tests` and
+:func:`parse_torus_bench`) and their results are reshaped into the shared
+schema documented at :data:`UNIFIED_CSV_FIELDNAMES` before being written to
+disk with :func:`rows_to_csv`. The module is also a CLI entry point (see
+:func:`main`) so it can be invoked directly from ``run_profile.sh`` on the
+profiling instance without any extra Python dependencies.
+
+Notes
+-----
+Stdlib-only by design: this script runs on a freshly provisioned EC2
+instance where installing a virtualenv is unwanted overhead. Only ``csv``,
+``argparse``, and ``pathlib`` (plus ``sys`` for the CLI entry point) are
+used, so it works under any plain ``python3`` >= 3.8.
+"""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import sys
+from pathlib import Path
+from typing import Union
+
+# Design: the unified schema is a module-level constant (rather than being
+# implicit in whatever keys happen to be in the first row dict) so that
+# rows_to_csv() always emits a stable, predictable column order regardless
+# of which source produced the rows, and so the notebook can rely on the
+# header never silently reordering itself as this module evolves.
+UNIFIED_CSV_FIELDNAMES: list[str] = [
+ "source",
+ "topology",
+ "dims",
+ "collective",
+ "size_bytes",
+ "count",
+ "dtype",
+ "time_us",
+ "algbw_GBps",
+ "busbw_GBps",
+ "wrong",
+]
+
+# Minimum number of whitespace-separated tokens a nccl-tests data row must
+# have before it is considered parseable. A row always carries at least
+# size, count, type, redop, root (5 leading columns) plus the 8 trailing
+# out-of-place/in-place metric columns = 13 tokens; 10 is used as a looser
+# lower bound per the spec so that unexpected/future nccl-tests column
+# layouts with slightly fewer leading columns are still accepted as long as
+# the trailing-8 structure holds.
+_MIN_NCCL_TESTS_TOKENS = 10
+
+# Number of trailing tokens on a nccl-tests data row that carry the
+# out-of-place/in-place timing results. This is the anchor of the parsing
+# strategy; see parse_nccl_tests() docstring for the rationale.
+_TRAILING_METRIC_TOKENS = 8
+
+# Sentinel line prefix emitted by torus_bench for each data point. Chosen
+# by the sibling work package specifically so it is trivial to grep/parse
+# out of interleaved '#'-prefixed human-readable log noise.
+_TORUS_SENTINEL_PREFIX = "TORUSBENCH,"
+
+
+def parse_nccl_tests(text: str) -> list[dict]:
+ """Parse the stdout of an nccl-tests collective benchmark binary.
+
+ nccl-tests binaries (``all_reduce_perf``, ``all_gather_perf``,
+ ``reduce_scatter_perf``, ``alltoall_perf``, ``broadcast_perf``,
+ ``sendrecv_perf``) share a common output shape: a block of ``#``-prefixed
+ header/comment lines, followed by one data row per message size, followed
+ by ``#``-prefixed summary lines. The *leading* columns of a data row vary
+ per collective (e.g. ``redop``/``root`` are meaningless for
+ ``alltoall_perf`` and print as ``none``/``-1``), but the *trailing* eight
+ columns are always, in order: out-of-place ``time``, ``algbw``, ``busbw``,
+ ``#wrong``, then in-place ``time``, ``algbw``, ``busbw``, ``#wrong``.
+
+ Parameters
+ ----------
+ text : str
+ Raw stdout captured from an nccl-tests binary invocation. May
+ contain blank lines and ``#``-prefixed comment/header/summary lines
+ interleaved with data rows.
+
+ Returns
+ -------
+ list of dict
+ One dict per parsed data row, in file order, with keys:
+ ``size_bytes`` (int), ``count`` (int), ``dtype`` (str),
+ ``time_us`` (float), ``algbw_GBps`` (float), ``busbw_GBps`` (float),
+ ``wrong`` (str; ``"0"``, another digit string, or ``"N/A"`` when
+ validation was disabled for the run). Only the *out-of-place*
+ metrics are kept, matching the spec's trailing-token convention;
+ the in-place metrics are intentionally discarded since the
+ correlation study only needs one consistent number per size.
+
+ Notes
+ -----
+ Design: rather than hand-writing a distinct column layout per collective
+ (which would need to track every nccl-tests release), this uses a single
+ robust rule anchored on the *trailing* 8 tokens, which nccl-tests has
+ kept stable across collectives and versions even as leading columns
+ (redop, root) have been added/repurposed. A line is treated as a data
+ row only if tokens[0] and tokens[1] both parse as int -- this
+ distinguishes real data rows (which always start with two integers:
+ size in bytes, element count) from stray non-'#' lines (blank-ish
+ whitespace, malformed output, or future header formats) without needing
+ to hard-code the '#' comment convention as the *only* skip signal.
+
+ This function does not raise on malformed input; unparseable lines are
+ silently skipped so that a partially-corrupt log (e.g. truncated by a
+ crashed run) still yields whatever valid rows it contains.
+
+ Examples
+ --------
+ >>> text = (
+ ... "# nThread 1 nGpus 8\\n"
+ ... " 1048576 262144 float sum -1 "
+ ... "98.52 10.64 18.62 0 97.11 10.80 18.90 0\\n"
+ ... )
+ >>> rows = parse_nccl_tests(text)
+ >>> rows[0]["size_bytes"], rows[0]["time_us"], rows[0]["wrong"]
+ (1048576, 98.52, '0')
+ """
+ rows: list[dict] = []
+ for line in text.splitlines():
+ stripped = line.strip()
+ # Skip blank lines and '#'-prefixed header/comment/summary lines.
+ if not stripped or stripped.startswith("#"):
+ continue
+
+ tokens = stripped.split()
+ if len(tokens) < _MIN_NCCL_TESTS_TOKENS:
+ continue
+
+ try:
+ size_bytes = int(tokens[0])
+ count = int(tokens[1])
+ except ValueError:
+ # Not a data row (e.g. stray non-'#' text); skip rather than
+ # raise so one bad line doesn't abort parsing of an otherwise
+ # good log.
+ continue
+
+ dtype = tokens[2]
+ try:
+ time_us = float(tokens[-_TRAILING_METRIC_TOKENS])
+ algbw_gbps = float(tokens[-_TRAILING_METRIC_TOKENS + 1])
+ busbw_gbps = float(tokens[-_TRAILING_METRIC_TOKENS + 2])
+ except ValueError:
+ # The trailing columns didn't parse as floats -- not a real
+ # data row (defensive; shouldn't happen given the int checks
+ # above already filtered most non-data lines).
+ continue
+ wrong = tokens[-_TRAILING_METRIC_TOKENS + 3]
+
+ rows.append(
+ {
+ "size_bytes": size_bytes,
+ "count": count,
+ "dtype": dtype,
+ "time_us": time_us,
+ "algbw_GBps": algbw_gbps,
+ "busbw_GBps": busbw_gbps,
+ "wrong": wrong,
+ }
+ )
+ return rows
+
+
+def parse_torus_bench(text: str) -> list[dict]:
+ """Parse the stdout of the custom ``torus_bench`` binary.
+
+ ``torus_bench`` prints ``#``-prefixed human-readable header lines plus
+ machine-parseable sentinel lines of the exact form::
+
+ TORUSBENCH,,,,,
+
+ Parameters
+ ----------
+ text : str
+ Raw stdout captured from a ``torus_bench`` invocation.
+
+ Returns
+ -------
+ list of dict
+ One dict per ``TORUSBENCH,`` sentinel line, in file order, with
+ keys: ``collective`` (str), ``dims`` (str, e.g. ``"2x2x2"``),
+ ``size_bytes`` (int), ``time_us`` (float), ``wrong`` (str; one of
+ ``"0"`` (check passed), ``"1"`` (check failed), or ``"N/A"``
+ (validation was not run for this data point)).
+
+ Notes
+ -----
+ Design: the sentinel line is comma-delimited (unlike nccl-tests'
+ whitespace-delimited columns) specifically so the sibling work package
+ could emit it without worrying about column-alignment padding; this
+ parser simply looks for the fixed ``TORUSBENCH,`` prefix and splits on
+ commas, ignoring every other line (including the human-readable ``#``
+ header). This makes the parser forward-compatible with additional
+ ``#``-prefixed diagnostic lines torus_bench might add later.
+
+ The ``check`` field's three-way encoding (``1``/``0``/``-``) is
+ remapped onto the same ``wrong`` vocabulary nccl-tests uses (a per-row
+ "wrongness" indicator string) so downstream consumers (rows_to_csv,
+ the notebook) can treat the ``wrong`` column uniformly across sources:
+ ``check == "1"`` (validation ran and passed) maps to ``wrong = "0"``
+ (zero wrong elements); ``check == "-"`` (validation was skipped for
+ this run) maps to ``wrong = "N/A"``, mirroring nccl-tests' own "N/A"
+ convention for validation-disabled runs; anything else (i.e.
+ ``check == "0"``, validation ran and failed) maps to ``wrong = "1"``.
+
+ Malformed sentinel lines (wrong field count, non-numeric size/time) are
+ silently skipped rather than raising, for the same reasons as
+ :func:`parse_nccl_tests`.
+ """
+ rows: list[dict] = []
+ for line in text.splitlines():
+ stripped = line.strip()
+ if not stripped.startswith(_TORUS_SENTINEL_PREFIX):
+ continue
+
+ fields = stripped.split(",")
+ # TORUSBENCH,,,,, = 6 fields.
+ if len(fields) != 6:
+ continue
+
+ _, collective, dims, size_bytes_str, time_us_str, check = fields
+ try:
+ size_bytes = int(size_bytes_str)
+ time_us = float(time_us_str)
+ except ValueError:
+ continue
+
+ if check == "1":
+ wrong = "0"
+ elif check == "-":
+ wrong = "N/A"
+ else:
+ wrong = "1"
+
+ rows.append(
+ {
+ "collective": collective,
+ "dims": dims,
+ "size_bytes": size_bytes,
+ "time_us": time_us,
+ "wrong": wrong,
+ }
+ )
+ return rows
+
+
+def rows_to_csv(rows: list[dict], out_path: Union[str, Path]) -> None:
+ """Write unified-schema rows to a CSV file.
+
+ Parameters
+ ----------
+ rows : list of dict
+ Rows already reshaped into the unified schema (see
+ :data:`UNIFIED_CSV_FIELDNAMES` for the exact column set and order).
+ Each dict must contain every key in ``UNIFIED_CSV_FIELDNAMES``;
+ missing keys are written as empty cells by :class:`csv.DictWriter`
+ default behavior is NOT relied upon here -- callers are expected to
+ supply complete rows (see :func:`main` for how CLI callers build
+ them). Extra keys beyond the unified schema are rejected by
+ :class:`csv.DictWriter` (``extrasaction="raise"``, the default) so
+ schema drift is caught early rather than silently dropped.
+ out_path : str or pathlib.Path
+ Destination file path. Parent directories are NOT created by this
+ function; callers must ensure the directory exists.
+
+ Returns
+ -------
+ None
+
+ Raises
+ ------
+ ValueError
+ If a row dict contains a key not present in
+ :data:`UNIFIED_CSV_FIELDNAMES` (raised by the underlying
+ :class:`csv.DictWriter`).
+ OSError
+ If ``out_path`` cannot be opened for writing (e.g. parent directory
+ does not exist, permission denied).
+
+ Notes
+ -----
+ Opens the file with ``newline=""`` as recommended by the :mod:`csv`
+ module docs, so that the csv module's own line-ending handling is used
+ verbatim rather than being double-translated by Python's text-mode
+ newline translation.
+ """
+ out_path = Path(out_path)
+ with out_path.open("w", newline="", encoding="utf-8") as f:
+ writer = csv.DictWriter(f, fieldnames=UNIFIED_CSV_FIELDNAMES)
+ writer.writeheader()
+ writer.writerows(rows)
+
+
+def _build_arg_parser() -> argparse.ArgumentParser:
+ """Construct the CLI argument parser for this module.
+
+ Returns
+ -------
+ argparse.ArgumentParser
+ Parser accepting the raw log path plus labeling/validation flags
+ described in the module CLI usage (see :func:`main`).
+ """
+ parser = argparse.ArgumentParser(
+ prog="parse_nccl.py",
+ description=(
+ "Parse a raw nccl-tests or torus_bench profiling log into the "
+ "unified CSV schema consumed by the correlation notebook."
+ ),
+ )
+ parser.add_argument(
+ "raw_log",
+ type=Path,
+ help="Path to the raw stdout log captured from the profiling binary.",
+ )
+ parser.add_argument(
+ "--source",
+ required=True,
+ choices=["nccl-tests", "torus_bench"],
+ help="Which tool produced raw_log.",
+ )
+ parser.add_argument(
+ "--collective",
+ required=True,
+ help=(
+ "Collective name. For --source nccl-tests this labels every "
+ "output row directly (the tool's own stdout does not name the "
+ "collective). For --source torus_bench this is instead "
+ "cross-checked against the collective embedded in each "
+ "TORUSBENCH sentinel line; a mismatch is an error."
+ ),
+ )
+ parser.add_argument(
+ "--topology",
+ required=True,
+ choices=["fc", "torus"],
+ help="Fabric topology label to stamp onto every output row.",
+ )
+ parser.add_argument(
+ "--dims",
+ default=None,
+ help=(
+ "Dimension string (e.g. '2x2x2'). For --source nccl-tests this "
+ "overrides the default dims label of '8' (the fixed GPU count "
+ "of a single NVSwitch-connected node). For --source "
+ "torus_bench, if given, it is cross-checked against the dims "
+ "embedded in each TORUSBENCH sentinel line; a mismatch is an "
+ "error. If omitted for torus_bench, the sentinel's own dims "
+ "value is used unchecked."
+ ),
+ )
+ parser.add_argument(
+ "--out",
+ required=True,
+ type=Path,
+ help="Destination CSV path.",
+ )
+ return parser
+
+
+# Default dims label for nccl-tests rows when --dims is not supplied on the
+# CLI. FC-leg runs are always against a single 8x-H100 NVSwitch node, so "8"
+# (the GPU count) is the natural default; --dims exists mainly to let the
+# CLI stay uniform with the torus_bench invocation and to support future
+# multi-node FC runs without changing this module.
+_DEFAULT_NCCL_TESTS_DIMS = "8"
+
+
+def main(argv: Union[list, None] = None) -> None:
+ """CLI entry point: parse a raw log and write the unified CSV.
+
+ Parameters
+ ----------
+ argv : list of str, optional
+ Argument vector to parse in place of ``sys.argv[1:]``. Primarily
+ useful for testing; production invocations (from ``run_profile.sh``)
+ pass ``None`` and rely on ``sys.argv``.
+
+ Returns
+ -------
+ None
+
+ Raises
+ ------
+ SystemExit
+ Raised by :mod:`argparse` on invalid/missing arguments (exit code
+ 2), or explicitly via ``parser.error()`` when a ``--collective``/
+ ``--dims`` value supplied on the CLI disagrees with the value
+ embedded in a torus_bench sentinel line (exit code 2). Also raised
+ implicitly if ``raw_log`` cannot be read (propagates as an
+ unhandled :class:`OSError`, not caught here -- a missing/unreadable
+ input log is a hard setup error the caller (``run_profile.sh``)
+ should see immediately rather than have masked).
+ """
+ parser = _build_arg_parser()
+ args = parser.parse_args(argv)
+
+ text = args.raw_log.read_text(encoding="utf-8")
+
+ if args.source == "torus_bench":
+ parsed = parse_torus_bench(text)
+ # Design: validate CLI-supplied --collective/--dims against what
+ # the sentinel lines actually say, rather than trusting the CLI
+ # blindly. This catches operator error in run_profile.sh (e.g. a
+ # copy-paste mistake wiring the wrong collective's log into the
+ # wrong parse invocation) at parse time instead of silently
+ # mislabeling data that later gets combined into the notebook.
+ for parsed_row in parsed:
+ if parsed_row["collective"] != args.collective:
+ parser.error(
+ f"--collective {args.collective!r} does not match "
+ f"collective {parsed_row['collective']!r} found in "
+ f"{args.raw_log}"
+ )
+ if args.dims is not None and parsed_row["dims"] != args.dims:
+ parser.error(
+ f"--dims {args.dims!r} does not match dims "
+ f"{parsed_row['dims']!r} found in {args.raw_log}"
+ )
+ rows = [
+ {
+ "source": "torus_bench",
+ "topology": args.topology,
+ "dims": parsed_row["dims"],
+ "collective": parsed_row["collective"],
+ "size_bytes": parsed_row["size_bytes"],
+ # Design: count/dtype/algbw/busbw are left empty for torus
+ # rows per spec -- bandwidth conventions for the torus
+ # topology (e.g. what counts as "algorithm bandwidth" when
+ # hops differ per link) are derived in the notebook from
+ # size_bytes/time_us/dims, not computed here, to keep this
+ # parser topology-agnostic.
+ "count": "",
+ "dtype": "",
+ "time_us": parsed_row["time_us"],
+ "algbw_GBps": "",
+ "busbw_GBps": "",
+ "wrong": parsed_row["wrong"],
+ }
+ for parsed_row in parsed
+ ]
+ else:
+ parsed = parse_nccl_tests(text)
+ dims = args.dims if args.dims is not None else _DEFAULT_NCCL_TESTS_DIMS
+ rows = [
+ {
+ "source": "nccl-tests",
+ "topology": args.topology,
+ "dims": dims,
+ "collective": args.collective,
+ "size_bytes": parsed_row["size_bytes"],
+ "count": parsed_row["count"],
+ "dtype": parsed_row["dtype"],
+ "time_us": parsed_row["time_us"],
+ "algbw_GBps": parsed_row["algbw_GBps"],
+ "busbw_GBps": parsed_row["busbw_GBps"],
+ "wrong": parsed_row["wrong"],
+ }
+ for parsed_row in parsed
+ ]
+
+ args.out.parent.mkdir(parents=True, exist_ok=True)
+ rows_to_csv(rows, args.out)
+
+
+if __name__ == "__main__":
+ main(sys.argv[1:])
diff --git a/notebooks/astrasim2_correlation/correlation/provision.py b/notebooks/astrasim2_correlation/correlation/provision.py
new file mode 100644
index 00000000..c6ca9fa4
--- /dev/null
+++ b/notebooks/astrasim2_correlation/correlation/provision.py
@@ -0,0 +1,895 @@
+"""Provision one p5.48xlarge (8x H100, NVSwitch) EC2 instance for NCCL profiling.
+
+This script is the "up" half of the correlation study's empirical leg: it
+launches exactly one spot-first, on-demand-fallback instance, waits for it
+to be SSH-reachable, and records everything needed to find/tear it down
+again in a small JSON state file. See ``teardown.py`` for the "down" half
+and ``README.md`` for the full runbook.
+
+Every function below is written to be independently importable and
+testable: ``orchestrate.py`` (a sibling work package, written separately)
+imports these functions directly rather than shelling out to this file, so
+their signatures are part of this module's public contract and must not
+change without updating that caller too.
+
+Design: boto3 is not a repo dependency
+----------------------------------------
+accelforge's ``pyproject.toml`` does not (and per this work package's scope
+must not) depend on ``boto3`` -- only this AWS-provisioning corner of one
+notebook's correlation study needs it. The import below is therefore
+guarded: importing this module never fails just because boto3 is missing,
+so ``python provision.py --help`` keeps working in any environment. Actual
+AWS calls fail fast with a clear "pip install boto3" message via
+:func:`_require_boto3`, called once at the top of :func:`main` before any
+client is constructed.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import socket
+import sys
+import time
+import urllib.error
+import urllib.request
+from pathlib import Path
+from typing import Optional
+
+from config import Config
+
+try:
+ import boto3
+ from botocore.exceptions import ClientError
+except ImportError: # pragma: no cover - exercised only when boto3 truly absent
+ boto3 = None
+ # Design: fall back to plain Exception as a placeholder so that
+ # `except ClientError:` clauses elsewhere in this module remain valid
+ # Python (no NameError at import time) even when boto3 is missing.
+ # Those clauses are only ever reached after _require_boto3() has
+ # already raised, so this placeholder is never actually matched in
+ # practice -- it exists purely to keep module import side-effect-free.
+ class ClientError(Exception): # type: ignore[no-redef]
+ pass
+
+
+def _require_boto3() -> None:
+ """Raise a clear, actionable error if boto3 is not installed.
+
+ Raises
+ ------
+ SystemExit
+ Always, if ``boto3`` failed to import. The message tells the user
+ exactly how to fix it rather than surfacing a bare
+ ``ModuleNotFoundError`` traceback.
+ """
+ if boto3 is None:
+ raise SystemExit(
+ "boto3 is required for AWS provisioning but is not installed in "
+ "this Python environment.\n"
+ "Install it with: pip install boto3"
+ )
+
+
+# ClientError codes under which "spot-then-ondemand" purchasing retries on
+# demand instead of failing the whole run. All represent spot-market
+# scarcity/limits rather than a request-shape problem, so retrying the same
+# request as on-demand is expected to succeed. p5.48xlarge (8x H100) is a
+# scarce, high-demand instance type, so hitting these in practice is not
+# unusual and should not be treated as fatal when the operator has opted
+# into a fallback.
+_SPOT_FALLBACK_ERROR_CODES = frozenset(
+ {
+ "InsufficientInstanceCapacity",
+ "SpotMaxPriceTooLow",
+ "MaxSpotInstanceCountExceeded",
+ "Unsupported",
+ "InstanceLimitExceeded",
+ }
+)
+
+_SSH_POLL_INTERVAL_S = 5.0
+_SSH_POLL_TIMEOUT_S = 5 * 60.0
+
+_COST_WARNING = (
+ "COST WARNING: p5.48xlarge on-demand pricing is roughly $30-55/hr "
+ "depending on region and current AWS pricing. VERIFY CURRENT PRICING "
+ "before proceeding: https://aws.amazon.com/ec2/pricing/on-demand/"
+)
+
+
+def resolve_ami(ssm_client, parameter: str) -> str:
+ """Resolve an AMI id from a public SSM parameter.
+
+ Parameters
+ ----------
+ ssm_client : botocore.client.BaseClient
+ A boto3 ``ssm`` client (or a stub thereof).
+ parameter : str
+ Fully-qualified SSM parameter name, e.g.
+ ``"/aws/service/deeplearning/ami/x86_64/.../latest/ami-id"``.
+
+ Returns
+ -------
+ str
+ The AMI id stored at ``parameter``.
+
+ Raises
+ ------
+ botocore.exceptions.ClientError
+ If ``parameter`` does not exist or the caller lacks
+ ``ssm:GetParameter`` permission; propagated unmodified so callers
+ see AWS's own error code and message.
+ """
+ response = ssm_client.get_parameter(Name=parameter)
+ return response["Parameter"]["Value"]
+
+
+def caller_ip() -> str:
+ """Discover the caller's public IPv4 address via checkip.amazonaws.com.
+
+ Used by :func:`main` to scope the provisioned security group's SSH
+ ingress rule to just this machine, when the operator has not supplied
+ an explicit ``--ssh-cidr``.
+
+ Returns
+ -------
+ str
+ The caller's public IP as a dotted-quad string.
+
+ Raises
+ ------
+ RuntimeError
+ If the IP could not be discovered for any reason (network error,
+ timeout, or an empty response body). See Notes for why this fails
+ loudly rather than falling back to any sentinel value.
+
+ Notes
+ -----
+ Design: fails *fast* (raises) rather than falling back to a sentinel.
+ An earlier version of this function failed *open* to the sentinel
+ ``"0.0.0.0"`` on discovery failure, reasoning that a transient DNS blip
+ or checkip.amazonaws.com outage shouldn't hard-fail an
+ otherwise-working run. That reasoning had a bug: every caller turns
+ this return value into a CIDR via ``f"{ip}/32"``, so the sentinel
+ actually produced ``"0.0.0.0/32"`` -- a CIDR matching no address at
+ all -- which locks *everyone*, including the operator, out over SSH,
+ the opposite of what the old warning text claimed ("falling back to
+ 0.0.0.0/0", i.e. open to the world). Worse, that silent misconfiguration
+ was only discoverable *after* a real key pair, security group, and
+ instance had already been created and billing had already started.
+ Raising here instead is strictly better on both axes this function
+ cares about -- SAFE (no accidental everyone-blocked security group) and
+ SECURE (no accidental world-open one either) -- and it fires before any
+ AWS resource exists or any money is spent: both call sites
+ (``provision.main`` and ``orchestrate.main``) invoke this function only
+ when ``--ssh-cidr`` was not supplied, and always before
+ ``ensure_key_pair``/``ensure_security_group``/``launch_instance``. The
+ error message tells the operator exactly how to proceed: re-run with
+ ``--ssh-cidr /32``.
+ """
+ try:
+ with urllib.request.urlopen("https://checkip.amazonaws.com", timeout=10) as resp:
+ ip = resp.read().decode("utf-8").strip()
+ if not ip:
+ raise ValueError("empty response body from checkip.amazonaws.com")
+ return ip
+ except (urllib.error.URLError, ValueError, OSError) as exc:
+ raise RuntimeError(
+ "Could not determine your public IP via checkip.amazonaws.com "
+ f"({exc!r}). No AWS resources have been created yet, so there is "
+ "nothing to clean up -- re-run with --ssh-cidr /32 to "
+ "supply your CIDR explicitly instead of relying on auto-detection."
+ ) from exc
+
+
+def ensure_key_pair(ec2_client, key_name: str, key_dir: Path) -> Path:
+ """Create (or reuse) an EC2 key pair and its local PEM file.
+
+ Parameters
+ ----------
+ ec2_client : botocore.client.BaseClient
+ A boto3 ``ec2`` client (or a stub thereof).
+ key_name : str
+ Name to give the key pair in AWS.
+ key_dir : pathlib.Path
+ Local directory to write ``.pem`` into. Created if it
+ does not already exist.
+
+ Returns
+ -------
+ pathlib.Path
+ Path to the local PEM file (either freshly written, or the
+ existing one being reused).
+
+ Raises
+ ------
+ RuntimeError
+ If AWS reports the key pair already exists (``ClientError`` code
+ ``InvalidKeyPair.Duplicate``) but no local PEM file is present.
+ AWS never returns private key material for a pre-existing key
+ pair, so there is no way to recover the PEM in this situation --
+ the caller must delete the AWS-side key pair or choose a
+ different ``run_id``.
+ botocore.exceptions.ClientError
+ For any other ``create_key_pair`` failure, propagated unmodified.
+
+ Notes
+ -----
+ Side effect: writes a file to disk at ``/.pem`` with
+ ``0o600`` permissions (owner read/write only, matching what ``ssh``
+ requires of private key files).
+ """
+ key_dir.mkdir(parents=True, exist_ok=True)
+ key_path = key_dir / f"{key_name}.pem"
+
+ try:
+ response = ec2_client.create_key_pair(
+ KeyName=key_name, KeyType="rsa", KeyFormat="pem"
+ )
+ except ClientError as exc:
+ error_code = exc.response.get("Error", {}).get("Code", "")
+ if error_code == "InvalidKeyPair.Duplicate":
+ if key_path.exists():
+ print(
+ f"Key pair {key_name!r} already exists in AWS and a local "
+ f"PEM was found at {key_path}; reusing it."
+ )
+ return key_path
+ raise RuntimeError(
+ f"Key pair {key_name!r} already exists in AWS, but no local "
+ f"PEM file was found at {key_path}. AWS never returns private "
+ "key material for a pre-existing key pair, so it cannot be "
+ "recovered. Either delete the AWS-side key pair "
+ f"(aws ec2 delete-key-pair --key-name {key_name}) and re-run, "
+ "or pass a different --run-id so a fresh key pair name is used."
+ ) from exc
+ raise
+
+ key_material = response["KeyMaterial"]
+ # Design: use os.open with the 0o600 mode baked into file creation
+ # (rather than write-then-chmod) so the PEM is never briefly readable
+ # at default (often world-readable) permissions between those two
+ # steps.
+ fd = os.open(key_path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
+ with os.fdopen(fd, "w") as fh:
+ fh.write(key_material)
+ return key_path
+
+
+def ensure_security_group(
+ ec2_client, group_name: str, ssh_cidr: str, tag_project: str, run_id: str
+) -> str:
+ """Create a security group in the default VPC allowing SSH from one CIDR.
+
+ Parameters
+ ----------
+ ec2_client : botocore.client.BaseClient
+ A boto3 ``ec2`` client (or a stub thereof).
+ group_name : str
+ Name to give the new security group.
+ ssh_cidr : str
+ CIDR block (e.g. ``"1.2.3.4/32"`` or ``"0.0.0.0/0"``) to allow
+ inbound TCP/22 from.
+ tag_project : str
+ Value for the ``Project`` tag on the new group.
+ run_id : str
+ Value for the ``RunId`` tag on the new group, and included in its
+ description.
+
+ Returns
+ -------
+ str
+ The new security group's id.
+
+ Raises
+ ------
+ RuntimeError
+ If the region has no default VPC (``describe_vpcs`` returns no
+ results for ``isDefault=true``). NCCL profiling on a single node
+ has no cross-VPC requirements, so this script deliberately does
+ not attempt to create or select a non-default VPC -- that is out
+ of scope for a short-lived profiling instance.
+ botocore.exceptions.ClientError
+ For any ``create_security_group``/``authorize_security_group_ingress``/
+ ``create_tags`` failure, propagated unmodified.
+ """
+ vpcs = ec2_client.describe_vpcs(Filters=[{"Name": "isDefault", "Values": ["true"]}])
+ vpc_list = vpcs.get("Vpcs", [])
+ if not vpc_list:
+ raise RuntimeError(
+ "No default VPC found in this region. Create one "
+ "(aws ec2 create-default-vpc) or provision a VPC manually, then "
+ "re-run."
+ )
+ vpc_id = vpc_list[0]["VpcId"]
+
+ create_resp = ec2_client.create_security_group(
+ GroupName=group_name,
+ Description=f"accelforge correlation study SG for run {run_id}",
+ VpcId=vpc_id,
+ )
+ sg_id = create_resp["GroupId"]
+
+ ec2_client.authorize_security_group_ingress(
+ GroupId=sg_id,
+ IpPermissions=[
+ {
+ "IpProtocol": "tcp",
+ "FromPort": 22,
+ "ToPort": 22,
+ "IpRanges": [
+ {
+ "CidrIp": ssh_cidr,
+ "Description": "SSH access for accelforge correlation study",
+ }
+ ],
+ }
+ ],
+ )
+
+ ec2_client.create_tags(
+ Resources=[sg_id],
+ Tags=[
+ {"Key": "Project", "Value": tag_project},
+ {"Key": "RunId", "Value": run_id},
+ {"Key": "Name", "Value": group_name},
+ ],
+ )
+ return sg_id
+
+
+def _build_run_instances_kwargs(
+ cfg: Config, ami_id: str, sg_id: str, key_name: str, dry_run: bool, use_spot: bool
+) -> dict:
+ """Build the ``run_instances`` kwargs shared by the spot and on-demand paths.
+
+ Parameters
+ ----------
+ cfg : Config
+ Run configuration.
+ ami_id : str
+ AMI id resolved by :func:`resolve_ami`.
+ sg_id : str
+ Security group id from :func:`ensure_security_group`.
+ key_name : str
+ Key pair name from :func:`ensure_key_pair`.
+ dry_run : bool
+ Whether to set ``DryRun=True`` on the request.
+ use_spot : bool
+ Whether to request a spot instance (adds ``InstanceMarketOptions``)
+ or an on-demand one.
+
+ Returns
+ -------
+ dict
+ Keyword arguments ready to pass to ``ec2_client.run_instances(**kwargs)``.
+
+ Notes
+ -----
+ Factored out of :func:`launch_instance` so the spot attempt and the
+ on-demand fallback attempt build their request the same way apart from
+ the one ``InstanceMarketOptions`` difference -- avoids the two paths
+ silently drifting apart (e.g. one attempt forgetting a tag) as this
+ function evolves.
+ """
+ kwargs = {
+ "ImageId": ami_id,
+ "InstanceType": cfg.instance_type,
+ "KeyName": key_name,
+ "SecurityGroupIds": [sg_id],
+ "MinCount": 1,
+ "MaxCount": 1,
+ # "terminate" (not "stop") so an in-instance `shutdown` -- e.g. the
+ # dead-man timer armed by setup scripts -- fully releases the
+ # instance rather than leaving it (and its EBS billing) stopped
+ # but still provisioned.
+ "InstanceInitiatedShutdownBehavior": "terminate",
+ "BlockDeviceMappings": [
+ {
+ "DeviceName": "/dev/sda1",
+ "Ebs": {
+ "VolumeSize": cfg.root_volume_gb,
+ "VolumeType": "gp3",
+ "DeleteOnTermination": True,
+ },
+ }
+ ],
+ "TagSpecifications": [
+ {
+ "ResourceType": "instance",
+ "Tags": [
+ {"Key": "Project", "Value": cfg.tag_project},
+ {"Key": "RunId", "Value": cfg.run_id},
+ {"Key": "Name", "Value": f"{cfg.tag_project}-{cfg.run_id}"},
+ ],
+ },
+ {
+ "ResourceType": "volume",
+ "Tags": [
+ {"Key": "Project", "Value": cfg.tag_project},
+ {"Key": "RunId", "Value": cfg.run_id},
+ {"Key": "Name", "Value": f"{cfg.tag_project}-{cfg.run_id}"},
+ ],
+ },
+ ],
+ "DryRun": dry_run,
+ }
+ if cfg.availability_zone:
+ kwargs["Placement"] = {"AvailabilityZone": cfg.availability_zone}
+ if use_spot:
+ kwargs["InstanceMarketOptions"] = {
+ "MarketType": "spot",
+ "SpotOptions": {
+ "SpotInstanceType": "one-time",
+ "InstanceInterruptionBehavior": "terminate",
+ },
+ }
+ return kwargs
+
+
+def launch_instance(
+ ec2_client,
+ cfg: Config,
+ ami_id: str,
+ sg_id: str,
+ key_name: str,
+ dry_run: bool = False,
+) -> dict:
+ """Launch exactly one instance, honoring ``cfg.purchasing``.
+
+ Parameters
+ ----------
+ ec2_client : botocore.client.BaseClient
+ A boto3 ``ec2`` client (or a stub thereof).
+ cfg : Config
+ Run configuration; ``cfg.purchasing`` selects the strategy below.
+ ami_id : str
+ AMI id resolved by :func:`resolve_ami`.
+ sg_id : str
+ Security group id from :func:`ensure_security_group`.
+ key_name : str
+ Key pair name from :func:`ensure_key_pair`.
+ dry_run : bool, default False
+ If ``True``, sets ``DryRun=True`` on every ``run_instances`` call.
+ AWS answers a dry run with an error either way: ``DryRunOperation``
+ means the call would have succeeded, ``UnauthorizedOperation``
+ means the caller lacks permission. This function treats those two
+ codes accordingly rather than as generic failures.
+
+ Returns
+ -------
+ dict
+ ``{"instance_id": str or None, "purchasing_used": "spot" or "ondemand"}``.
+ ``instance_id`` is ``None`` when ``dry_run=True`` and the
+ authorization check succeeded, since no instance was actually
+ created in that case.
+
+ Raises
+ ------
+ RuntimeError
+ If a dry run reports ``UnauthorizedOperation`` (the configured
+ credentials cannot launch this instance type/configuration).
+ botocore.exceptions.ClientError
+ - If ``cfg.purchasing == "spot"`` and the spot request fails for
+ any reason (no fallback is attempted in this mode).
+ - If ``cfg.purchasing == "spot-then-ondemand"`` and the spot
+ request fails with a code *not* in
+ :data:`_SPOT_FALLBACK_ERROR_CODES` (that set is deliberately
+ narrow -- e.g. a malformed request should fail loudly rather
+ than silently retrying as on-demand and masking the bug).
+ - If the (possibly-fallback) on-demand request itself fails.
+
+ Notes
+ -----
+ Purchasing strategies:
+
+ - ``"ondemand"``: on-demand only, no spot attempt.
+ - ``"spot"``: spot only; any failure propagates without a fallback.
+ - ``"spot-then-ondemand"`` (the default): attempts spot first. If that
+ attempt fails with one of :data:`_SPOT_FALLBACK_ERROR_CODES` --
+ capacity/limit/market conditions rather than a malformed request --
+ it prints the failure and retries once as on-demand. Any other
+ ``ClientError`` code (e.g. a parameter validation error) propagates
+ immediately without a fallback attempt, since retrying on-demand
+ would not fix a malformed request and would only obscure the real
+ error.
+ """
+
+ def _run(use_spot: bool) -> dict:
+ kwargs = _build_run_instances_kwargs(cfg, ami_id, sg_id, key_name, dry_run, use_spot)
+ try:
+ response = ec2_client.run_instances(**kwargs)
+ except ClientError as exc:
+ code = exc.response.get("Error", {}).get("Code", "")
+ if code == "DryRunOperation":
+ # AWS's DryRun contract: this specific error code means
+ # "you WOULD have been authorized to make this call" -- it
+ # is deliberately raised as an error even on the success
+ # path, so seeing it here IS the successful outcome of a
+ # dry run, not a failure.
+ print("dry-run OK: authorized")
+ return {
+ "instance_id": None,
+ "purchasing_used": "spot" if use_spot else "ondemand",
+ }
+ if code == "UnauthorizedOperation":
+ raise RuntimeError(
+ "AWS denied the run_instances permission check "
+ "(UnauthorizedOperation). The configured credentials lack "
+ f"ec2:RunInstances (or a related) permission for "
+ f"{cfg.instance_type}."
+ ) from exc
+ raise
+ instance_id = response["Instances"][0]["InstanceId"]
+ return {
+ "instance_id": instance_id,
+ "purchasing_used": "spot" if use_spot else "ondemand",
+ }
+
+ if cfg.purchasing == "ondemand":
+ return _run(use_spot=False)
+
+ if cfg.purchasing == "spot":
+ return _run(use_spot=True)
+
+ # cfg.purchasing == "spot-then-ondemand" (validated by Config.__post_init__
+ # to be one of exactly these three values).
+ try:
+ return _run(use_spot=True)
+ except ClientError as exc:
+ code = exc.response.get("Error", {}).get("Code", "")
+ if code in _SPOT_FALLBACK_ERROR_CODES:
+ print(f"Spot request failed ({code}); falling back to on-demand.")
+ return _run(use_spot=False)
+ raise
+
+
+def wait_for_instance(ec2_client, instance_id: str) -> str:
+ """Block until an instance is running and accepting TCP connections on port 22.
+
+ Parameters
+ ----------
+ ec2_client : botocore.client.BaseClient
+ A boto3 ``ec2`` client (or a stub thereof).
+ instance_id : str
+ Id of the instance to wait for.
+
+ Returns
+ -------
+ str
+ The instance's public IPv4 address.
+
+ Raises
+ ------
+ RuntimeError
+ If ``describe_instances`` returns no matching instance, or the
+ instance has no public IP address (e.g. it landed in a subnet that
+ does not auto-assign one).
+ TimeoutError
+ If port 22 does not become reachable within
+ :data:`_SSH_POLL_TIMEOUT_S` seconds of the instance reaching the
+ ``running`` state.
+ botocore.exceptions.WaiterError
+ If the ``instance_running`` waiter itself times out or the
+ instance transitions to a terminal failure state.
+
+ Notes
+ -----
+ Two-stage wait, because "EC2 says running" and "sshd is accepting
+ connections" are different events with a real gap between them (boot,
+ cloud-init, driver/NCCL setup on the deep learning AMI): first the
+ ``instance_running`` waiter (AWS-side state), then a plain TCP connect
+ poll against port 22 (this script's own liveness check), every
+ :data:`_SSH_POLL_INTERVAL_S` seconds for up to
+ :data:`_SSH_POLL_TIMEOUT_S`.
+ """
+ waiter = ec2_client.get_waiter("instance_running")
+ waiter.wait(InstanceIds=[instance_id])
+
+ describe = ec2_client.describe_instances(InstanceIds=[instance_id])
+ reservations = describe.get("Reservations", [])
+ if not reservations or not reservations[0].get("Instances"):
+ raise RuntimeError(
+ f"describe_instances returned no data for instance {instance_id!r}"
+ )
+ instance = reservations[0]["Instances"][0]
+ public_ip = instance.get("PublicIpAddress")
+ if not public_ip:
+ raise RuntimeError(
+ f"Instance {instance_id} is running but has no public IP address. "
+ "Check that its subnet auto-assigns public IPs."
+ )
+
+ _wait_for_ssh_port(public_ip)
+ return public_ip
+
+
+def _wait_for_ssh_port(
+ host: str,
+ port: int = 22,
+ interval_s: float = _SSH_POLL_INTERVAL_S,
+ timeout_s: float = _SSH_POLL_TIMEOUT_S,
+) -> None:
+ """Poll a TCP port until it accepts a connection or a timeout elapses.
+
+ Parameters
+ ----------
+ host : str
+ Hostname or IP address to connect to.
+ port : int, default 22
+ TCP port to poll.
+ interval_s : float, default 5.0
+ Seconds to sleep between connection attempts.
+ timeout_s : float, default 300.0
+ Total seconds to poll before giving up.
+
+ Raises
+ ------
+ TimeoutError
+ If no connection succeeds within ``timeout_s`` seconds.
+
+ Notes
+ -----
+ Uses ``socket.create_connection`` (a plain TCP connect/close) rather
+ than an actual SSH handshake -- sufficient to confirm sshd is up
+ without adding a paramiko/fabric dependency for a single boolean
+ liveness check.
+ """
+ deadline = time.monotonic() + timeout_s
+ while time.monotonic() < deadline:
+ try:
+ with socket.create_connection((host, port), timeout=interval_s):
+ return
+ except OSError:
+ pass
+ time.sleep(interval_s)
+ raise TimeoutError(
+ f"Timed out after {timeout_s}s waiting for {host}:{port} to accept "
+ "TCP connections (SSH not yet reachable)."
+ )
+
+
+def write_state(state_dir, run_id: str, state: dict) -> Path:
+ """Write a run's provisioning state to ``/.json``.
+
+ Parameters
+ ----------
+ state_dir : str or pathlib.Path
+ Directory to write the state file into. Created if it does not
+ already exist.
+ run_id : str
+ Run identifier; also the state file's basename (without ``.json``).
+ state : dict
+ JSON-serializable state to write. Expected (by ``teardown.py``) to
+ contain ``run_id``, ``region``, ``instance_id``, ``sg_id``,
+ ``key_name``, ``key_path``, ``public_ip``, ``purchasing_used``, and
+ ``ami_id``, but this function itself does not validate the shape
+ of ``state`` -- it is a thin, schema-agnostic writer so callers
+ (including future ones) are free to add fields.
+
+ Returns
+ -------
+ pathlib.Path
+ Path to the written state file.
+
+ Notes
+ -----
+ Side effect: writes ``/.json``, creating
+ ``state_dir`` if needed. Uses ``json.dump(..., default=str)`` so a
+ stray non-JSON-native value (e.g. if a caller forgets to stringify a
+ ``pathlib.Path``) is coerced to its string form instead of raising a
+ ``TypeError`` deep in a provisioning run.
+ """
+ state_dir = Path(state_dir)
+ state_dir.mkdir(parents=True, exist_ok=True)
+ state_path = state_dir / f"{run_id}.json"
+ with open(state_path, "w") as fh:
+ json.dump(state, fh, indent=2, default=str)
+ return state_path
+
+
+def _prompt_yes_no(prompt: str) -> bool:
+ """Ask an interactive yes/no question, returning ``True`` only for "yes".
+
+ Parameters
+ ----------
+ prompt : str
+ Text to show before the input cursor.
+
+ Returns
+ -------
+ bool
+ ``True`` if the user typed exactly ``"yes"`` (case-insensitive,
+ surrounding whitespace ignored); ``False`` for anything else,
+ including EOF/blank input. Requiring the full word "yes" (not just
+ "y") is a deliberate speed bump before an action that costs real
+ money.
+ """
+ try:
+ answer = input(prompt)
+ except EOFError:
+ return False
+ return answer.strip().lower() == "yes"
+
+
+def main(argv: Optional[list] = None) -> int:
+ """Parse CLI args and provision one instance end to end.
+
+ Parameters
+ ----------
+ argv : list[str] or None, default None
+ Argument list, as passed to ``argparse``'s ``parse_args``. ``None``
+ reads from ``sys.argv[1:]``.
+
+ Returns
+ -------
+ int
+ Process exit code: ``0`` on success (or a completed dry run),
+ ``1`` if the user declined the cost confirmation prompt.
+
+ Notes
+ -----
+ Side effects (skipped entirely when ``--dry-run`` is passed, except
+ for the ``DryRun=True`` API calls themselves): creates an SSH key pair
+ and local PEM, creates a security group, launches an instance, waits
+ for it to be SSH-reachable, and writes a state JSON file -- TWICE: once
+ immediately after launch (``public_ip`` as an explicit placeholder) and
+ again once ``wait_for_instance`` returns a real IP, so a crash or
+ interruption during the (potentially multi-minute) SSH wait still
+ leaves a local record of a running, billing instance. See the inline
+ "Design/WHY" comment at that first ``write_state`` call for the full
+ rationale. Prints a cost warning and requires interactive ``"yes"``
+ confirmation before any of that happens, unless ``--yes`` is passed.
+
+ This function deliberately never tears anything down itself, even on a
+ ``wait_for_instance`` failure -- ``provision.py``'s contract is
+ provision-and-leave-running (see ``README.md``'s "advanced / manual
+ control" section); recovery in that failure case is a printed
+ ``teardown.py`` command for the operator to run, not an automatic call.
+ """
+ parser = argparse.ArgumentParser(
+ prog="provision.py",
+ description=(
+ "Provision one p5.48xlarge (8x H100, NVSwitch) EC2 instance, "
+ "spot-first with on-demand fallback, for NCCL profiling."
+ ),
+ )
+ Config.add_args(parser)
+ parser.add_argument(
+ "--dry-run",
+ action="store_true",
+ help="Perform DryRun=True authorization checks only; launch nothing.",
+ )
+ parser.add_argument(
+ "--ssh-cidr",
+ type=str,
+ default=None,
+ help=(
+ "CIDR block allowed SSH access, e.g. 1.2.3.4/32. Overrides "
+ "auto-detected caller IP (see caller_ip())."
+ ),
+ )
+ parser.add_argument(
+ "--yes",
+ action="store_true",
+ help="Skip the interactive cost-confirmation prompt.",
+ )
+ args = parser.parse_args(argv)
+ cfg = Config.from_parsed(args)
+
+ # Fail fast on a missing boto3 before printing anything else -- no
+ # point asking the operator to confirm a cost warning for a run that
+ # cannot possibly proceed.
+ _require_boto3()
+
+ print(f"=== accelforge correlation-study provisioning: run_id={cfg.run_id} ===")
+ print(_COST_WARNING)
+ print(
+ f"Purchasing mode: {cfg.purchasing}. Instance type: {cfg.instance_type}. "
+ f"Region: {cfg.region}."
+ )
+
+ if not args.dry_run and not args.yes:
+ if not _prompt_yes_no("Proceed with provisioning? [yes/N]: "):
+ print("Aborted by user.")
+ return 1
+
+ ec2_client = boto3.client("ec2", region_name=cfg.region)
+ ssm_client = boto3.client("ssm", region_name=cfg.region)
+
+ ssh_cidr = args.ssh_cidr
+ if not ssh_cidr:
+ ip = caller_ip()
+ ssh_cidr = f"{ip}/32"
+ print(f"SSH will be allowed from: {ssh_cidr}")
+
+ ami_id = resolve_ami(ssm_client, cfg.ami_ssm_parameter)
+ print(f"Resolved AMI: {ami_id}")
+
+ key_name = f"{cfg.tag_project}-{cfg.run_id}"
+ key_path = ensure_key_pair(ec2_client, key_name, cfg.key_dir)
+ print(f"Key pair ready: {key_name} -> {key_path}")
+
+ sg_name = f"{cfg.tag_project}-{cfg.run_id}-sg"
+ sg_id = ensure_security_group(ec2_client, sg_name, ssh_cidr, cfg.tag_project, cfg.run_id)
+ print(f"Security group ready: {sg_id}")
+
+ launch_result = launch_instance(
+ ec2_client, cfg, ami_id, sg_id, key_name, dry_run=args.dry_run
+ )
+ if args.dry_run:
+ print(
+ f"Dry run complete (purchasing checked: {launch_result['purchasing_used']}); "
+ "no instance was launched."
+ )
+ return 0
+
+ instance_id = launch_result["instance_id"]
+ purchasing_used = launch_result["purchasing_used"]
+ print(f"Launched instance {instance_id} ({purchasing_used}).")
+
+ # Design/WHY (closes an orphan-instance window): write the state file --
+ # with public_ip as an explicit placeholder -- IMMEDIATELY after
+ # launch_instance returns, before wait_for_instance is even called.
+ # wait_for_instance can block for several minutes (instance boot, then
+ # polling for SSH) and can itself raise (TimeoutError, WaiterError, or
+ # simply be interrupted by Ctrl-C/a killed process). Before this fix, a
+ # crash in that window left an instance running and billing with NO
+ # local record of it at all, since write_state was only ever called
+ # once, after wait_for_instance returned successfully. Writing state
+ # now -- and rewriting it once the real public_ip is known below --
+ # means teardown.py can always find and tear down this run from local
+ # state alone, even if this process never gets past wait_for_instance.
+ state = {
+ "run_id": cfg.run_id,
+ "region": cfg.region,
+ "instance_id": instance_id,
+ "sg_id": sg_id,
+ "key_name": key_name,
+ "key_path": str(key_path),
+ "public_ip": None,
+ "purchasing_used": purchasing_used,
+ "ami_id": ami_id,
+ }
+ state_path = write_state(cfg.state_dir, cfg.run_id, state)
+ print(f"State written to: {state_path} (public_ip pending SSH reachability).")
+
+ print("Waiting for it to become reachable...")
+ try:
+ public_ip = wait_for_instance(ec2_client, instance_id)
+ except Exception:
+ # The instance is still running (and billing) regardless of why
+ # wait_for_instance failed -- print a loud, impossible-to-miss block
+ # naming exactly how to find and recover it, then re-raise
+ # unmodified so this failure still surfaces as a nonzero exit code
+ # (never silently swallowed).
+ print("=" * 70, file=sys.stderr)
+ print("ERROR: wait_for_instance failed (see traceback below).", file=sys.stderr)
+ print(f"Instance {instance_id} IS STILL RUNNING AND BILLING.", file=sys.stderr)
+ print(f"State file: {state_path}", file=sys.stderr)
+ print(
+ f"Recover with: python teardown.py --run-id {cfg.run_id} --region {cfg.region}",
+ file=sys.stderr,
+ )
+ print("=" * 70, file=sys.stderr)
+ raise
+
+ print(f"Instance is running and SSH-reachable at {public_ip}")
+
+ state["public_ip"] = public_ip
+ state_path = write_state(cfg.state_dir, cfg.run_id, state)
+
+ print(f"State written to: {state_path}")
+ print(f"Public IP: {public_ip}")
+ print(f"SSH command: ssh -i {key_path} {cfg.ssh_user}@{public_ip}")
+ print(
+ "Remember: tear this down when done "
+ f"(python teardown.py --run-id {cfg.run_id})."
+ )
+ return 0
+
+
+if __name__ == "__main__":
+ raise SystemExit(main())
diff --git a/notebooks/astrasim2_correlation/correlation/run_profile.sh b/notebooks/astrasim2_correlation/correlation/run_profile.sh
new file mode 100644
index 00000000..372f580e
--- /dev/null
+++ b/notebooks/astrasim2_correlation/correlation/run_profile.sh
@@ -0,0 +1,182 @@
+#!/usr/bin/env bash
+#
+# run_profile.sh -- run one profiling leg (FC or torus) for one or more
+# NCCL collectives on a provisioned 8x-H100 instance, and parse each raw
+# log into the unified CSV schema consumed by the correlation notebook.
+#
+# This script runs ON the EC2 instance (not locally); it has no GPUs to
+# talk to when checked out anywhere else, so it is verified here by
+# `bash -n` (syntax check) only -- see the WP2 report for details.
+#
+# Usage:
+# ./run_profile.sh \
+# [...]
+#
+# Env overrides:
+# NCCL_TESTS_DIR Path to a built nccl-tests checkout. Default: $HOME/nccl-tests
+# TORUS_BENCH_BIN Path to the built torus_bench binary. Default: $HOME/torus_bench/torus_bench
+# PARSE_NCCL Path to parse_nccl.py. Default: the copy next to this script.
+# WARMUP nccl-tests/torus_bench warmup iteration count. Default: 5
+# ITERS nccl-tests/torus_bench measured iteration count. Default: 20
+#
+# Design: fail fast and loud (set -euo pipefail) rather than silently
+# continuing past a failed collective run or a failed parse -- a partial,
+# uncaught failure here would otherwise show up much later as a confusing
+# gap in the correlation notebook's data rather than as a build/run error
+# on the instance where it's cheap to diagnose.
+set -euo pipefail
+
+# ---------------------------------------------------------------------------
+# Argument parsing
+# ---------------------------------------------------------------------------
+if [[ $# -lt 6 ]]; then
+ echo "Usage: $0 [...]" >&2
+ exit 1
+fi
+
+results_dir="$1"; shift
+topology="$1"; shift
+min_bytes="$1"; shift
+max_bytes="$1"; shift
+dims="$1"; shift
+# Remaining positional args are the list of collectives to profile in this leg.
+collectives=("$@")
+
+if [[ "$topology" != "fc" && "$topology" != "torus" ]]; then
+ echo "ERROR: must be 'fc' or 'torus', got '$topology'" >&2
+ exit 1
+fi
+
+# ---------------------------------------------------------------------------
+# Environment / defaults
+# ---------------------------------------------------------------------------
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+
+NCCL_TESTS_DIR="${NCCL_TESTS_DIR:-$HOME/nccl-tests}"
+TORUS_BENCH_BIN="${TORUS_BENCH_BIN:-$HOME/torus_bench/torus_bench}"
+# Design: PARSE_NCCL defaults to the copy sitting next to this script
+# (rather than requiring it on PATH or hard-coding an absolute install
+# path) so the pair of files can be scp'd to the instance as a unit and
+# just work.
+PARSE_NCCL="${PARSE_NCCL:-$SCRIPT_DIR/parse_nccl.py}"
+WARMUP="${WARMUP:-5}"
+ITERS="${ITERS:-20}"
+
+RAW_DIR="$results_dir/raw"
+CSV_DIR="$results_dir/csv"
+mkdir -p "$RAW_DIR" "$CSV_DIR"
+
+# ---------------------------------------------------------------------------
+# metadata.txt: written exactly once per invocation (not once per
+# collective), since it captures machine/software state that doesn't
+# change across the collectives loop below.
+# ---------------------------------------------------------------------------
+metadata_file="$results_dir/metadata.txt"
+{
+ echo "=== date ==="
+ date -u
+ echo
+ echo "=== uname -a ==="
+ uname -a
+ echo
+ echo "=== nvidia-smi --query-gpu=name,driver_version --format=csv ==="
+ nvidia-smi --query-gpu=name,driver_version --format=csv
+ echo
+ echo "=== nvidia-smi topo -m ==="
+ nvidia-smi topo -m
+ echo
+ echo "=== nccl-tests git rev ==="
+ # Best-effort: nccl-tests may not be a git checkout (e.g. if a tarball
+ # was scp'd instead), so a failed `git rev-parse` here must not abort
+ # the whole script under `set -e`.
+ git -C "$NCCL_TESTS_DIR" rev-parse HEAD 2>/dev/null || echo "unknown (not a git checkout or NCCL_TESTS_DIR missing)"
+} > "$metadata_file"
+echo "Wrote $metadata_file"
+
+# ---------------------------------------------------------------------------
+# FC leg: stock nccl-tests binaries.
+# ---------------------------------------------------------------------------
+run_fc_leg() {
+ local collective="$1"
+ local bin="$NCCL_TESTS_DIR/build/${collective}_perf"
+ local raw_log="$RAW_DIR/fc_${collective}.log"
+ local out_csv="$CSV_DIR/fc_${collective}.csv"
+
+ if [[ ! -x "$bin" ]]; then
+ echo "ERROR: nccl-tests binary not found: $bin" >&2
+ echo " Build nccl-tests first, e.g. via setup_node.sh, or:" >&2
+ echo " make -j -C \"$NCCL_TESTS_DIR\" MPI=0 CUDA_HOME=\"\${CUDA_HOME:-/usr/local/cuda}\"" >&2
+ exit 1
+ fi
+
+ echo "=== FC leg: $collective ==="
+ "$bin" -b "$min_bytes" -e "$max_bytes" -f 2 -g 8 -w "$WARMUP" -n "$ITERS" -c 1 | tee "$raw_log"
+ python3 "$PARSE_NCCL" "$raw_log" \
+ --source nccl-tests \
+ --collective "$collective" \
+ --topology fc \
+ --dims "$dims" \
+ --out "$out_csv"
+
+ # Fail loudly on a header-only CSV: parse_nccl.py can exit 0 while emitting zero data rows
+ # (e.g. every swept size hit a divisibility skip for this dims/size combination), which
+ # would otherwise look identical to a real, successful leg -- a silent data gap discovered
+ # only much later in the correlation notebook, rather than here, where the raw log needed
+ # to diagnose it is still on disk and cheap to inspect.
+ csv_lines="$(wc -l < "$out_csv")"
+ if [[ "$csv_lines" -le 1 ]]; then
+ echo "ERROR: $out_csv has no data rows (header-only or empty) -- see $raw_log" >&2
+ exit 1
+ fi
+ echo "Wrote $out_csv"
+}
+
+# ---------------------------------------------------------------------------
+# Torus leg: custom torus_bench binary (built by the sibling work package).
+# ---------------------------------------------------------------------------
+run_torus_leg() {
+ local collective="$1"
+ local raw_log="$RAW_DIR/torus_${collective}.log"
+ local out_csv="$CSV_DIR/torus_${collective}.csv"
+
+ if [[ ! -x "$TORUS_BENCH_BIN" ]]; then
+ echo "ERROR: torus_bench binary not found: $TORUS_BENCH_BIN" >&2
+ echo " Build torus_bench first (see setup_node.sh), e.g.:" >&2
+ echo " make -C \"\$TORUS_BENCH_DIR\" torus_bench" >&2
+ exit 1
+ fi
+
+ echo "=== Torus leg: $collective ==="
+ "$TORUS_BENCH_BIN" --collective "$collective" --dims "$dims" \
+ -b "$min_bytes" -e "$max_bytes" -f 2 -w "$WARMUP" -n "$ITERS" --check \
+ | tee "$raw_log"
+ python3 "$PARSE_NCCL" "$raw_log" \
+ --source torus_bench \
+ --collective "$collective" \
+ --topology torus \
+ --dims "$dims" \
+ --out "$out_csv"
+
+ # See run_fc_leg's identical check above for the full rationale: a header-only CSV here
+ # (e.g. every swept size hit torus_bench's own divisibility SKIP for this dims/size
+ # combination) must fail the script loudly rather than silently passing as "done".
+ csv_lines="$(wc -l < "$out_csv")"
+ if [[ "$csv_lines" -le 1 ]]; then
+ echo "ERROR: $out_csv has no data rows (header-only or empty) -- see $raw_log" >&2
+ exit 1
+ fi
+ echo "Wrote $out_csv"
+}
+
+# ---------------------------------------------------------------------------
+# Main loop: one leg (selected by $topology), all requested collectives.
+# ---------------------------------------------------------------------------
+for collective in "${collectives[@]}"; do
+ if [[ "$topology" == "fc" ]]; then
+ run_fc_leg "$collective"
+ else
+ run_torus_leg "$collective"
+ fi
+done
+
+echo "Done. Results in $results_dir"
diff --git a/notebooks/astrasim2_correlation/correlation/setup_node.sh b/notebooks/astrasim2_correlation/correlation/setup_node.sh
new file mode 100644
index 00000000..ec407e99
--- /dev/null
+++ b/notebooks/astrasim2_correlation/correlation/setup_node.sh
@@ -0,0 +1,137 @@
+#!/usr/bin/env bash
+#
+# setup_node.sh -- idempotent one-time (but safe-to-rerun) setup for an
+# EC2 p5.48xlarge (8x H100, NVSwitch) profiling instance: arms a dead-man
+# shutdown, verifies the GPU fabric is visible, builds nccl-tests and
+# torus_bench if not already built, and prints a versions summary.
+#
+# This script runs ON the EC2 instance; it is verified here by `bash -n`
+# (syntax check) only -- see the WP2 report for details.
+#
+# Env:
+# DEADMAN_MINUTES Minutes until the dead-man shutdown fires. Default: 120
+# NCCL_TESTS_DIR Where to clone/build nccl-tests. Default: $HOME/nccl-tests
+# TORUS_BENCH_DIR Where torus_bench is scp'd/built. Default: $HOME/torus_bench
+# CUDA_HOME CUDA toolkit root used to build nccl-tests. Default: /usr/local/cuda
+set -euo pipefail
+
+DEADMAN_MINUTES="${DEADMAN_MINUTES:-120}"
+NCCL_TESTS_DIR="${NCCL_TESTS_DIR:-$HOME/nccl-tests}"
+TORUS_BENCH_DIR="${TORUS_BENCH_DIR:-$HOME/torus_bench}"
+CUDA_HOME="${CUDA_HOME:-/usr/local/cuda}"
+
+# ---------------------------------------------------------------------------
+# Step 1: arm the dead-man switch FIRST, before anything else can fail or
+# hang.
+#
+# Design/WHY this must be first: the instance is launched with
+# shutdown-behavior=terminate, so a `shutdown -P` here is what actually
+# terminates (not just stops) the instance and stops billing. If setup
+# were to fail, hang (e.g. a stuck `make`, a stalled git clone over a flaky
+# network), or if the orchestrating controller process on the caller's side
+# dies/loses connectivity, this is the only backstop that guarantees the
+# (expensive, 8x H100) instance doesn't run forever. Arming it before any
+# other step -- including the GPU sanity check below, which could itself
+# hang on a broken driver -- ensures the cost cap applies unconditionally
+# from the very start of setup, not only after setup "succeeds".
+# ---------------------------------------------------------------------------
+echo "Arming dead-man shutdown: instance will terminate in ${DEADMAN_MINUTES} minutes unless this script (or a later run of it) is used to push it back further."
+# Cancel any already-pending shutdown before arming a new one: issuing a second `shutdown`
+# while one is already pending errors on some systemd versions, and cancel-then-rearm also
+# makes a re-run of this script push the deadline BACK (rather than erroring or stacking),
+# which is the desired semantics for a legitimate re-setup (e.g. extending a long-running
+# sweep with a fresh DEADMAN_MINUTES). `|| true` because there being no pending shutdown to
+# cancel (the common case, e.g. this script's first run) is not an error.
+sudo shutdown -c 2>/dev/null || true
+sudo shutdown -P "+${DEADMAN_MINUTES}"
+
+# ---------------------------------------------------------------------------
+# Step 2: verify nvidia-smi works and all 8 GPUs are visible. Fail loudly
+# (rather than proceeding to build against a broken/partial driver) since
+# every downstream profiling run depends on this.
+# ---------------------------------------------------------------------------
+if ! command -v nvidia-smi >/dev/null 2>&1; then
+ echo "ERROR: nvidia-smi not found on PATH. Is the NVIDIA driver installed?" >&2
+ exit 1
+fi
+
+if ! nvidia-smi >/dev/null; then
+ echo "ERROR: nvidia-smi is present but failed to run. Driver/GPU problem?" >&2
+ exit 1
+fi
+
+gpu_count="$(nvidia-smi -L | wc -l)"
+if [[ "$gpu_count" -ne 8 ]]; then
+ echo "ERROR: expected 8 GPUs (p5.48xlarge), found $gpu_count. Aborting." >&2
+ exit 1
+fi
+echo "OK: nvidia-smi reports $gpu_count GPUs."
+
+# ---------------------------------------------------------------------------
+# Step 3: build nccl-tests if it isn't already built.
+#
+# Idempotent: only clones if NCCL_TESTS_DIR doesn't exist yet, and only
+# (re)builds if the all_reduce_perf binary is missing -- a rerun of this
+# script after a successful first run is a fast no-op here.
+# ---------------------------------------------------------------------------
+if [[ -x "$NCCL_TESTS_DIR/build/all_reduce_perf" ]]; then
+ echo "OK: nccl-tests already built at $NCCL_TESTS_DIR."
+else
+ if [[ ! -d "$NCCL_TESTS_DIR" ]]; then
+ echo "Cloning nccl-tests into $NCCL_TESTS_DIR ..."
+ git clone https://github.com/NVIDIA/nccl-tests "$NCCL_TESTS_DIR"
+ fi
+ echo "Building nccl-tests (MPI=0, CUDA_HOME=$CUDA_HOME) ..."
+ make -j -C "$NCCL_TESTS_DIR" MPI=0 CUDA_HOME="$CUDA_HOME"
+fi
+
+# ---------------------------------------------------------------------------
+# Step 4: build torus_bench if it isn't already built.
+#
+# TORUS_BENCH_DIR is scp'd onto the instance by the orchestrator (a sibling
+# work package owns the torus_bench source); if it hasn't landed yet, warn
+# and continue rather than failing -- the FC leg (nccl-tests) can still run
+# without it, and setup_node.sh may legitimately run before the orchestrator
+# has finished copying torus_bench over.
+# ---------------------------------------------------------------------------
+torus_bench_bin="$TORUS_BENCH_DIR/torus_bench"
+if [[ -x "$torus_bench_bin" ]]; then
+ echo "OK: torus_bench already built at $torus_bench_bin."
+elif [[ -d "$TORUS_BENCH_DIR" ]]; then
+ echo "Building torus_bench ..."
+ make -C "$TORUS_BENCH_DIR" torus_bench
+else
+ echo "WARNING: $TORUS_BENCH_DIR not found (expected to be scp'd there by the orchestrator)." >&2
+ echo " Skipping torus_bench build; the FC leg can still run without it." >&2
+fi
+
+# ---------------------------------------------------------------------------
+# Step 5: print a versions summary for the run's metadata/provenance.
+# Every lookup here is best-effort (guarded so a missing tool doesn't abort
+# the script under `set -e`), since this is diagnostic output, not a hard
+# requirement.
+# ---------------------------------------------------------------------------
+echo "=== Versions ==="
+
+if command -v nvcc >/dev/null 2>&1; then
+ nvcc --version | tail -n 1
+else
+ echo "nvcc: not found on PATH"
+fi
+
+nvidia-smi --query-gpu=driver_version --format=csv,noheader | head -n 1 | sed 's/^/driver_version: /'
+
+# NCCL version discovery: prefer asking a Python-visible torch build (most
+# accurate for the environment that will actually run collectives), and
+# fall back to scanning the linker cache for libnccl if torch isn't
+# importable. Both are best-effort.
+if python3 -c "import torch; print(torch.cuda.nccl.version())" 2>/dev/null; then
+ :
+elif ldconfig -p | grep -qi libnccl; then
+ echo "libnccl found via ldconfig:"
+ ldconfig -p | grep -i libnccl
+else
+ echo "NCCL version: could not be determined (no torch, no libnccl in ldconfig cache)"
+fi
+
+echo "setup_node.sh complete."
diff --git a/notebooks/astrasim2_correlation/correlation/teardown.py b/notebooks/astrasim2_correlation/correlation/teardown.py
new file mode 100644
index 00000000..e533125a
--- /dev/null
+++ b/notebooks/astrasim2_correlation/correlation/teardown.py
@@ -0,0 +1,850 @@
+"""Tear down AWS resources created by ``provision.py``.
+
+This script is the "down" half of the correlation study's empirical leg
+(see ``provision.py`` for the "up" half and ``README.md`` for the full
+runbook). It is designed to be safe to run more than once, and to work
+even when its own local state is missing or stale, because it discovers
+targets two ways and reconciles them:
+
+1. Local state files under ``.state/*.json``, written by
+ ``provision.py``'s ``write_state`` -- the fast, detailed path, since a
+ state file already has the security group id and key pair name without
+ any extra API calls.
+2. A live ``describe_instances`` search filtered on the ``Project`` tag
+ (and ``RunId`` tag, when ``--run-id`` is given) -- the authoritative
+ path, since it reflects what AWS actually has running right now even if
+ a state file was deleted, never written (a crash mid-provision), or the
+ run was started from a different machine/checkout.
+
+Every function below is written to be independently importable and
+testable, matching ``provision.py``'s convention; see that module's
+docstring for why boto3 is imported guarded rather than as a hard
+dependency.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+import time
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+from config import Config
+
+try:
+ import boto3
+ from botocore.exceptions import ClientError
+except ImportError: # pragma: no cover - exercised only when boto3 truly absent
+ boto3 = None
+
+ # See provision.py's identical placeholder for why this exists: keeps
+ # `except ClientError:` clauses valid Python even without boto3
+ # installed, without ever actually being reachable (main() always
+ # calls _require_boto3() first).
+ class ClientError(Exception): # type: ignore[no-redef]
+ pass
+
+
+def _require_boto3() -> None:
+ """Raise a clear, actionable error if boto3 is not installed.
+
+ Raises
+ ------
+ SystemExit
+ Always, if ``boto3`` failed to import. See ``provision._require_boto3``
+ for the identical rationale; kept as a separate copy here (rather
+ than importing it from ``provision``) so this module has no
+ import-time dependency on ``provision.py`` at all.
+ """
+ if boto3 is None:
+ raise SystemExit(
+ "boto3 is required for AWS teardown but is not installed in this "
+ "Python environment.\n"
+ "Install it with: pip install boto3"
+ )
+
+
+# Instance states worth discovering/tearing down. Deliberately excludes
+# "shutting-down" and "terminated": those instances are already on their
+# way out or gone and do not need (and, for "terminated", cannot receive)
+# a terminate_instances call.
+_ACTIVE_STATES = ("pending", "running", "stopping", "stopped")
+
+# EC2 does not release the ENI-to-security-group association the instant
+# terminate_instances returns (or the instance_terminated waiter is
+# satisfied); delete_security_group can fail with DependencyViolation for a
+# short window afterward while that teardown finishes propagating. Retrying
+# with a fixed backoff is expected to succeed within a few attempts rather
+# than being a genuine, permanent conflict.
+_SG_DELETE_MAX_RETRIES = 5
+_SG_DELETE_RETRY_SLEEP_S = 5.0
+
+
+def find_tagged_instances(
+ ec2_client, tag_project: str, run_id: Optional[str] = None
+) -> List[dict]:
+ """Find EC2 instances tagged for this study, optionally scoped to one run.
+
+ Parameters
+ ----------
+ ec2_client : botocore.client.BaseClient
+ A boto3 ``ec2`` client (or a stub thereof).
+ tag_project : str
+ Value the ``Project`` tag must match.
+ run_id : str or None, default None
+ If given, additionally require the ``RunId`` tag to match this
+ value. If ``None``, instances from every run under ``tag_project``
+ are returned.
+
+ Returns
+ -------
+ list[dict]
+ Raw ``Instance`` dicts (the ``Reservations[].Instances[]`` shape
+ returned by ``describe_instances``), for instances currently in
+ one of :data:`_ACTIVE_STATES`. Empty list if none match.
+
+ Notes
+ -----
+ Paginates via ``NextToken`` manually (rather than
+ ``ec2_client.get_paginator(...)``) so this function works identically
+ against a plain client and a ``botocore.stub.Stubber``-wrapped one used
+ in tests, without needing the Stubber to understand paginator internals.
+ """
+ filters = [
+ {"Name": "tag:Project", "Values": [tag_project]},
+ {"Name": "instance-state-name", "Values": list(_ACTIVE_STATES)},
+ ]
+ if run_id:
+ filters.append({"Name": "tag:RunId", "Values": [run_id]})
+
+ instances: List[dict] = []
+ kwargs: Dict[str, Any] = {"Filters": filters}
+ while True:
+ response = ec2_client.describe_instances(**kwargs)
+ for reservation in response.get("Reservations", []):
+ instances.extend(reservation.get("Instances", []))
+ next_token = response.get("NextToken")
+ if not next_token:
+ break
+ kwargs["NextToken"] = next_token
+ return instances
+
+
+def _delete_security_group_with_retry(
+ ec2_client,
+ sg_id: str,
+ max_retries: int = _SG_DELETE_MAX_RETRIES,
+ retry_sleep_s: float = _SG_DELETE_RETRY_SLEEP_S,
+) -> None:
+ """Delete a security group, retrying on ``DependencyViolation``.
+
+ Parameters
+ ----------
+ ec2_client : botocore.client.BaseClient
+ A boto3 ``ec2`` client (or a stub thereof).
+ sg_id : str
+ Security group id to delete.
+ max_retries : int, default 5
+ Maximum number of ``delete_security_group`` attempts.
+ retry_sleep_s : float, default 5.0
+ Seconds to sleep between retries.
+
+ Raises
+ ------
+ RuntimeError
+ If every attempt fails with ``DependencyViolation`` (the ENI
+ association never cleared in time).
+ botocore.exceptions.ClientError
+ For any ``ClientError`` code other than ``DependencyViolation`` or
+ ``InvalidGroup.NotFound``, propagated immediately without retry.
+
+ Notes
+ -----
+ See the module-level comment on :data:`_SG_DELETE_MAX_RETRIES` for why
+ ``DependencyViolation`` specifically is retried rather than treated as
+ fatal on the first failure.
+ """
+ last_exc: Optional[ClientError] = None
+ for attempt in range(1, max_retries + 1):
+ try:
+ ec2_client.delete_security_group(GroupId=sg_id)
+ print(f"Deleted security group {sg_id}.")
+ return
+ except ClientError as exc:
+ code = exc.response.get("Error", {}).get("Code", "")
+ if code == "InvalidGroup.NotFound":
+ print(f"Security group {sg_id} already gone.")
+ return
+ if code != "DependencyViolation":
+ raise
+ last_exc = exc
+ if attempt < max_retries:
+ print(
+ f"delete_security_group({sg_id}) hit DependencyViolation "
+ f"(attempt {attempt}/{max_retries}); retrying in "
+ f"{retry_sleep_s}s..."
+ )
+ time.sleep(retry_sleep_s)
+ raise RuntimeError(
+ f"Failed to delete security group {sg_id} after {max_retries} attempts "
+ "due to a persistent DependencyViolation."
+ ) from last_exc
+
+
+def teardown_run(ec2_client, state: dict, delete_key: bool) -> None:
+ """Terminate one run's instance and clean up its security group and key.
+
+ Parameters
+ ----------
+ ec2_client : botocore.client.BaseClient
+ A boto3 ``ec2`` client (or a stub thereof).
+ state : dict
+ Per-run state, either loaded from a ``.state/*.json`` file (see
+ ``provision.write_state``) or synthesized from a live
+ ``describe_instances`` result (see ``_derive_state_from_instance``).
+ Recognized keys: ``instance_id``, ``sg_id``, ``key_name``,
+ ``key_path``. All are optional -- a missing key simply skips that
+ cleanup step, so a partially-populated state (e.g. derived from
+ AWS alone, with no known ``key_path``) still tears down whatever it
+ can.
+ delete_key : bool
+ If ``True``, also delete the AWS-side key pair (and the local PEM,
+ if ``state["key_path"]`` is known and exists on disk).
+
+ Raises
+ ------
+ RuntimeError
+ If security group deletion exhausts its retries (see
+ :func:`_delete_security_group_with_retry`).
+ botocore.exceptions.ClientError
+ For any AWS failure other than the specific "already gone" codes
+ this function is written to tolerate (``InvalidInstanceID.NotFound``
+ for the instance, ``InvalidGroup.NotFound`` for the security
+ group), propagated unmodified.
+ botocore.exceptions.WaiterError
+ If the ``instance_terminated`` waiter times out or the instance
+ reaches an unexpected terminal state.
+
+ Notes
+ -----
+ Order matters: instance termination is started and waited on *before*
+ security group deletion is attempted, because the security group
+ cannot be deleted while an instance's network interface still
+ references it (see :data:`_SG_DELETE_MAX_RETRIES`'s docstring).
+ """
+ instance_id = state.get("instance_id")
+ if instance_id:
+ try:
+ ec2_client.terminate_instances(InstanceIds=[instance_id])
+ except ClientError as exc:
+ code = exc.response.get("Error", {}).get("Code", "")
+ if code != "InvalidInstanceID.NotFound":
+ raise
+ print(
+ f"Instance {instance_id} already gone "
+ "(InvalidInstanceID.NotFound); continuing teardown."
+ )
+ else:
+ print(
+ f"Termination requested for {instance_id}; waiting for it to "
+ "fully terminate..."
+ )
+ waiter = ec2_client.get_waiter("instance_terminated")
+ waiter.wait(InstanceIds=[instance_id])
+ print(f"Instance {instance_id} terminated.")
+
+ sg_id = state.get("sg_id")
+ if sg_id:
+ _delete_security_group_with_retry(ec2_client, sg_id)
+
+ if delete_key:
+ key_name = state.get("key_name")
+ if key_name:
+ try:
+ ec2_client.delete_key_pair(KeyName=key_name)
+ print(f"Deleted AWS key pair {key_name}.")
+ except ClientError as exc:
+ # Best-effort: a stale/already-deleted key pair should not
+ # block the rest of teardown, but it is still reported, not
+ # silently dropped.
+ print(f"WARNING: failed to delete AWS key pair {key_name}: {exc}")
+
+ key_path = state.get("key_path")
+ if key_path:
+ local_path = Path(key_path)
+ if local_path.exists():
+ local_path.unlink()
+ print(f"Deleted local PEM {local_path}.")
+ elif key_name:
+ print(
+ f"No local PEM path known for key pair {key_name!r} (this run's "
+ "state was derived from AWS alone, not a local state file); "
+ "only the AWS-side key pair was deleted. If a PEM for it exists "
+ "on this or another machine, remove it manually."
+ )
+
+
+def _derive_state_from_instance(instance: dict) -> dict:
+ """Reconstruct a minimal teardown ``state`` dict from a live instance.
+
+ Used when a tagged instance is discovered via :func:`find_tagged_instances`
+ but has no matching local ``.state/*.json`` file (deleted, never
+ written, or written on a different machine). EC2 instance descriptions
+ already carry everything ``teardown_run`` needs except the local PEM
+ path, which cannot be recovered this way.
+
+ Parameters
+ ----------
+ instance : dict
+ One ``Instance`` dict as returned by ``describe_instances``.
+
+ Returns
+ -------
+ dict
+ ``{"run_id", "instance_id", "sg_id", "key_name", "key_path"}``,
+ with ``key_path`` always ``None`` (see above) and ``run_id`` taken
+ from the instance's ``RunId`` tag, falling back to the instance id
+ itself if that tag is somehow missing.
+ """
+ tags = {t["Key"]: t["Value"] for t in instance.get("Tags", [])}
+ security_groups = instance.get("SecurityGroups", [])
+ return {
+ "run_id": tags.get("RunId", instance["InstanceId"]),
+ "instance_id": instance["InstanceId"],
+ "sg_id": security_groups[0]["GroupId"] if security_groups else None,
+ "key_name": instance.get("KeyName"),
+ "key_path": None,
+ }
+
+
+def _state_dir_files(state_dir: Path, run_id: Optional[str]) -> List[Path]:
+ """List local state files relevant to this teardown invocation.
+
+ Parameters
+ ----------
+ state_dir : pathlib.Path
+ Directory containing ``.json`` state files.
+ run_id : str or None
+ If given, look only for ``/.json``. If
+ ``None``, return every ``*.json`` file in ``state_dir``.
+
+ Returns
+ -------
+ list[pathlib.Path]
+ Matching, existing file paths, sorted for deterministic output.
+ Empty list if ``state_dir`` does not exist or nothing matches.
+ """
+ if not state_dir.exists():
+ return []
+ if run_id:
+ candidate = state_dir / f"{run_id}.json"
+ return [candidate] if candidate.exists() else []
+ return sorted(state_dir.glob("*.json"))
+
+
+def _load_state_file(path: Path) -> dict:
+ """Load one state JSON file.
+
+ Parameters
+ ----------
+ path : pathlib.Path
+ Path to a ``.state/.json`` file.
+
+ Returns
+ -------
+ dict
+ The parsed JSON content.
+
+ Raises
+ ------
+ OSError
+ If ``path`` cannot be opened.
+ json.JSONDecodeError
+ If ``path`` does not contain valid JSON.
+ """
+ with open(path, "r") as fh:
+ return json.load(fh)
+
+
+def _prompt_yes_no(prompt: str) -> bool:
+ """Ask an interactive yes/no question, returning ``True`` only for "yes".
+
+ Parameters
+ ----------
+ prompt : str
+ Text to show before the input cursor.
+
+ Returns
+ -------
+ bool
+ ``True`` only if the user typed exactly ``"yes"``
+ (case-insensitive); ``False`` otherwise, including on EOF.
+ """
+ try:
+ answer = input(prompt)
+ except EOFError:
+ return False
+ return answer.strip().lower() == "yes"
+
+
+def _verify(ec2_client, tag_project: str, run_id: Optional[str]) -> int:
+ """Audit for any still-running tagged instances, without tearing anything down.
+
+ Parameters
+ ----------
+ ec2_client : botocore.client.BaseClient
+ A boto3 ``ec2`` client (or a stub thereof).
+ tag_project : str
+ Value the ``Project`` tag must match.
+ run_id : str or None
+ If given, scope the audit to just this run.
+
+ Returns
+ -------
+ int
+ ``0`` (and prints ``"no running instances"``) if nothing tagged
+ remains in :data:`_ACTIVE_STATES`; ``1`` (and prints a table) if
+ anything does. Intended for use as a CI/cron safety check after a
+ teardown, so a stuck resource is caught rather than silently
+ left running and accruing cost.
+ """
+ instances = find_tagged_instances(ec2_client, tag_project, run_id)
+ if not instances:
+ print("no running instances")
+ return 0
+
+ print("Tagged instances still present:")
+ print(f"{'InstanceId':<21} {'State':<12} {'RunId'}")
+ for instance in instances:
+ tags = {t["Key"]: t["Value"] for t in instance.get("Tags", [])}
+ state_name = instance.get("State", {}).get("Name", "unknown")
+ print(f"{instance['InstanceId']:<21} {state_name:<12} {tags.get('RunId', '?')}")
+ return 1
+
+
+# ---------------------------------------------------------------------------
+# Region resolution (Fix 4a)
+# ---------------------------------------------------------------------------
+#
+# Design/WHY: before this fix, teardown.py always defaulted --region to
+# "us-east-1" and used exactly that one region for every discovery/teardown
+# call, regardless of what region a run's own state file recorded. A run
+# provisioned in any other region (e.g. because capacity/quota forced a
+# different --region at provision time) was therefore invisible to
+# `teardown.py --all` and to a bare `teardown.py --run-id ` unless the
+# operator remembered to pass --region explicitly every time -- silently
+# leaving that run's instance running and billing. The fix: --region now
+# defaults to None (so this module can tell "the user explicitly asked for
+# us-east-1" apart from "the user said nothing"), and region resolution
+# follows this priority, per run:
+# 1. an explicit --region flag always wins (it is an explicit override of
+# whatever a state file might say, e.g. for recovering a run whose
+# state file was hand-edited or lost a region field);
+# 2. otherwise, a matching local state file's own recorded "region" field
+# is authoritative (this is what makes --all correctly span multiple
+# regions in one invocation);
+# 3. otherwise (no --region, no matching/region-bearing state file --
+# e.g. an instance discovered via live AWS tags alone, with no local
+# state at all), fall back to Config's own default region
+# ("us-east-1"), matching this module's pre-fix behavior for that one
+# case.
+# resolve_regions() is a pure function (no AWS calls, no I/O) precisely so
+# this priority logic is unit-testable on its own, per the work-package
+# spec's explicit ask.
+
+
+def resolve_regions(args, states: Dict[str, dict]) -> Dict[str, List[str]]:
+ """Determine which AWS region(s) to operate in, and which run_ids live in each.
+
+ See the "Region resolution" design comment immediately above this
+ function for the full priority rationale (explicit ``--region`` flag,
+ then a state file's own recorded region, then Config's default).
+
+ Parameters
+ ----------
+ args : argparse.Namespace or any object with ``.region`` and ``.run_id``
+ Only ``args.region`` (str or None) and ``args.run_id`` (str or
+ None) are consulted; duck-typed so a test can pass a minimal stand-in
+ without building a full parsed CLI namespace.
+ states : dict[str, dict]
+ Every locally known state dict, keyed by run_id, as loaded from
+ ``.state/*.json`` files -- NOT pre-filtered to ``args.run_id``; this
+ function does that scoping itself.
+
+ Returns
+ -------
+ dict[str, list[str]]
+ Mapping of region -> list of run_ids (drawn from `states`) resolved
+ to that region.
+
+ - If ``args.run_id`` is set: exactly one key (the resolved region
+ for that one run), mapping to ``[args.run_id]``.
+ - Otherwise (``--all`` or a bare ``--verify``): one key per distinct
+ region recorded across every entry in `states` (grouped), PLUS
+ the default/fallback region (``args.region`` if given, else
+ Config's default) as its own key even if no state file happens to
+ record it -- so a caller iterating this mapping's keys always
+ still searches that region too, for tag-only discovery of
+ instances with no matching local state file at all (e.g. a crash
+ before ``provision.write_state`` ever ran).
+
+ Notes
+ -----
+ Pure function: makes no AWS calls and performs no I/O, so it is
+ directly unit-testable without a Stubber or any boto3 client.
+
+ Examples
+ --------
+ >>> import argparse
+ >>> args = argparse.Namespace(region=None, run_id=None)
+ >>> resolve_regions(args, {"run-a": {"region": "us-west-2"}})
+ {'us-east-1': [], 'us-west-2': ['run-a']}
+ """
+ default_region = args.region or Config().region
+
+ if args.run_id:
+ state = states.get(args.run_id)
+ region = args.region or (state.get("region") if state else None) or default_region
+ return {region: [args.run_id]}
+
+ regions: Dict[str, List[str]] = {default_region: []}
+ for run_id, state in states.items():
+ region = args.region or state.get("region") or default_region
+ regions.setdefault(region, []).append(run_id)
+ return regions
+
+
+def _client_for_region(clients: Dict[str, Any], region: str):
+ """Return a cached boto3 ``ec2`` client for `region`, creating it on first use.
+
+ Parameters
+ ----------
+ clients : dict[str, botocore.client.BaseClient]
+ Mutable cache, keyed by region name; mutated in place on a cache
+ miss. Callers own the dict's lifetime (typically one per
+ :func:`main` invocation).
+ region : str
+ AWS region name to build (or reuse) a client for.
+
+ Returns
+ -------
+ botocore.client.BaseClient
+ A boto3 ``ec2`` client bound to `region`, reused across calls that
+ pass the same `region` and the same `clients` dict.
+
+ Notes
+ -----
+ Design: ``--all`` may need to talk to several regions in one
+ invocation (one per distinct region recorded across this study's state
+ files -- see :func:`resolve_regions`). Building a client lazily, keyed
+ by region, keeps the number of client objects (and any connection-pool
+ overhead) proportional to the number of *distinct regions* actually
+ involved in one run, rather than the number of runs or the number of
+ times a region happens to be revisited.
+ """
+ if region not in clients:
+ clients[region] = boto3.client("ec2", region_name=region)
+ return clients[region]
+
+
+def _teardown_one_region(
+ ec2_client,
+ tag_project: str,
+ region: str,
+ run_id_filter: Optional[str],
+ file_states_by_run: Dict[str, dict],
+ file_paths_by_run: Dict[str, Path],
+ delete_key: bool,
+ yes: bool,
+) -> int:
+ """Discover, stale-clean, confirm, and tear down every matching resource in one region.
+
+ Factored out of :func:`main` so :func:`main` itself only needs to
+ resolve regions (:func:`resolve_regions`) and loop over them; this is
+ also what makes the stale-state cleanup path (Fix 4b) directly
+ unit-testable with a single ``botocore.stub.Stubber``-wrapped client,
+ without needing to drive the CLI/argparse layer or monkeypatch
+ ``boto3.client`` at all.
+
+ Parameters
+ ----------
+ ec2_client : botocore.client.BaseClient
+ A boto3 ``ec2`` client (or a stub thereof) already bound to
+ `region`.
+ tag_project : str
+ Value the ``Project`` tag must match (``Config().tag_project``).
+ region : str
+ The AWS region `ec2_client` is bound to; used only in printed
+ messages (this function performs no region validation of its own).
+ run_id_filter : str or None
+ If given, :func:`find_tagged_instances` is scoped to just this one
+ run (mirrors the original single-``--run-id`` behavior); ``None``
+ discovers every tagged run in `region`.
+ file_states_by_run : dict[str, dict]
+ Locally known state dicts for the run_ids :func:`resolve_regions`
+ assigned to this region, keyed by run_id.
+ file_paths_by_run : dict[str, pathlib.Path]
+ The corresponding ``.state/.json`` paths, same keys as
+ `file_states_by_run`.
+ delete_key : bool
+ Forwarded to every :func:`teardown_run` call in this region.
+ yes : bool
+ If ``False``, prompts for interactive confirmation before tearing
+ down any LIVE target found in this region. Stale-state cleanup
+ (see Notes) is never gated on this prompt: it only reclaims
+ resources this function has already determined are orphaned (no
+ matching live instance), not anything newly discovered as
+ still-in-use.
+
+ Returns
+ -------
+ int
+ ``0`` if every teardown/cleanup in this region succeeded (including
+ "nothing to do here"); ``1`` if any individual teardown/cleanup
+ raised, or if the user declined confirmation for this region's live
+ targets.
+
+ Notes
+ -----
+ Order of operations, and why (Fix 4b): (1) live AWS discovery via
+ :func:`find_tagged_instances`, reconciled against `file_states_by_run`
+ exactly as the pre-fix module-level code did; (2) any state file in
+ `file_states_by_run` with NO matching live instance is "stale" --
+ rather than just unlinking its state file (the pre-fix behavior), it is
+ now routed through :func:`teardown_run` FIRST. ``teardown_run``'s
+ terminate step already tolerates ``InvalidInstanceID.NotFound``, so a
+ truly-dead instance is a safe no-op there -- but its security group and
+ (optionally) key pair do NOT disappear just because the instance is
+ gone, e.g. via the on-instance dead-man timer terminating it out from
+ under a local state file that was never cleaned up. Routing through
+ ``teardown_run`` first closes that SG/key-pair leak; the state file is
+ removed only if that call did NOT raise, so a failed stale cleanup
+ leaves the state file in place for a future retry instead of silently
+ losing track of the leak. (3) only THEN are any remaining LIVE targets
+ confirmed and torn down, matching the pre-fix confirmation UX.
+ """
+ file_states_by_instance: Dict[str, dict] = {
+ s["instance_id"]: s for s in file_states_by_run.values() if s.get("instance_id")
+ }
+ aws_instances = find_tagged_instances(ec2_client, tag_project, run_id_filter)
+
+ targets: List[dict] = []
+ for instance in aws_instances:
+ instance_id = instance["InstanceId"]
+ if instance_id in file_states_by_instance:
+ targets.append(file_states_by_instance[instance_id])
+ else:
+ targets.append(_derive_state_from_instance(instance))
+
+ exit_code = 0
+
+ # --- Stale-state cleanup: see the "Order of operations" note above. ---
+ live_run_ids = {t.get("run_id") for t in targets if t.get("run_id")}
+ for run_id_key in set(file_states_by_run) - live_run_ids:
+ print(
+ f"State file for run {run_id_key!r} (region {region}) has no matching live AWS "
+ "instance; routing it through teardown to reclaim any leftover security group/key "
+ "pair before removing the state file."
+ )
+ try:
+ teardown_run(ec2_client, file_states_by_run[run_id_key], delete_key=delete_key)
+ except Exception as exc: # noqa: BLE001
+ # Design: deliberately broad, matching the identical rationale
+ # on the live-target teardown loop below -- one stale run's
+ # cleanup failure must not abort the rest of a --all batch.
+ print(f"ERROR cleaning up stale run {run_id_key!r}: {exc}", file=sys.stderr)
+ exit_code = 1
+ continue
+ state_path = file_paths_by_run.get(run_id_key)
+ if state_path and state_path.exists():
+ state_path.unlink()
+ print(f"Removed state file {state_path}.")
+
+ if not targets:
+ if not file_states_by_run:
+ print(f"No matching resources found in {region}.")
+ return exit_code
+
+ print(f"The following will be torn down in {region}:")
+ for t in targets:
+ print(
+ f" run_id={t.get('run_id', '?')} instance_id={t.get('instance_id', '?')} "
+ f"sg_id={t.get('sg_id', '?')}"
+ )
+
+ if not yes and not _prompt_yes_no("Proceed with teardown? [yes/N]: "):
+ print("Aborted by user.")
+ return 1
+
+ for t in targets:
+ try:
+ teardown_run(ec2_client, t, delete_key=delete_key)
+ except Exception as exc: # noqa: BLE001
+ # Design: deliberately broad. One run's teardown failure should
+ # not abort the rest of a --all batch; report it and keep
+ # going, then reflect the failure in the process exit code
+ # rather than swallowing it silently.
+ print(f"ERROR tearing down run {t.get('run_id', '?')}: {exc}", file=sys.stderr)
+ exit_code = 1
+ continue
+
+ run_id_key = t.get("run_id")
+ state_path = file_paths_by_run.get(run_id_key)
+ if state_path and state_path.exists():
+ state_path.unlink()
+ print(f"Removed state file {state_path}.")
+
+ return exit_code
+
+
+def main(argv: Optional[list] = None) -> int:
+ """Parse CLI args and tear down matching resources.
+
+ Parameters
+ ----------
+ argv : list[str] or None, default None
+ Argument list, as passed to ``argparse``'s ``parse_args``. ``None``
+ reads from ``sys.argv[1:]``.
+
+ Returns
+ -------
+ int
+ Process exit code. For ``--verify``: ``0`` if every audited region
+ is clean, ``1`` if tagged instances remain in any of them.
+ Otherwise: ``0`` on a fully successful teardown across every region
+ involved (or nothing to do anywhere), ``1`` if the user declined
+ confirmation in any region or if any individual run's
+ teardown/stale-cleanup raised.
+
+ Notes
+ -----
+ ``--tag-project`` is deliberately not a flag here (this CLI's flag set
+ is fixed by the work-package spec to
+ ``[--region] [--run-id | --all] [--verify] [--delete-key] [--yes]``):
+ the ``Project`` tag value used for discovery is always
+ ``Config().tag_project`` (i.e. ``Config``'s dataclass default,
+ ``"accelforge-correlation"``). A run provisioned with a *custom*
+ ``--tag-project`` cannot be found by this CLI and must instead be torn
+ down by calling :func:`find_tagged_instances`/:func:`teardown_run`
+ directly with that project value -- both are plain importable
+ functions for exactly this reason.
+
+ ``--region`` now defaults to ``None`` (Fix 4a), NOT ``"us-east-1"``: see
+ the "Region resolution" design comment above :func:`resolve_regions`
+ for the full rationale and priority order. This function may therefore
+ talk to *more than one* region in a single invocation (e.g. ``--all``
+ spanning every region this study's local state knows about, or
+ ``--verify`` auditing all of them) -- :func:`_client_for_region` keeps
+ one cached boto3 client per distinct region actually needed.
+ """
+ parser = argparse.ArgumentParser(
+ prog="teardown.py",
+ description="Tear down accelforge correlation-study AWS resources.",
+ )
+ parser.add_argument(
+ "--region",
+ type=str,
+ default=None,
+ help=(
+ "AWS region. If omitted, each targeted run's own state-file-recorded "
+ "region is used when known, else Config's default region (us-east-1); "
+ "an explicit --region always overrides both. See resolve_regions()."
+ ),
+ )
+ target = parser.add_mutually_exclusive_group()
+ target.add_argument("--run-id", type=str, default=None, help="Tear down only this run.")
+ target.add_argument(
+ "--all",
+ action="store_true",
+ help=(
+ "Tear down every accelforge-correlation-tagged run this study's local "
+ "state knows about, across every region those runs were provisioned in."
+ ),
+ )
+ parser.add_argument(
+ "--verify",
+ action="store_true",
+ help="Audit only: report any still-running tagged instances and exit non-zero if any remain.",
+ )
+ parser.add_argument(
+ "--delete-key",
+ action="store_true",
+ help="Also delete the AWS key pair and local PEM.",
+ )
+ parser.add_argument("--yes", action="store_true", help="Skip interactive confirmation.")
+ args = parser.parse_args(argv)
+
+ if not args.verify and not args.run_id and not args.all:
+ parser.error("one of --run-id, --all, or --verify is required")
+
+ _require_boto3()
+
+ # Design: instantiate a plain Config() rather than duplicating its
+ # tag_project/state_dir default logic here. See the "Notes" above on
+ # why --tag-project is not a teardown.py flag.
+ defaults = Config()
+ tag_project = defaults.tag_project
+ state_dir = defaults.state_dir
+
+ # Load EVERY locally known state file up front, regardless of scope --
+ # resolve_regions() needs the full set to group by region (Fix 4a);
+ # scope filtering (--run-id vs --all/--verify) happens inside
+ # resolve_regions() and the per-region loop below, not here.
+ all_states_by_run: Dict[str, dict] = {}
+ all_paths_by_run: Dict[str, Path] = {}
+ for path in _state_dir_files(state_dir, None):
+ try:
+ loaded = _load_state_file(path)
+ except (OSError, json.JSONDecodeError) as exc:
+ print(f"WARNING: could not read state file {path}: {exc}; skipping.")
+ continue
+ run_id_key = loaded.get("run_id", path.stem)
+ all_states_by_run[run_id_key] = loaded
+ all_paths_by_run[run_id_key] = path
+
+ region_map = resolve_regions(args, all_states_by_run)
+ clients: Dict[str, Any] = {}
+
+ if args.verify:
+ # Fix 4c: audit every region resolve_regions() knows about (not
+ # only the --region flag/fallback), so a run provisioned in a
+ # different region than the operator happens to be thinking about
+ # is not silently skipped by an audit meant to catch exactly that.
+ exit_code = 0
+ for region in sorted(region_map):
+ print(f"--- Verifying region {region} ---")
+ client = _client_for_region(clients, region)
+ if _verify(client, tag_project, args.run_id) != 0:
+ exit_code = 1
+ return exit_code
+
+ exit_code = 0
+ for region in sorted(region_map):
+ ec2_client = _client_for_region(clients, region)
+ run_ids_here = region_map[region]
+ file_states_by_run = {
+ rid: all_states_by_run[rid] for rid in run_ids_here if rid in all_states_by_run
+ }
+ file_paths_by_run = {
+ rid: all_paths_by_run[rid] for rid in run_ids_here if rid in all_paths_by_run
+ }
+ region_exit_code = _teardown_one_region(
+ ec2_client,
+ tag_project,
+ region,
+ args.run_id,
+ file_states_by_run,
+ file_paths_by_run,
+ args.delete_key,
+ args.yes,
+ )
+ if region_exit_code != 0:
+ exit_code = 1
+
+ return exit_code
+
+
+if __name__ == "__main__":
+ raise SystemExit(main())
diff --git a/tests/not_working/isl/mapper/__init__.py b/notebooks/astrasim2_correlation/correlation/tests/__init__.py
similarity index 100%
rename from tests/not_working/isl/mapper/__init__.py
rename to notebooks/astrasim2_correlation/correlation/tests/__init__.py
diff --git a/notebooks/astrasim2_correlation/correlation/tests/test_orchestrate.py b/notebooks/astrasim2_correlation/correlation/tests/test_orchestrate.py
new file mode 100644
index 00000000..79c73050
--- /dev/null
+++ b/notebooks/astrasim2_correlation/correlation/tests/test_orchestrate.py
@@ -0,0 +1,580 @@
+"""Tests for orchestrate.py.
+
+Like ``test_provision_teardown.py``, these tests never touch real AWS,
+never touch the network, and never spawn a real ``ssh``/``scp`` process.
+Where ``test_provision_teardown.py`` achieves that with
+``botocore.stub.Stubber`` (because it exercises ``provision.py``'s and
+``teardown.py``'s own boto3-request-building code directly), this module
+takes a different, coarser-grained approach: ``orchestrate.py`` mostly
+*sequences* those already-tested functions rather than building its own
+boto3 requests, so the provisioning functions themselves
+(``resolve_ami``, ``ensure_key_pair``, ``ensure_security_group``,
+``launch_instance``, ``wait_for_instance``), ``teardown.teardown_run``,
+``orchestrate.caller_ip``, ``orchestrate.boto3``, and
+``orchestrate.subprocess.run`` are all monkeypatched with lightweight
+fakes that record what they were called with. This exercises
+``orchestrate.py``'s own sequencing/argv-building logic -- the part this
+work package is actually responsible for -- without needing Stubber
+responses for calls this module never makes itself.
+
+Import strategy
+-----------------
+Mirrors ``test_provision_teardown.py`` exactly: ``correlation/`` has no
+``__init__.py`` (deliberately not a package), so it is inserted onto
+``sys.path`` explicitly before importing ``config``/``orchestrate``,
+rather than relying on pytest's own rootdir-insertion behavior.
+"""
+
+from __future__ import annotations
+
+import json
+import subprocess
+import sys
+from pathlib import Path
+from typing import Any, Dict, List
+
+import pytest
+
+# boto3 is not (and must not become) an accelforge package dependency; skip
+# this whole module rather than error if it is not installed in the
+# environment running the tests. orchestrate.py's own import of boto3 is
+# guarded the same way provision.py's/teardown.py's is, but the tests below
+# exercise real code paths that assume boto3 (and its stub-friendly
+# botocore internals) are present, matching test_provision_teardown.py's
+# identical importorskip.
+boto3 = pytest.importorskip("boto3")
+
+_CORRELATION_DIR = Path(__file__).resolve().parent.parent
+if str(_CORRELATION_DIR) not in sys.path:
+ sys.path.insert(0, str(_CORRELATION_DIR))
+
+import config # noqa: E402
+import orchestrate # noqa: E402
+
+Config = config.Config
+
+
+# ---------------------------------------------------------------------------
+# legs_for
+# ---------------------------------------------------------------------------
+
+
+def test_legs_for_both_returns_fc_then_torus():
+ """"both" expands to the FC leg followed by the torus leg, in that order."""
+ assert orchestrate.legs_for("both") == ["fc", "torus"]
+
+
+def test_legs_for_fc_returns_single_leg():
+ assert orchestrate.legs_for("fc") == ["fc"]
+
+
+def test_legs_for_torus_returns_single_leg():
+ assert orchestrate.legs_for("torus") == ["torus"]
+
+
+# ---------------------------------------------------------------------------
+# build_ssh_cmd / build_scp_cmd
+# ---------------------------------------------------------------------------
+
+
+def test_build_ssh_cmd_exact_argv(tmp_path):
+ """build_ssh_cmd's argv matches the exact option set/order the spec requires."""
+ key_path = tmp_path / "keys" / "run.pem"
+ known_hosts_path = tmp_path / ".state" / "known_hosts"
+
+ cmd = orchestrate.build_ssh_cmd(key_path, "ubuntu", "203.0.113.9", "echo hi", known_hosts_path)
+
+ assert cmd == [
+ "ssh",
+ "-i",
+ str(key_path),
+ "-o",
+ "StrictHostKeyChecking=accept-new",
+ "-o",
+ f"UserKnownHostsFile={known_hosts_path}",
+ "-o",
+ "ConnectTimeout=30",
+ "ubuntu@203.0.113.9",
+ "echo hi",
+ ]
+
+
+def test_build_scp_cmd_non_recursive_push_argv(tmp_path):
+ """A multi-source, non-recursive push builds argv with sources then dest, no -r."""
+ key_path = tmp_path / "keys" / "run.pem"
+ known_hosts_path = tmp_path / ".state" / "known_hosts"
+ sources = ["setup_node.sh", "run_profile.sh", "parse_nccl.py"]
+ dest = "ubuntu@203.0.113.9:~/"
+
+ cmd = orchestrate.build_scp_cmd(key_path, sources, dest, known_hosts_path)
+
+ assert cmd == [
+ "scp",
+ "-i",
+ str(key_path),
+ "-o",
+ "StrictHostKeyChecking=accept-new",
+ "-o",
+ f"UserKnownHostsFile={known_hosts_path}",
+ "-o",
+ "ConnectTimeout=30",
+ "setup_node.sh",
+ "run_profile.sh",
+ "parse_nccl.py",
+ "ubuntu@203.0.113.9:~/",
+ ]
+ assert "-r" not in cmd
+
+
+def test_build_scp_cmd_recursive_flag_is_first_positional_after_scp(tmp_path):
+ """recursive=True inserts -r immediately after the program name, before -i."""
+ key_path = tmp_path / "keys" / "run.pem"
+ known_hosts_path = tmp_path / ".state" / "known_hosts"
+
+ cmd = orchestrate.build_scp_cmd(
+ key_path,
+ ["ubuntu@203.0.113.9:~/results_fc"],
+ "/local/data/run/fc",
+ known_hosts_path,
+ recursive=True,
+ )
+
+ assert cmd[0] == "scp"
+ assert cmd[1] == "-r"
+ assert cmd[2] == "-i"
+ assert cmd[-2] == "ubuntu@203.0.113.9:~/results_fc"
+ assert cmd[-1] == "/local/data/run/fc"
+
+
+def test_build_scp_cmd_rejects_empty_sources(tmp_path):
+ with pytest.raises(ValueError):
+ orchestrate.build_scp_cmd(tmp_path / "k.pem", [], "dest", tmp_path / "known_hosts")
+
+
+# ---------------------------------------------------------------------------
+# End-to-end orchestration (main()), everything monkeypatched
+# ---------------------------------------------------------------------------
+
+
+class _FakeBoto3:
+ """Stand-in for the `boto3` module, only supplying `.client(...)`.
+
+ Design: monkeypatched onto `orchestrate.boto3` specifically (not the
+ real, globally-shared `boto3` module) so this fake never leaks into
+ any other module's view of boto3. Returns a plain sentinel string
+ rather than a real client, since every function orchestrate.py passes
+ that "client" to (resolve_ami, ensure_key_pair, ...) is itself
+ monkeypatched below and never actually calls a botocore method on it.
+ """
+
+ @staticmethod
+ def client(service_name: str, region_name: str = None):
+ return f"fake-{service_name}-client[{region_name}]"
+
+
+@pytest.fixture
+def orchestrate_fakes(tmp_path, monkeypatch):
+ """Monkeypatch every AWS/network/subprocess seam orchestrate.py has.
+
+ Returns
+ -------
+ dict
+ ``{"calls": list of (name, args) tuples recording every fake
+ invocation in order, "run_cmds": list of argv lists recorded by
+ the fake subprocess.run, "teardown_calls": list of state dicts
+ teardown_run was called with}``.
+
+ Notes
+ -----
+ Also points ``cfg``'s ``--key-dir``/``--state-dir`` and
+ ``orchestrate._DATA_DIR`` at ``tmp_path`` subdirectories (see
+ ``orchestrate._DATA_DIR``'s module docstring comment for why that
+ constant exists) so a full ``orchestrate.main(...)`` run in these
+ tests writes real state/CSV files only under pytest's ephemeral
+ ``tmp_path``, never into the real repository tree.
+ """
+ calls: List[tuple] = []
+ run_cmds: List[List[str]] = []
+ teardown_calls: List[Dict[str, Any]] = []
+
+ monkeypatch.setattr(orchestrate, "boto3", _FakeBoto3)
+ monkeypatch.setattr(orchestrate, "caller_ip", lambda: (_ for _ in ()).throw(
+ AssertionError("caller_ip() should never be called when --ssh-cidr is passed")
+ ))
+
+ def fake_resolve_ami(ssm_client, parameter):
+ calls.append(("resolve_ami", parameter))
+ return "ami-fake0123456789"
+
+ def fake_ensure_key_pair(ec2_client, key_name, key_dir):
+ calls.append(("ensure_key_pair", key_name))
+ key_dir.mkdir(parents=True, exist_ok=True)
+ return key_dir / f"{key_name}.pem"
+
+ def fake_ensure_security_group(ec2_client, group_name, ssh_cidr, tag_project, run_id):
+ calls.append(("ensure_security_group", group_name, ssh_cidr))
+ return "sg-fake0123456789"
+
+ def fake_launch_instance(ec2_client, cfg, ami_id, sg_id, key_name, dry_run=False):
+ calls.append(("launch_instance", dry_run))
+ if dry_run:
+ return {"instance_id": None, "purchasing_used": "ondemand"}
+ return {"instance_id": "i-fake0123456789", "purchasing_used": "spot"}
+
+ def fake_wait_for_instance(ec2_client, instance_id):
+ # Ordering assertion baked into the fake itself (rather than only
+ # checked after main() returns): the state file must already
+ # exist, with this instance_id recorded, by the time
+ # wait_for_instance is called -- this is the exact "state written
+ # before the SSH wait" ordering the work-package spec requires.
+ # See orchestrate.main's docstring "Design/WHY" note.
+ state_path = Path(_last_state_dir[0]) / f"{_last_run_id[0]}.json"
+ assert state_path.exists(), "state file must be written before wait_for_instance is called"
+ written = json.loads(state_path.read_text())
+ assert written["instance_id"] == instance_id
+ assert written["public_ip"] is None
+ calls.append(("wait_for_instance", instance_id))
+ return "203.0.113.9"
+
+ def fake_teardown_run(ec2_client, state, delete_key):
+ teardown_calls.append(dict(state))
+ calls.append(("teardown_run", state.get("instance_id"), delete_key))
+
+ def fake_run(cmd, check=True, **kwargs):
+ run_cmds.append(list(cmd))
+ return subprocess.CompletedProcess(cmd, 0)
+
+ # _last_state_dir / _last_run_id let fake_wait_for_instance locate the
+ # state file without needing main()'s local `cfg` in scope; populated
+ # by the test itself right before calling orchestrate.main(...).
+ _last_state_dir: List[Path] = [None]
+ _last_run_id: List[str] = [None]
+
+ monkeypatch.setattr(orchestrate, "resolve_ami", fake_resolve_ami)
+ monkeypatch.setattr(orchestrate, "ensure_key_pair", fake_ensure_key_pair)
+ monkeypatch.setattr(orchestrate, "ensure_security_group", fake_ensure_security_group)
+ monkeypatch.setattr(orchestrate, "launch_instance", fake_launch_instance)
+ monkeypatch.setattr(orchestrate, "wait_for_instance", fake_wait_for_instance)
+ monkeypatch.setattr(orchestrate, "teardown_run", fake_teardown_run)
+ monkeypatch.setattr(orchestrate.subprocess, "run", fake_run)
+
+ data_dir = tmp_path / "data"
+ monkeypatch.setattr(orchestrate, "_DATA_DIR", data_dir)
+
+ return {
+ "calls": calls,
+ "run_cmds": run_cmds,
+ "teardown_calls": teardown_calls,
+ "state_dir_holder": _last_state_dir,
+ "run_id_holder": _last_run_id,
+ "data_dir": data_dir,
+ }
+
+
+def _base_argv(tmp_path, run_id: str, extra: List[str] = None) -> List[str]:
+ """Shared CLI args for the end-to-end tests below.
+
+ Parameters
+ ----------
+ tmp_path : pathlib.Path
+ pytest's per-test temp directory; key-dir/state-dir are pointed
+ here so no test ever touches the real correlation/keys or
+ correlation/.state directories.
+ run_id : str
+ Deterministic run id so tests can locate the state file/data
+ directory by name instead of discovering a generated one.
+ extra : list[str] or None
+ Additional argv to append (e.g. ``["--keep-alive"]``).
+
+ Returns
+ -------
+ list[str]
+ argv suitable for ``orchestrate.main(...)``. Always includes
+ ``--ssh-cidr`` explicitly so ``caller_ip()`` (a real network call)
+ is never reached, and ``--yes`` so no interactive prompt blocks
+ the test.
+ """
+ argv = [
+ "--yes",
+ "--topology",
+ "both",
+ "--run-id",
+ run_id,
+ "--key-dir",
+ str(tmp_path / "keys"),
+ "--state-dir",
+ str(tmp_path / "state"),
+ "--ssh-cidr",
+ "203.0.113.5/32",
+ "--collectives",
+ "all_reduce,alltoall",
+ "--min-mib",
+ "1",
+ "--max-mib",
+ "2",
+ "--torus-dims",
+ "2x2x2",
+ ]
+ if extra:
+ argv += extra
+ return argv
+
+
+def test_main_end_to_end_happy_path(tmp_path, orchestrate_fakes, monkeypatch):
+ """Full main() run: state ordering, per-leg profiling, fetch, and teardown.
+
+ Asserts, per the work-package spec's test #3:
+ - the state file is written after launch (before wait_for_instance is
+ called -- enforced inside the fake_wait_for_instance itself) and
+ contains instance_id,
+ - the setup ssh command is executed with DEADMAN_MINUTES set,
+ - exactly one run_profile.sh invocation per leg, with the correct
+ dims ("8" then "2x2x2") and byte bounds (1 MiB / 2 MiB here),
+ - one scp fetch per leg,
+ - teardown_run is called exactly once, at the end,
+ - the state file is removed afterward.
+ """
+ run_id = "test-run-e2e"
+ state_dir = tmp_path / "state"
+ orchestrate_fakes["state_dir_holder"][0] = state_dir
+ orchestrate_fakes["run_id_holder"][0] = run_id
+
+ exit_code = orchestrate.main(_base_argv(tmp_path, run_id))
+
+ assert exit_code == 0
+
+ # --- state file lifecycle -------------------------------------------------
+ state_path = state_dir / f"{run_id}.json"
+ assert not state_path.exists(), "state file must be removed after a successful teardown"
+
+ # --- setup command ----------------------------------------------------------
+ run_cmds = orchestrate_fakes["run_cmds"]
+ setup_cmds = [c for c in run_cmds if "setup_node.sh" in c[-1]]
+ assert len(setup_cmds) == 1
+ assert "DEADMAN_MINUTES=120" in setup_cmds[0][-1]
+ assert setup_cmds[0][0] == "ssh"
+
+ # --- one run_profile.sh invocation per leg, correct dims/bytes -----------
+ profile_cmds = [c for c in run_cmds if "run_profile.sh" in c[-1]]
+ assert len(profile_cmds) == 2
+ fc_cmd, torus_cmd = profile_cmds[0][-1], profile_cmds[1][-1]
+ assert "results_fc fc 1048576 2097152 8 all_reduce alltoall" in fc_cmd
+ assert "results_torus torus 1048576 2097152 2x2x2 all_reduce alltoall" in torus_cmd
+
+ # --- fetch per leg ------------------------------------------------------
+ # Recursive scp fetch argv shape (per build_scp_cmd): [..., source, dest],
+ # so the source (a "user@ip:~/results_" string) is always the
+ # second-to-last element.
+ scp_fetch_cmds = [c for c in run_cmds if c[0] == "scp" and "-r" in c and "results_" in c[-2]]
+ fetch_sources = {c[-2] for c in scp_fetch_cmds}
+ assert any("results_fc" in s for s in fetch_sources)
+ assert any("results_torus" in s for s in fetch_sources)
+
+ # --- push commands happened before setup/profiling -----------------------
+ # A push argv is a flat list of local file-path elements followed by a
+ # remote dest string, so membership needs a substring scan across
+ # elements rather than an exact-element containment check (the pushed
+ # sources are full absolute paths, not the bare "setup_node.sh").
+ def _any_elem_contains(cmd: List[str], needle: str) -> bool:
+ return any(needle in elem for elem in cmd)
+
+ push_cmds = [c for c in run_cmds if c[0] == "scp" and _any_elem_contains(c, "setup_node.sh")]
+ assert len(push_cmds) == 1
+ torus_push_cmds = [c for c in run_cmds if c[0] == "scp" and "-r" in c and "torus_bench" in c[-1]]
+ assert len(torus_push_cmds) == 1
+
+ # --- ordering: push -> setup -> (profile -> fetch) x legs ----------------
+ def _first_index(predicate):
+ return next(i for i, c in enumerate(run_cmds) if predicate(c))
+
+ push_idx = _first_index(lambda c: c[0] == "scp" and _any_elem_contains(c, "setup_node.sh"))
+ setup_idx = _first_index(lambda c: c[0] == "ssh" and "setup_node.sh" in c[-1])
+ fc_profile_idx = _first_index(lambda c: c[0] == "ssh" and "run_profile.sh" in c[-1] and " fc " in c[-1])
+ assert push_idx < setup_idx < fc_profile_idx
+
+ # --- teardown called exactly once, at the very end -----------------------
+ assert len(orchestrate_fakes["teardown_calls"]) == 1
+ assert orchestrate_fakes["teardown_calls"][0]["instance_id"] == "i-fake0123456789"
+ call_names = [c[0] for c in orchestrate_fakes["calls"]]
+ assert call_names[-1] == "teardown_run"
+
+
+# ---------------------------------------------------------------------------
+# Fix 1 (BLOCKER): _run_leg must create the fetch's PARENT directory (not the
+# leaf) before scp runs, and must refuse to fetch into an already-existing
+# leaf directory.
+# ---------------------------------------------------------------------------
+
+
+def test_main_creates_run_data_dir_before_fetch_scp(tmp_path, orchestrate_fakes, monkeypatch):
+ """run_data_dir (the fetch's parent) exists by the time each leg's fetch scp runs.
+
+ Regression test for Fix 1: without ``run_data_dir.mkdir(...)`` before the
+ fetch, `scp -r user@host:~/results_fc /fc` would fail
+ outright the first time a run_id is used, since its parent directory
+ would not exist yet. This test captures os-path existence from *inside*
+ the fake ``subprocess.run`` at the exact moment a fetch command is
+ recorded, per the work-package spec's test #1.
+ """
+ run_id = "test-run-mkdir"
+ state_dir = tmp_path / "state"
+ orchestrate_fakes["state_dir_holder"][0] = state_dir
+ orchestrate_fakes["run_id_holder"][0] = run_id
+ data_dir = orchestrate_fakes["data_dir"]
+
+ run_data_dir_existed_at_fetch: List[bool] = []
+
+ def fake_run(cmd, check=True, **kwargs):
+ orchestrate_fakes["run_cmds"].append(list(cmd))
+ # A fetch is a recursive scp whose source (second-to-last argv
+ # element, per build_scp_cmd's argv shape) names a remote
+ # results_ path -- distinguishes it from the (also recursive)
+ # torus_bench/ push, whose source is a local path instead.
+ if cmd[0] == "scp" and "-r" in cmd and "results_" in cmd[-2]:
+ run_data_dir_existed_at_fetch.append((data_dir / run_id).exists())
+ return subprocess.CompletedProcess(cmd, 0)
+
+ monkeypatch.setattr(orchestrate.subprocess, "run", fake_run)
+
+ exit_code = orchestrate.main(_base_argv(tmp_path, run_id))
+
+ assert exit_code == 0
+ # One fetch per leg (fc, torus); run_data_dir must already exist at both.
+ assert len(run_data_dir_existed_at_fetch) == 2
+ assert all(run_data_dir_existed_at_fetch)
+
+
+def test_run_leg_raises_if_leaf_dir_already_exists(tmp_path, monkeypatch):
+ """_run_leg refuses to re-fetch into an already-existing run_data_dir/leg.
+
+ Direct unit test of :func:`orchestrate._run_leg` (rather than a full
+ ``main()`` run) per the work-package spec's test #2, exercising the
+ RuntimeError in isolation. ``_run_streaming`` is monkeypatched to a
+ recording no-op so no real ssh/scp subprocess is ever attempted; the
+ remote profiling ssh command runs (it happens before the leaf-dir check
+ in ``_run_leg``'s body), but the fetch scp must never be reached.
+ """
+ run_cmds: List[List[str]] = []
+ monkeypatch.setattr(orchestrate, "_run_streaming", lambda cmd: run_cmds.append(cmd))
+
+ cfg = Config(run_id="test-run-leaf-exists", torus_dims=(2, 2, 2), collectives="all_reduce")
+ run_data_dir = tmp_path / "data" / cfg.run_id
+ local_leg_dir = run_data_dir / "fc"
+ local_leg_dir.mkdir(parents=True)
+
+ with pytest.raises(RuntimeError, match="already exists"):
+ orchestrate._run_leg(
+ cfg,
+ "fc",
+ tmp_path / "key.pem",
+ "203.0.113.9",
+ tmp_path / "known_hosts",
+ run_data_dir,
+ )
+
+ # Only the remote profiling ssh command (which precedes the leaf-dir
+ # check in _run_leg's body) ran; the fetch scp was never attempted.
+ assert len(run_cmds) == 1
+ assert run_cmds[0][0] == "ssh"
+
+
+def test_main_dry_run_stops_before_launch_and_ssh(tmp_path, orchestrate_fakes):
+ """--dry-run creates key+SG+a dry launch, but never waits/pushes/profiles/tears down."""
+ run_id = "test-run-dry"
+ state_dir = tmp_path / "state"
+ orchestrate_fakes["state_dir_holder"][0] = state_dir
+ orchestrate_fakes["run_id_holder"][0] = run_id
+
+ exit_code = orchestrate.main(_base_argv(tmp_path, run_id, extra=["--dry-run"]))
+
+ assert exit_code == 0
+ call_names = [c[0] for c in orchestrate_fakes["calls"]]
+ assert "ensure_key_pair" in call_names
+ assert "ensure_security_group" in call_names
+ assert "launch_instance" in call_names
+ assert "wait_for_instance" not in call_names
+ assert "teardown_run" not in call_names
+ assert orchestrate_fakes["run_cmds"] == []
+ assert not (state_dir / f"{run_id}.json").exists()
+
+
+def test_main_profiling_failure_still_tears_down_and_propagates(tmp_path, orchestrate_fakes, monkeypatch):
+ """A failed profiling subprocess still triggers teardown, and the error propagates."""
+ run_id = "test-run-fail"
+ state_dir = tmp_path / "state"
+ orchestrate_fakes["state_dir_holder"][0] = state_dir
+ orchestrate_fakes["run_id_holder"][0] = run_id
+
+ def failing_run(cmd, check=True, **kwargs):
+ orchestrate_fakes["run_cmds"].append(list(cmd))
+ if "run_profile.sh" in cmd[-1]:
+ raise subprocess.CalledProcessError(returncode=1, cmd=cmd)
+ return subprocess.CompletedProcess(cmd, 0)
+
+ monkeypatch.setattr(orchestrate.subprocess, "run", failing_run)
+
+ with pytest.raises(subprocess.CalledProcessError):
+ orchestrate.main(_base_argv(tmp_path, run_id))
+
+ # Teardown must still have run despite the propagating exception.
+ assert len(orchestrate_fakes["teardown_calls"]) == 1
+ assert orchestrate_fakes["teardown_calls"][0]["instance_id"] == "i-fake0123456789"
+ # And the state file was still cleaned up by that successful teardown.
+ assert not (state_dir / f"{run_id}.json").exists()
+
+
+def test_main_teardown_failure_does_not_mask_original_exception(tmp_path, orchestrate_fakes, monkeypatch):
+ """If teardown_run ALSO raises, the original profiling exception still propagates."""
+ run_id = "test-run-double-fail"
+ state_dir = tmp_path / "state"
+ orchestrate_fakes["state_dir_holder"][0] = state_dir
+ orchestrate_fakes["run_id_holder"][0] = run_id
+
+ def failing_run(cmd, check=True, **kwargs):
+ orchestrate_fakes["run_cmds"].append(list(cmd))
+ if "run_profile.sh" in cmd[-1]:
+ raise subprocess.CalledProcessError(returncode=1, cmd=cmd)
+ return subprocess.CompletedProcess(cmd, 0)
+
+ def failing_teardown(ec2_client, state, delete_key):
+ orchestrate_fakes["teardown_calls"].append(dict(state))
+ raise RuntimeError("simulated teardown failure (e.g. DependencyViolation exhausted)")
+
+ monkeypatch.setattr(orchestrate.subprocess, "run", failing_run)
+ monkeypatch.setattr(orchestrate, "teardown_run", failing_teardown)
+
+ # The ORIGINAL exception (CalledProcessError from profiling) must be
+ # what propagates, not the teardown's RuntimeError -- this is the
+ # "don't mask the original exception" behavior _teardown_and_cleanup
+ # documents.
+ with pytest.raises(subprocess.CalledProcessError):
+ orchestrate.main(_base_argv(tmp_path, run_id))
+
+ assert len(orchestrate_fakes["teardown_calls"]) == 1
+ # Teardown failed, so the state file must NOT have been removed --
+ # teardown.py --run-id needs it to find leftover resources later.
+ assert (state_dir / f"{run_id}.json").exists()
+
+
+def test_main_keep_alive_skips_teardown_and_keeps_state_file(tmp_path, orchestrate_fakes, capsys):
+ """--keep-alive leaves the instance up: no teardown_run call, state file remains."""
+ run_id = "test-run-keep-alive"
+ state_dir = tmp_path / "state"
+ orchestrate_fakes["state_dir_holder"][0] = state_dir
+ orchestrate_fakes["run_id_holder"][0] = run_id
+
+ exit_code = orchestrate.main(_base_argv(tmp_path, run_id, extra=["--keep-alive"]))
+
+ assert exit_code == 0
+ assert orchestrate_fakes["teardown_calls"] == []
+ call_names = [c[0] for c in orchestrate_fakes["calls"]]
+ assert "teardown_run" not in call_names
+
+ state_path = state_dir / f"{run_id}.json"
+ assert state_path.exists()
+ written = json.loads(state_path.read_text())
+ assert written["public_ip"] == "203.0.113.9"
+
+ out = capsys.readouterr().out
+ assert "STILL BEING BILLED" in out
+ assert "ssh -i" in out
diff --git a/notebooks/astrasim2_correlation/correlation/tests/test_parse_nccl.py b/notebooks/astrasim2_correlation/correlation/tests/test_parse_nccl.py
new file mode 100644
index 00000000..322e682b
--- /dev/null
+++ b/notebooks/astrasim2_correlation/correlation/tests/test_parse_nccl.py
@@ -0,0 +1,288 @@
+"""Tests for parse_nccl.py: the nccl-tests / torus_bench log parser.
+
+Fixtures under ``tests/fixtures/`` hold representative raw stdout captured
+from the two profiling tools (see the ``correlation`` package's WP2 spec
+for the exact log formats). These tests exercise the parsing functions
+directly (:func:`parse_nccl.parse_nccl_tests`,
+:func:`parse_nccl.parse_torus_bench`), the CSV writer
+(:func:`parse_nccl.rows_to_csv`), and the CLI entry point end-to-end via
+:mod:`subprocess`, using ``sys.executable`` so the tests run under whatever
+interpreter is running pytest itself (matching how ``run_profile.sh``
+invokes this script with a plain ``python3``).
+"""
+
+from __future__ import annotations
+
+import csv
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+
+from parse_nccl import (
+ UNIFIED_CSV_FIELDNAMES,
+ parse_nccl_tests,
+ parse_torus_bench,
+ rows_to_csv,
+)
+
+# Design: resolve fixture/script paths relative to this test file (not the
+# CWD) so the suite passes regardless of where pytest is invoked from, per
+# the spec's instruction to load fixtures with pathlib relative to the test
+# file.
+TESTS_DIR = Path(__file__).resolve().parent
+FIXTURES_DIR = TESTS_DIR / "fixtures"
+PARSE_NCCL_SCRIPT = TESTS_DIR.parent / "parse_nccl.py"
+
+FC_ALL_REDUCE_LOG = FIXTURES_DIR / "fc_all_reduce.log"
+FC_ALLTOALL_LOG = FIXTURES_DIR / "fc_alltoall.log"
+TORUS_ALL_REDUCE_LOG = FIXTURES_DIR / "torus_all_reduce.log"
+
+
+def test_parse_nccl_tests_all_reduce():
+ """fc_all_reduce.log parses to 2 rows with the expected first-row values.
+
+ Exercises the common case: a well-formed all_reduce_perf log with a
+ standard 13-token data row (size, count, type, redop, root, then the
+ trailing-8 out-of-place/in-place metrics).
+ """
+ text = FC_ALL_REDUCE_LOG.read_text(encoding="utf-8")
+ rows = parse_nccl_tests(text)
+
+ assert len(rows) == 2
+ first = rows[0]
+ assert first["size_bytes"] == 1048576
+ assert first["count"] == 262144
+ assert first["dtype"] == "float"
+ assert first["time_us"] == pytest.approx(98.52)
+ assert first["algbw_GBps"] == pytest.approx(10.64)
+ assert first["busbw_GBps"] == pytest.approx(18.62)
+ assert first["wrong"] == "0"
+
+
+def test_parse_nccl_tests_alltoall_na_wrong():
+ """fc_alltoall.log parses to 2 rows and exercises the 'N/A' #wrong path.
+
+ alltoall_perf prints redop="none" and root="-1" instead of a real
+ reduction op/root -- this test confirms the trailing-8-token rule
+ parses those rows correctly regardless, and that an "N/A" out-of-place
+ #wrong value (validation disabled for that data point) is preserved
+ as the literal string "N/A" rather than raising or being coerced to a
+ number.
+ """
+ text = FC_ALLTOALL_LOG.read_text(encoding="utf-8")
+ rows = parse_nccl_tests(text)
+
+ assert len(rows) == 2
+ assert rows[0]["wrong"] == "N/A"
+ assert rows[0]["size_bytes"] == 1048576
+ assert rows[0]["time_us"] == pytest.approx(120.44)
+ # Second row is a normal (non-N/A) row, confirming N/A handling on row 0
+ # didn't leak into subsequent parsing.
+ assert rows[1]["wrong"] == "0"
+ assert rows[1]["size_bytes"] == 2097152
+
+
+def test_parse_torus_bench_all_reduce():
+ """torus_all_reduce.log parses to 3 rows with the expected first row.
+
+ Exercises the comma-delimited TORUSBENCH sentinel format and the
+ check-field-to-wrong-string remapping (check "1" -> wrong "0").
+ """
+ text = TORUS_ALL_REDUCE_LOG.read_text(encoding="utf-8")
+ rows = parse_torus_bench(text)
+
+ assert len(rows) == 3
+ first = rows[0]
+ assert first["collective"] == "all_reduce"
+ assert first["dims"] == "2x2x2"
+ assert first["size_bytes"] == 1048576
+ assert first["time_us"] == pytest.approx(142.11)
+ assert first["wrong"] == "0"
+
+
+def test_rows_to_csv_round_trip(tmp_path):
+ """rows_to_csv() writes the exact unified header, in order, and empty
+ algbw/busbw cells for torus rows round-trip as empty strings.
+
+ Builds one nccl-tests-shaped row and one torus_bench-shaped row (as
+ main() would produce them) and confirms the on-disk CSV, when read back
+ with csv.DictReader, has fieldnames matching UNIFIED_CSV_FIELDNAMES
+ exactly (order included) and that the torus row's bandwidth columns are
+ empty rather than "None" or some other stand-in.
+ """
+ rows = [
+ {
+ "source": "nccl-tests",
+ "topology": "fc",
+ "dims": "8",
+ "collective": "all_reduce",
+ "size_bytes": 1048576,
+ "count": 262144,
+ "dtype": "float",
+ "time_us": 98.52,
+ "algbw_GBps": 10.64,
+ "busbw_GBps": 18.62,
+ "wrong": "0",
+ },
+ {
+ "source": "torus_bench",
+ "topology": "torus",
+ "dims": "2x2x2",
+ "collective": "all_reduce",
+ "size_bytes": 1048576,
+ "count": "",
+ "dtype": "",
+ "time_us": 142.11,
+ "algbw_GBps": "",
+ "busbw_GBps": "",
+ "wrong": "0",
+ },
+ ]
+ out_path = tmp_path / "unified.csv"
+
+ rows_to_csv(rows, out_path)
+
+ with out_path.open(newline="", encoding="utf-8") as f:
+ reader = csv.DictReader(f)
+ assert reader.fieldnames == UNIFIED_CSV_FIELDNAMES
+ read_rows = list(reader)
+
+ assert len(read_rows) == 2
+ assert read_rows[1]["source"] == "torus_bench"
+ assert read_rows[1]["algbw_GBps"] == ""
+ assert read_rows[1]["busbw_GBps"] == ""
+
+
+def _run_cli(*args: str) -> subprocess.CompletedProcess:
+ """Invoke parse_nccl.py's CLI as a subprocess.
+
+ Parameters
+ ----------
+ *args : str
+ Arguments to pass after the script path, e.g. the raw log path and
+ ``--source``/``--out``/etc. flags.
+
+ Returns
+ -------
+ subprocess.CompletedProcess
+ Result of the invocation, with stdout/stderr captured as text.
+
+ Notes
+ -----
+ Design: uses ``sys.executable`` (not a hard-coded ``python3``) so the
+ subprocess runs under the exact interpreter executing the test suite,
+ matching pytest's own environment rather than risking a PATH mismatch.
+ """
+ return subprocess.run(
+ [sys.executable, str(PARSE_NCCL_SCRIPT), *args],
+ capture_output=True,
+ text=True,
+ )
+
+
+def test_cli_nccl_tests_end_to_end(tmp_path):
+ """CLI parses an nccl-tests log to a CSV with the expected row count."""
+ out_path = tmp_path / "fc_all_reduce.csv"
+ result = _run_cli(
+ str(FC_ALL_REDUCE_LOG),
+ "--source",
+ "nccl-tests",
+ "--collective",
+ "all_reduce",
+ "--topology",
+ "fc",
+ "--out",
+ str(out_path),
+ )
+
+ assert result.returncode == 0, result.stderr
+ assert out_path.exists()
+
+ with out_path.open(newline="", encoding="utf-8") as f:
+ read_rows = list(csv.DictReader(f))
+ assert len(read_rows) == 2
+ assert all(row["source"] == "nccl-tests" for row in read_rows)
+ assert all(row["dims"] == "8" for row in read_rows)
+
+
+def test_cli_torus_bench_end_to_end(tmp_path):
+ """CLI parses a torus_bench log to a CSV with the expected row count."""
+ out_path = tmp_path / "torus_all_reduce.csv"
+ result = _run_cli(
+ str(TORUS_ALL_REDUCE_LOG),
+ "--source",
+ "torus_bench",
+ "--collective",
+ "all_reduce",
+ "--topology",
+ "torus",
+ "--dims",
+ "2x2x2",
+ "--out",
+ str(out_path),
+ )
+
+ assert result.returncode == 0, result.stderr
+ assert out_path.exists()
+
+ with out_path.open(newline="", encoding="utf-8") as f:
+ read_rows = list(csv.DictReader(f))
+ assert len(read_rows) == 3
+ assert all(row["source"] == "torus_bench" for row in read_rows)
+ assert all(row["algbw_GBps"] == "" for row in read_rows)
+
+
+def test_cli_torus_bench_collective_mismatch_errors(tmp_path):
+ """CLI exits non-zero when --collective disagrees with the log content.
+
+ Bonus coverage beyond the spec's 6 mandated cases: confirms the
+ validate-CLI-against-sentinel behavior documented in parse_nccl.main()
+ actually triggers a hard failure rather than silently mislabeling data,
+ since a silent mismatch here would corrupt the correlation notebook's
+ inputs without any visible signal.
+ """
+ out_path = tmp_path / "should_not_be_created.csv"
+ result = _run_cli(
+ str(TORUS_ALL_REDUCE_LOG),
+ "--source",
+ "torus_bench",
+ "--collective",
+ "all_gather", # deliberately wrong; fixture is all_reduce
+ "--topology",
+ "torus",
+ "--out",
+ str(out_path),
+ )
+
+ assert result.returncode != 0
+ assert not out_path.exists()
+
+
+@pytest.mark.parametrize(
+ "log_text",
+ [
+ "# just a comment\n\n# another comment\n",
+ "hello world\n",
+ "# nThread 1 nGpus 8\nhello world\n# trailing comment\n",
+ ],
+)
+def test_parse_nccl_tests_skips_malformed_lines(log_text):
+ """Comment-only, blank, and garbage lines are skipped without raising."""
+ assert parse_nccl_tests(log_text) == []
+
+
+@pytest.mark.parametrize(
+ "log_text",
+ [
+ "# torus_bench collective=all_reduce dims=2x2x2\n",
+ "hello world\n",
+ "# header\nhello world\nTORUSBENCH,incomplete,field\n",
+ ],
+)
+def test_parse_torus_bench_skips_malformed_lines(log_text):
+ """Comment-only, blank, garbage, and malformed sentinel lines are
+ skipped without raising (including a TORUSBENCH, line with too few
+ comma-separated fields)."""
+ assert parse_torus_bench(log_text) == []
diff --git a/notebooks/astrasim2_correlation/correlation/tests/test_provision_teardown.py b/notebooks/astrasim2_correlation/correlation/tests/test_provision_teardown.py
new file mode 100644
index 00000000..0e8caad2
--- /dev/null
+++ b/notebooks/astrasim2_correlation/correlation/tests/test_provision_teardown.py
@@ -0,0 +1,802 @@
+"""Tests for config.py, provision.py, and teardown.py.
+
+These tests never touch real AWS: every boto3 client call is intercepted
+by ``botocore.stub.Stubber``, which asserts on the exact request
+parameters and returns a canned response (or raises a canned
+``ClientError``) instead of making a network call. This lets the tests
+exercise real botocore request-building and error-handling code paths --
+including exact parameter shape assertions -- with zero network access and
+zero AWS credentials, matching this work package's hard constraint that
+nothing may call AWS.
+
+Import strategy
+-----------------
+``config.py``/``provision.py``/``teardown.py`` live directly under
+``correlation/`` (one level above this ``tests/`` package), and
+``correlation/`` is deliberately *not* a Python package (no
+``__init__.py`` there -- see the work-package spec). ``provision.py`` and
+``teardown.py`` both do a plain ``from config import Config``, which only
+resolves if ``correlation/`` is on ``sys.path``. We therefore insert that
+directory onto ``sys.path`` explicitly before importing any of the three
+modules under test, rather than relying on pytest's own rootdir-insertion
+behavior (which happens to also achieve this here, but only for the
+specific "prepend" import mode and package-marker layout currently in
+place -- an explicit insert is robust to either changing).
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+import pytest
+
+# boto3 is not (and must not become) an accelforge package dependency; skip
+# this whole module rather than error if it is not installed in the
+# environment running the tests.
+boto3 = pytest.importorskip("boto3")
+
+from botocore.exceptions import ClientError # noqa: E402 (after importorskip)
+from botocore.stub import Stubber # noqa: E402
+
+_CORRELATION_DIR = Path(__file__).resolve().parent.parent
+if str(_CORRELATION_DIR) not in sys.path:
+ sys.path.insert(0, str(_CORRELATION_DIR))
+
+import config # noqa: E402
+import provision # noqa: E402
+import teardown # noqa: E402
+
+Config = config.Config
+
+
+def _make_client(service_name: str):
+ """Build a boto3 client with dummy credentials, safe for Stubber use.
+
+ Parameters
+ ----------
+ service_name : str
+ E.g. ``"ec2"`` or ``"ssm"``.
+
+ Returns
+ -------
+ botocore.client.BaseClient
+ A real boto3 client object, never used to make a real network
+ call in these tests (every call site below is wrapped in a
+ ``Stubber`` context).
+
+ Notes
+ -----
+ Design: ``Stubber`` intercepts the HTTP send step, but botocore still
+ runs its normal request-signing step first, which raises
+ ``NoCredentialsError`` if no credentials are configured anywhere
+ (env vars, profile, instance metadata, ...). This test environment
+ intentionally has none, so every client is built with harmless dummy
+ static credentials purely to satisfy the signer -- these are never
+ sent anywhere, since Stubber never performs a real HTTP request.
+ """
+ return boto3.client(
+ service_name,
+ region_name="us-east-1",
+ aws_access_key_id="testing",
+ aws_secret_access_key="testing",
+ )
+
+
+# ---------------------------------------------------------------------------
+# resolve_ami
+# ---------------------------------------------------------------------------
+
+
+def test_resolve_ami_returns_stubbed_parameter_value():
+ """resolve_ami extracts Parameter.Value from the SSM response."""
+ ssm_client = _make_client("ssm")
+ stubber = Stubber(ssm_client)
+ parameter_name = (
+ "/aws/service/deeplearning/ami/x86_64/"
+ "base-oss-nvidia-driver-gpu-ubuntu-22.04/latest/ami-id"
+ )
+ stubber.add_response(
+ "get_parameter",
+ {
+ "Parameter": {
+ "Name": parameter_name,
+ "Value": "ami-0123456789abcdef0",
+ "Type": "String",
+ }
+ },
+ {"Name": parameter_name},
+ )
+
+ with stubber:
+ result = provision.resolve_ami(ssm_client, parameter_name)
+
+ stubber.assert_no_pending_responses()
+ assert result == "ami-0123456789abcdef0"
+
+
+# ---------------------------------------------------------------------------
+# launch_instance
+# ---------------------------------------------------------------------------
+
+
+def test_launch_instance_spot_then_ondemand_falls_back_on_capacity_error():
+ """A spot InsufficientInstanceCapacity error triggers an on-demand retry.
+
+ Also asserts (per the work-package spec) that the first, failing
+ request was a spot request -- i.e. it carried InstanceMarketOptions --
+ both directly (inspecting the kwargs dict) and indirectly (via
+ Stubber's expected_params, which would fail the test if
+ launch_instance's real spot request didn't match).
+
+ Design (Fix 10): the FIRST call's expected_params is an independently
+ hardcoded literal dict, not derived from
+ ``provision._build_run_instances_kwargs`` -- deriving it from the same
+ helper the code under test calls would make this test circular (a bug
+ in that helper's request-shape would go undetected, since the test's
+ expectation and the code's actual request would drift together). The
+ literal below pins the exact request shape against the work-package
+ spec instead of against the code's own helper. The second (on-demand
+ fallback) call keeps using the helper-derived ``ondemand_kwargs`` for
+ convenience, since its shape isn't the focus of this particular test.
+ """
+ ec2_client = _make_client("ec2")
+ stubber = Stubber(ec2_client)
+
+ cfg = Config(purchasing="spot-then-ondemand", run_id="test-run")
+ ami_id = "ami-0123456789abcdef0"
+ sg_id = "sg-0123456789abcdef0"
+ key_name = "accelforge-correlation-test-run"
+
+ # Independently hardcoded, per the work-package spec's exact field list
+ # -- see the docstring above for why this must NOT be derived from
+ # provision._build_run_instances_kwargs.
+ spot_kwargs_literal = {
+ "ImageId": ami_id,
+ "InstanceType": "p5.48xlarge",
+ "KeyName": key_name,
+ "SecurityGroupIds": [sg_id],
+ "MinCount": 1,
+ "MaxCount": 1,
+ "InstanceInitiatedShutdownBehavior": "terminate",
+ "BlockDeviceMappings": [
+ {
+ "DeviceName": "/dev/sda1",
+ "Ebs": {
+ "VolumeSize": 200,
+ "VolumeType": "gp3",
+ "DeleteOnTermination": True,
+ },
+ }
+ ],
+ "TagSpecifications": [
+ {
+ "ResourceType": "instance",
+ "Tags": [
+ {"Key": "Project", "Value": "accelforge-correlation"},
+ {"Key": "RunId", "Value": "test-run"},
+ {"Key": "Name", "Value": "accelforge-correlation-test-run"},
+ ],
+ },
+ {
+ "ResourceType": "volume",
+ "Tags": [
+ {"Key": "Project", "Value": "accelforge-correlation"},
+ {"Key": "RunId", "Value": "test-run"},
+ {"Key": "Name", "Value": "accelforge-correlation-test-run"},
+ ],
+ },
+ ],
+ "DryRun": False,
+ "InstanceMarketOptions": {
+ "MarketType": "spot",
+ "SpotOptions": {
+ "SpotInstanceType": "one-time",
+ "InstanceInterruptionBehavior": "terminate",
+ },
+ },
+ }
+ # Kept helper-derived for the fallback call, per the spec ("if convenient").
+ ondemand_kwargs = provision._build_run_instances_kwargs(
+ cfg, ami_id, sg_id, key_name, dry_run=False, use_spot=False
+ )
+
+ # The spec's explicit ask: the FIRST request must have carried
+ # InstanceMarketOptions (spot), the fallback must not.
+ assert "InstanceMarketOptions" in spot_kwargs_literal
+ assert "InstanceMarketOptions" not in ondemand_kwargs
+
+ stubber.add_client_error(
+ "run_instances",
+ service_error_code="InsufficientInstanceCapacity",
+ service_message="There is no Spot capacity available.",
+ expected_params=spot_kwargs_literal,
+ )
+ stubber.add_response(
+ "run_instances",
+ {"Instances": [{"InstanceId": "i-0123456789abcdef0"}]},
+ expected_params=ondemand_kwargs,
+ )
+
+ with stubber:
+ result = provision.launch_instance(
+ ec2_client, cfg, ami_id, sg_id, key_name, dry_run=False
+ )
+
+ stubber.assert_no_pending_responses()
+ assert result == {
+ "instance_id": "i-0123456789abcdef0",
+ "purchasing_used": "ondemand",
+ }
+
+
+def test_launch_instance_spot_only_does_not_fall_back():
+ """purchasing="spot" propagates the ClientError instead of retrying on-demand."""
+ ec2_client = _make_client("ec2")
+ stubber = Stubber(ec2_client)
+
+ cfg = Config(purchasing="spot", run_id="test-run")
+ ami_id = "ami-0123456789abcdef0"
+ sg_id = "sg-0123456789abcdef0"
+ key_name = "accelforge-correlation-test-run"
+
+ spot_kwargs = provision._build_run_instances_kwargs(
+ cfg, ami_id, sg_id, key_name, dry_run=False, use_spot=True
+ )
+ # Only one response is ever queued: if launch_instance incorrectly
+ # attempted a second (fallback) call, Stubber itself would raise for
+ # having no more queued responses, which is not a ClientError -- so
+ # pytest.raises(ClientError) below would fail loudly in that case too.
+ stubber.add_client_error(
+ "run_instances",
+ service_error_code="InsufficientInstanceCapacity",
+ service_message="There is no Spot capacity available.",
+ expected_params=spot_kwargs,
+ )
+
+ with stubber:
+ with pytest.raises(ClientError):
+ provision.launch_instance(
+ ec2_client, cfg, ami_id, sg_id, key_name, dry_run=False
+ )
+
+ stubber.assert_no_pending_responses()
+
+
+def test_launch_instance_dry_run_success_does_not_raise(capsys):
+ """A DryRunOperation error is treated as a successful authorization check."""
+ ec2_client = _make_client("ec2")
+ stubber = Stubber(ec2_client)
+
+ cfg = Config(purchasing="ondemand", run_id="test-run")
+ ami_id = "ami-0123456789abcdef0"
+ sg_id = "sg-0123456789abcdef0"
+ key_name = "accelforge-correlation-test-run"
+
+ ondemand_kwargs = provision._build_run_instances_kwargs(
+ cfg, ami_id, sg_id, key_name, dry_run=True, use_spot=False
+ )
+ stubber.add_client_error(
+ "run_instances",
+ service_error_code="DryRunOperation",
+ service_message="Request would have succeeded, but DryRun flag is set.",
+ expected_params=ondemand_kwargs,
+ )
+
+ with stubber:
+ result = provision.launch_instance(
+ ec2_client, cfg, ami_id, sg_id, key_name, dry_run=True
+ )
+
+ stubber.assert_no_pending_responses()
+ assert result["purchasing_used"] == "ondemand"
+ # No instance was actually created during a dry run.
+ assert result["instance_id"] is None
+ assert "dry-run OK" in capsys.readouterr().out
+
+
+# ---------------------------------------------------------------------------
+# ensure_security_group
+# ---------------------------------------------------------------------------
+
+
+def test_ensure_security_group_authorizes_requested_cidr():
+ """ensure_security_group wires the given ssh_cidr into the ingress rule."""
+ ec2_client = _make_client("ec2")
+ stubber = Stubber(ec2_client)
+
+ vpc_id = "vpc-0123456789abcdef0"
+ sg_id = "sg-0123456789abcdef0"
+ ssh_cidr = "203.0.113.5/32"
+ group_name = "accelforge-correlation-test-run-sg"
+ run_id = "test-run"
+ tag_project = "accelforge-correlation"
+
+ stubber.add_response(
+ "describe_vpcs",
+ {"Vpcs": [{"VpcId": vpc_id, "IsDefault": True}]},
+ {"Filters": [{"Name": "isDefault", "Values": ["true"]}]},
+ )
+ stubber.add_response(
+ "create_security_group",
+ {"GroupId": sg_id},
+ {
+ "GroupName": group_name,
+ "Description": f"accelforge correlation study SG for run {run_id}",
+ "VpcId": vpc_id,
+ },
+ )
+ # The parameter that matters most here: expected_params pins the exact
+ # CidrIp ensure_security_group must send, so the test fails loudly if
+ # the wrong CIDR (or the wrong port) were ever authorized.
+ stubber.add_response(
+ "authorize_security_group_ingress",
+ {},
+ {
+ "GroupId": sg_id,
+ "IpPermissions": [
+ {
+ "IpProtocol": "tcp",
+ "FromPort": 22,
+ "ToPort": 22,
+ "IpRanges": [
+ {
+ "CidrIp": ssh_cidr,
+ "Description": "SSH access for accelforge correlation study",
+ }
+ ],
+ }
+ ],
+ },
+ )
+ stubber.add_response(
+ "create_tags",
+ {},
+ {
+ "Resources": [sg_id],
+ "Tags": [
+ {"Key": "Project", "Value": tag_project},
+ {"Key": "RunId", "Value": run_id},
+ {"Key": "Name", "Value": group_name},
+ ],
+ },
+ )
+
+ with stubber:
+ result = provision.ensure_security_group(
+ ec2_client, group_name, ssh_cidr, tag_project, run_id
+ )
+
+ stubber.assert_no_pending_responses()
+ assert result == sg_id
+
+
+# ---------------------------------------------------------------------------
+# teardown.find_tagged_instances
+# ---------------------------------------------------------------------------
+
+
+def test_find_tagged_instances_returns_ids_from_both_reservations():
+ """find_tagged_instances flattens across multiple Reservations entries."""
+ ec2_client = _make_client("ec2")
+ stubber = Stubber(ec2_client)
+
+ tag_project = "accelforge-correlation"
+ expected_filters = {
+ "Filters": [
+ {"Name": "tag:Project", "Values": [tag_project]},
+ {
+ "Name": "instance-state-name",
+ "Values": ["pending", "running", "stopping", "stopped"],
+ },
+ ]
+ }
+ stubber.add_response(
+ "describe_instances",
+ {
+ "Reservations": [
+ {
+ "Instances": [
+ {
+ "InstanceId": "i-aaaa000000000001",
+ "State": {"Name": "running"},
+ "Tags": [{"Key": "RunId", "Value": "run-a"}],
+ }
+ ]
+ },
+ {
+ "Instances": [
+ {
+ "InstanceId": "i-bbbb000000000002",
+ "State": {"Name": "pending"},
+ "Tags": [{"Key": "RunId", "Value": "run-b"}],
+ }
+ ]
+ },
+ ]
+ },
+ expected_filters,
+ )
+
+ with stubber:
+ result = teardown.find_tagged_instances(ec2_client, tag_project)
+
+ stubber.assert_no_pending_responses()
+ ids = {instance["InstanceId"] for instance in result}
+ assert ids == {"i-aaaa000000000001", "i-bbbb000000000002"}
+
+
+# ---------------------------------------------------------------------------
+# teardown security-group delete retry
+# ---------------------------------------------------------------------------
+
+
+def test_delete_security_group_retries_past_dependency_violation(monkeypatch):
+ """A DependencyViolation is retried (not fatal) and eventually succeeds."""
+ ec2_client = _make_client("ec2")
+ stubber = Stubber(ec2_client)
+ sg_id = "sg-0123456789abcdef0"
+
+ stubber.add_client_error(
+ "delete_security_group",
+ service_error_code="DependencyViolation",
+ service_message="resource sg-0123456789abcdef0 has a dependent object",
+ expected_params={"GroupId": sg_id},
+ )
+ stubber.add_response("delete_security_group", {}, {"GroupId": sg_id})
+
+ sleep_calls = []
+ # Patch time.sleep as seen through teardown's own `import time`, so the
+ # retry loop does not actually block the test suite for
+ # _SG_DELETE_RETRY_SLEEP_S seconds.
+ monkeypatch.setattr(teardown.time, "sleep", lambda seconds: sleep_calls.append(seconds))
+
+ with stubber:
+ teardown._delete_security_group_with_retry(ec2_client, sg_id)
+
+ stubber.assert_no_pending_responses()
+ assert len(sleep_calls) == 1
+
+
+# ---------------------------------------------------------------------------
+# Fix 3 (MAJOR): caller_ip() fails fast instead of fail-open
+# ---------------------------------------------------------------------------
+
+
+def test_caller_ip_raises_on_discovery_failure(monkeypatch):
+ """caller_ip() raises RuntimeError (mentioning --ssh-cidr) on any discovery failure.
+
+ Regression test for Fix 3: the old behavior returned the sentinel
+ "0.0.0.0" on failure, which every caller turned into the CIDR
+ "0.0.0.0/32" -- unreachable by anyone, including the operator -- only
+ after real AWS resources already existed and were already billing.
+ Failing fast here means the error surfaces before any of that happens.
+ """
+
+ def fake_urlopen(*args, **kwargs):
+ raise provision.urllib.error.URLError("simulated DNS failure")
+
+ monkeypatch.setattr(provision.urllib.request, "urlopen", fake_urlopen)
+
+ with pytest.raises(RuntimeError, match="--ssh-cidr"):
+ provision.caller_ip()
+
+
+# ---------------------------------------------------------------------------
+# Fix 2 (MAJOR): provision.main() closes the orphan-instance window
+# ---------------------------------------------------------------------------
+
+
+def test_provision_main_writes_state_before_wait_and_reports_recovery_on_failure(
+ tmp_path, monkeypatch, capsys
+):
+ """provision.main() writes state (public_ip=None) before wait_for_instance runs,
+ and on a wait_for_instance failure prints a loud recovery block (instance
+ id, STILL RUNNING AND BILLING, state file path, exact recovery command)
+ and re-raises rather than silently losing track of a running instance.
+
+ Every AWS-touching seam provision.main() has (resolve_ami,
+ ensure_key_pair, ensure_security_group, launch_instance,
+ wait_for_instance, and boto3.client itself) is monkeypatched with a
+ lightweight fake, mirroring test_orchestrate.py's ``orchestrate_fakes``
+ approach -- this exercises provision.main()'s own sequencing/state-file
+ logic (what Fix 2 changed) without touching real AWS or needing a
+ Stubber response sequence for calls this test never lets happen for
+ real.
+ """
+
+ class _FakeBoto3:
+ @staticmethod
+ def client(service_name, region_name=None):
+ return f"fake-{service_name}-client[{region_name}]"
+
+ monkeypatch.setattr(provision, "boto3", _FakeBoto3)
+ monkeypatch.setattr(provision, "resolve_ami", lambda ssm, param: "ami-fake0123456789")
+
+ def fake_ensure_key_pair(ec2, key_name, key_dir):
+ key_dir.mkdir(parents=True, exist_ok=True)
+ return key_dir / f"{key_name}.pem"
+
+ monkeypatch.setattr(provision, "ensure_key_pair", fake_ensure_key_pair)
+ monkeypatch.setattr(
+ provision,
+ "ensure_security_group",
+ lambda ec2, name, cidr, tag_project, run_id: "sg-fake0123456789",
+ )
+ monkeypatch.setattr(
+ provision,
+ "launch_instance",
+ lambda ec2, cfg, ami_id, sg_id, key_name, dry_run=False: {
+ "instance_id": "i-fake0123456789",
+ "purchasing_used": "spot",
+ },
+ )
+
+ state_dir = tmp_path / "state"
+ run_id = "test-run-orphan-window"
+
+ def failing_wait_for_instance(ec2, instance_id):
+ # By the time wait_for_instance is called, state must already be on
+ # disk with public_ip still the None placeholder -- exactly the
+ # ordering Fix 2 requires (state written BEFORE the SSH wait, not
+ # only after it succeeds).
+ state_path = state_dir / f"{run_id}.json"
+ assert state_path.exists(), "state file must exist before wait_for_instance is called"
+ written = json.loads(state_path.read_text())
+ assert written["instance_id"] == instance_id
+ assert written["public_ip"] is None
+ raise TimeoutError("simulated SSH-reachability timeout")
+
+ monkeypatch.setattr(provision, "wait_for_instance", failing_wait_for_instance)
+
+ argv = [
+ "--yes",
+ "--run-id",
+ run_id,
+ "--key-dir",
+ str(tmp_path / "keys"),
+ "--state-dir",
+ str(state_dir),
+ "--ssh-cidr",
+ "203.0.113.5/32",
+ ]
+
+ with pytest.raises(TimeoutError):
+ provision.main(argv)
+
+ # The state file must survive the failure -- teardown.py needs it to
+ # find and tear down the still-running instance later.
+ state_path = state_dir / f"{run_id}.json"
+ assert state_path.exists()
+
+ err = capsys.readouterr().err
+ assert "i-fake0123456789" in err
+ assert "STILL RUNNING AND BILLING" in err
+ assert str(state_path) in err
+ assert f"python teardown.py --run-id {run_id} --region" in err
+
+
+# ---------------------------------------------------------------------------
+# Fix 4a (MAJOR): teardown.resolve_regions -- pure region-resolution logic
+# ---------------------------------------------------------------------------
+#
+# These tests make no AWS calls and use no Stubber, per resolve_regions()'s
+# own design (a pure function of `args` and a plain dict of loaded state
+# files) -- `args` is a minimal argparse.Namespace stand-in, not a fully
+# parsed CLI invocation, since resolve_regions() only ever consults
+# `.region` and `.run_id`.
+
+
+def test_resolve_regions_explicit_flag_wins_over_state_region():
+ """An explicit --region always overrides a state file's own recorded region."""
+ args = argparse.Namespace(region="us-west-2", run_id="run-a")
+ states = {"run-a": {"region": "eu-central-1"}}
+
+ assert teardown.resolve_regions(args, states) == {"us-west-2": ["run-a"]}
+
+
+def test_resolve_regions_uses_state_file_region_when_no_flag():
+ """With no --region, a single --run-id's own recorded region is used."""
+ args = argparse.Namespace(region=None, run_id="run-a")
+ states = {"run-a": {"region": "ap-southeast-2"}}
+
+ assert teardown.resolve_regions(args, states) == {"ap-southeast-2": ["run-a"]}
+
+
+def test_resolve_regions_falls_back_to_config_default_when_unknown():
+ """With no --region and no matching state file, Config's default region is used."""
+ args = argparse.Namespace(region=None, run_id="run-with-no-state-file")
+ states: dict = {}
+
+ assert teardown.resolve_regions(args, states) == {
+ Config().region: ["run-with-no-state-file"]
+ }
+
+
+def test_resolve_regions_all_groups_by_recorded_region_and_keeps_default():
+ """--all (run_id=None) groups every known run by its own region, plus the default."""
+ args = argparse.Namespace(region=None, run_id=None)
+ states = {
+ "run-a": {"region": "us-west-2"},
+ "run-b": {"region": "us-west-2"},
+ "run-c": {"region": "eu-central-1"},
+ "run-d": {}, # no recorded region at all -> falls back to the default
+ }
+
+ result = teardown.resolve_regions(args, states)
+
+ assert set(result["us-west-2"]) == {"run-a", "run-b"}
+ assert result["eu-central-1"] == ["run-c"]
+ # The default/fallback region (Config's own default) is always present
+ # as a key, even though only run-d actually landed there via fallback,
+ # so a caller iterating this mapping's keys always still searches it
+ # for tag-only discovery of state-less instances.
+ assert "run-d" in result[Config().region]
+
+
+def test_resolve_regions_all_explicit_flag_overrides_every_run():
+ """An explicit --region with --all overrides every individual run's own region."""
+ args = argparse.Namespace(region="us-east-2", run_id=None)
+ states = {
+ "run-a": {"region": "us-west-2"},
+ "run-b": {"region": "eu-central-1"},
+ }
+
+ result = teardown.resolve_regions(args, states)
+
+ assert result == {"us-east-2": ["run-a", "run-b"]}
+
+
+# ---------------------------------------------------------------------------
+# Fix 4b (MAJOR): a stale state file (no matching live instance) is routed
+# through teardown_run -- SG/key-pair cleanup included -- BEFORE its local
+# state file is removed, instead of just being unlinked.
+# ---------------------------------------------------------------------------
+
+
+def test_teardown_one_region_stale_state_calls_delete_security_group_before_removing_file(
+ tmp_path,
+):
+ """A stale run's security group is deleted before its state file disappears.
+
+ Regression test for Fix 4b: the pre-fix behavior just unlinked a stale
+ state file without ever touching AWS, leaking the security group (and,
+ with --delete-key, the key pair) of any run whose instance died out
+ from under it -- e.g. via the on-instance dead-man timer firing -- before
+ teardown.py was ever run. Stubber's queued ``delete_security_group``
+ response is only consumed if ``_teardown_one_region``'s internal
+ ``teardown_run`` call actually invokes it (``stubber.assert_no_pending_
+ responses()`` below fails the test otherwise); combined with the state
+ file only being removed from disk AFTER that call returns without
+ raising, these two assertions together establish the required "SG
+ deleted before state file disappears" ordering.
+ """
+ ec2_client = _make_client("ec2")
+ stubber = Stubber(ec2_client)
+
+ run_id = "stale-run"
+ instance_id = "i-0123456789abcdef0"
+ sg_id = "sg-0123456789abcdef0"
+ tag_project = "accelforge-correlation"
+
+ state = {
+ "run_id": run_id,
+ "region": "us-east-1",
+ "instance_id": instance_id,
+ "sg_id": sg_id,
+ "key_name": "accelforge-correlation-stale-run",
+ "key_path": None,
+ }
+ state_path = tmp_path / f"{run_id}.json"
+ state_path.write_text(json.dumps(state))
+
+ # Discovery: no live instance matches this run -- makes it "stale".
+ stubber.add_response(
+ "describe_instances",
+ {"Reservations": []},
+ {
+ "Filters": [
+ {"Name": "tag:Project", "Values": [tag_project]},
+ {
+ "Name": "instance-state-name",
+ "Values": ["pending", "running", "stopping", "stopped"],
+ },
+ ]
+ },
+ )
+ # teardown_run's terminate step: the instance is already gone (e.g. the
+ # dead-man timer fired) -- tolerated as InvalidInstanceID.NotFound, per
+ # teardown_run's own existing contract.
+ stubber.add_client_error(
+ "terminate_instances",
+ service_error_code="InvalidInstanceID.NotFound",
+ service_message=f"The instance ID '{instance_id}' does not exist",
+ expected_params={"InstanceIds": [instance_id]},
+ )
+ # The assertion this test exists for: delete_security_group MUST be
+ # called (Stubber raises on assert_no_pending_responses() otherwise).
+ stubber.add_response("delete_security_group", {}, {"GroupId": sg_id})
+
+ with stubber:
+ exit_code = teardown._teardown_one_region(
+ ec2_client,
+ tag_project,
+ "us-east-1",
+ None,
+ {run_id: state},
+ {run_id: state_path},
+ delete_key=False,
+ yes=True,
+ )
+
+ stubber.assert_no_pending_responses()
+ assert exit_code == 0
+ # The state file is gone only now that teardown_run (including
+ # delete_security_group) has already completed successfully.
+ assert not state_path.exists()
+
+
+def test_teardown_one_region_stale_state_keeps_file_if_teardown_run_raises(tmp_path, monkeypatch):
+ """If the stale-cleanup teardown_run call itself raises, the state file survives.
+
+ Complements the happy-path stale-cleanup test above: a failed cleanup
+ must not lose track of the leak by removing the state file anyway.
+ """
+ ec2_client = _make_client("ec2")
+ stubber = Stubber(ec2_client)
+
+ run_id = "stale-run-cleanup-fails"
+ sg_id = "sg-0123456789abcdef0"
+ tag_project = "accelforge-correlation"
+
+ # No instance_id: teardown_run skips straight to security-group
+ # deletion, which is the call we make fail here.
+ state = {"run_id": run_id, "region": "us-east-1", "sg_id": sg_id, "key_name": None}
+ state_path = tmp_path / f"{run_id}.json"
+ state_path.write_text(json.dumps(state))
+
+ # Avoid actually sleeping between retries (see
+ # test_delete_security_group_retries_past_dependency_violation's
+ # identical rationale for patching teardown's own `import time`).
+ monkeypatch.setattr(teardown.time, "sleep", lambda seconds: None)
+
+ stubber.add_response(
+ "describe_instances",
+ {"Reservations": []},
+ {
+ "Filters": [
+ {"Name": "tag:Project", "Values": [tag_project]},
+ {
+ "Name": "instance-state-name",
+ "Values": ["pending", "running", "stopping", "stopped"],
+ },
+ ]
+ },
+ )
+ # A persistent DependencyViolation exhausts _delete_security_group_with_retry's
+ # retries and surfaces as a RuntimeError from teardown_run.
+ for _ in range(teardown._SG_DELETE_MAX_RETRIES):
+ stubber.add_client_error(
+ "delete_security_group",
+ service_error_code="DependencyViolation",
+ service_message="resource has a dependent object",
+ expected_params={"GroupId": sg_id},
+ )
+
+ with stubber:
+ exit_code = teardown._teardown_one_region(
+ ec2_client,
+ tag_project,
+ "us-east-1",
+ None,
+ {run_id: state},
+ {run_id: state_path},
+ delete_key=False,
+ yes=True,
+ )
+
+ stubber.assert_no_pending_responses()
+ assert exit_code == 1
+ assert state_path.exists()
diff --git a/notebooks/astrasim2_correlation/correlation/torus_bench/Makefile b/notebooks/astrasim2_correlation/correlation/torus_bench/Makefile
new file mode 100644
index 00000000..a53fd392
--- /dev/null
+++ b/notebooks/astrasim2_correlation/correlation/torus_bench/Makefile
@@ -0,0 +1,27 @@
+CUDA_HOME ?= /usr/local/cuda
+NCCL_HOME ?=
+NVCC ?= $(CUDA_HOME)/bin/nvcc
+GXX ?= g++
+ARCH ?= -arch=sm_90
+
+.PHONY: clean
+
+NCCL_INC := $(if $(NCCL_HOME),-I$(NCCL_HOME)/include)
+NCCL_LIB := $(if $(NCCL_HOME),-L$(NCCL_HOME)/lib)
+
+# Default target (first in this file): the GPU/NCCL executor. Requires nvcc plus a CUDA
+# toolkit and NCCL headers/libs; NOT buildable in this development environment (no nvcc here
+# -- see the work-package report). Provided so the GPU target compiles cleanly, by construction,
+# on the target AWS p5.48xlarge instance, where CUDA_HOME/NCCL_HOME should be set as needed.
+torus_bench: torus_bench.cu
+ $(NVCC) -O3 $(ARCH) $(NCCL_INC) $(NCCL_LIB) -o $@ $< -lnccl
+
+# Simulator target: forces g++ to treat the .cu file as plain C++ (-x c++), not CUDA source.
+# This is the local, GPU-free test vehicle for the schedule-building, edge-assertion, and
+# per-collective reduction/routing logic -- see the file header of torus_bench.cu for the full
+# schedule-as-data rationale. This is the target exercised by the acceptance criteria.
+torus_bench_sim: torus_bench.cu
+ $(GXX) -x c++ -std=c++17 -O2 -Wall -Wextra -DTORUS_SIM -o $@ $<
+
+clean:
+ rm -f torus_bench torus_bench_sim
diff --git a/notebooks/astrasim2_correlation/correlation/torus_bench/torus_bench.cu b/notebooks/astrasim2_correlation/correlation/torus_bench/torus_bench.cu
new file mode 100644
index 00000000..ed17c0b9
--- /dev/null
+++ b/notebooks/astrasim2_correlation/correlation/torus_bench/torus_bench.cu
@@ -0,0 +1,1681 @@
+// torus_bench.cu
+//
+// Single-process multi-GPU NCCL benchmark that runs collective algorithms restricted to the
+// EDGES OF A LOGICAL TORUS, on hardware that is physically fully-connected (one AWS
+// p5.48xlarge, 8x H100 over NVSwitch).
+//
+// SCIENTIFIC PURPOSE (drives the whole design)
+// ---------------------------------------------
+// This benchmark exists to correlate an analytical torus-network model against real
+// measurements. The measurement is only valid if EVERY inter-GPU transfer travels between
+// logical torus NEIGHBORS -- that constraint IS the experiment. A single misrouted transfer
+// (e.g. a "shortcut" the NVSwitch fabric would happily allow but the torus topology would not)
+// silently invalidates the correlation. We therefore do not trust ourselves to hand-write
+// per-collective CUDA/NCCL call sequences and eyeball their correctness; instead:
+//
+// SCHEDULE-AS-DATA: pure host code (no CUDA, no GPU) builds an explicit, fully materialized
+// step-by-step transfer schedule (`Schedule` = vector, each Step a set of concurrent
+// `Xfer`s followed by local reduce/copy `LocalOp`s). A single choke point -- the edge
+// assertion inside build_schedule() -- inspects every Xfer the schedule will ever contain and
+// aborts the program if any of them is not a torus-neighbor transfer. This assertion is the
+// scientific guarantee of this file and must never be disabled, in either build.
+//
+// TWO EXECUTORS, ONE SCHEDULE: the identical Schedule produced by build_schedule() is handed
+// to either of two interchangeable executors with matching function signatures:
+// - a host-memory SIMULATOR (compiled here, in this environment, with no GPU/nvcc
+// available -- this is the local test vehicle and is exercised by the acceptance
+// criteria), and
+// - an NCCL/CUDA executor (built on the target GPU instance; cannot be compiled or run in
+// this environment since no nvcc/CUDA toolkit is installed here -- see the project report
+// for confirmation of this constraint).
+// Both executors are driven by the exact same schedule-building and CLI code; only the
+// "how do I actually move these bytes" implementation differs, selected at compile time via
+// the TORUS_SIM macro. Compiling with `g++ -x c++ -DTORUS_SIM` yields the simulator binary;
+// compiling with `nvcc` (TORUS_SIM undefined) yields the GPU binary. All CUDA/NCCL-specific
+// code is fenced with `#ifndef TORUS_SIM` so a plain C++ compiler never sees CUDA syntax.
+//
+// TOPOLOGY CONVENTION
+// --------------------
+// dims = [d_0, ..., d_{K-1}] is a K-dimensional torus with product(dims) = N ranks. Rank <->
+// coordinate mapping is row-major with the LAST dimension fastest-varying (i.e. like a C array
+// of shape `dims`). neighbor(r, dim, +-1) wraps around (mod dims[dim]). For an extent-2
+// dimension, +1 and -1 land on the SAME neighbor -- the code below computes this generically via
+// modular arithmetic and never special-cases extent-2 dimensions.
+//
+// DOCUMENTATION CONVENTIONS USED IN THIS FILE
+// ---------------------------------------------
+// Each function has a comment block with: a one-line summary, Parameters, Returns, and (where
+// relevant) Preconditions/Notes -- the C++ analogue of NumPy-style docstrings. Each per-collective
+// schedule-builder additionally documents, phase by phase, the data-placement INVARIANT that
+// phase establishes; this is the load-bearing correctness argument for that collective and is
+// exactly what a reader needs to convince themselves the algorithm is right.
+//
+// DERIVATION NOTE ON reduce_scatter (see build_reduce_scatter() below for full detail): the
+// planning spec for this work package proposed a tentative per-step slot formula for
+// reduce_scatter and then explicitly flagged uncertainty about it ("hold on, mirror the
+// all_gather invariant exactly"). That tentative formula, worked through by hand and confirmed
+// with a throwaway Python simulation across dims=[8],[2,4],[4,2],[2,2,2], places each fully
+// reduced slot one hop short of its destination rank. The corrected formula and direction
+// (documented at build_reduce_scatter) were derived from a time-reversal argument against the
+// (spec-verified-correct) all_gather formula and independently confirmed by brute-force
+// simulation before being encoded here; the C++ simulator's --check flag re-verifies this at
+// runtime for every size in the sweep.
+
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+
+// =====================================================================================
+// Buffer identifiers and collective enum
+// =====================================================================================
+
+// Logical per-rank buffer roles used by every schedule builder and both executors. Both
+// executors allocate storage per (rank, BufId) according to buffer_sizes() below.
+enum BufId : int {
+ BUF_SEND = 0,
+ BUF_RECV = 1,
+ BUF_TMP = 2,
+ BUF_WORK_A = 3,
+ BUF_WORK_B = 4,
+ NUM_BUFS = 5
+};
+
+enum class Collective { ALL_REDUCE, ALL_GATHER, REDUCE_SCATTER, ALLTOALL, BROADCAST, SENDRECV };
+
+// Returns the canonical CLI/output-format name for a collective.
+//
+// Parameters
+// ----------
+// c : Collective
+//
+// Returns
+// -------
+// const char* -- a string literal, e.g. "all_reduce". Never null.
+const char* collective_name(Collective c) {
+ switch (c) {
+ case Collective::ALL_REDUCE: return "all_reduce";
+ case Collective::ALL_GATHER: return "all_gather";
+ case Collective::REDUCE_SCATTER: return "reduce_scatter";
+ case Collective::ALLTOALL: return "alltoall";
+ case Collective::BROADCAST: return "broadcast";
+ case Collective::SENDRECV: return "sendrecv";
+ }
+ return "?";
+}
+
+// Parses a --collective CLI argument.
+//
+// Parameters
+// ----------
+// s : const std::string& -- one of the six recognized collective names.
+// out : Collective& -- set on success; left unmodified on failure.
+//
+// Returns
+// -------
+// bool -- true if `s` was recognized.
+bool parse_collective(const std::string& s, Collective& out) {
+ if (s == "all_reduce") { out = Collective::ALL_REDUCE; return true; }
+ if (s == "all_gather") { out = Collective::ALL_GATHER; return true; }
+ if (s == "reduce_scatter") { out = Collective::REDUCE_SCATTER; return true; }
+ if (s == "alltoall") { out = Collective::ALLTOALL; return true; }
+ if (s == "broadcast") { out = Collective::BROADCAST; return true; }
+ if (s == "sendrecv") { out = Collective::SENDRECV; return true; }
+ return false;
+}
+
+// Returns a human-readable name for a BufId, used only in --check diagnostic output.
+const char* buf_name(int b) {
+ switch (b) {
+ case BUF_SEND: return "BUF_SEND";
+ case BUF_RECV: return "BUF_RECV";
+ case BUF_TMP: return "BUF_TMP";
+ case BUF_WORK_A: return "BUF_WORK_A";
+ case BUF_WORK_B: return "BUF_WORK_B";
+ }
+ return "?";
+}
+
+// =====================================================================================
+// Schedule IR (pure data -- no CUDA dependency anywhere in this section)
+// =====================================================================================
+
+// One point-to-point transfer between two DIFFERENT ranks. Every Xfer that will ever be
+// constructed by any builder in this file must satisfy is_torus_neighbor(src, dst) -- see
+// check_edge_or_abort() and the final validation pass in build_schedule(). src == dst is
+// forbidden by construction: same-rank data movement must be expressed as a LocalOp instead.
+struct Xfer {
+ int src, dst;
+ int src_buf, dst_buf;
+ size_t src_off, dst_off, bytes;
+};
+
+// A same-rank operation: either a byte-for-byte copy (add == false) or an elementwise
+// float accumulate dst[i] += src[i] (add == true), applied over `bytes` bytes (i.e.
+// bytes/sizeof(float) floats -- bytes must be a multiple of 4 whenever add == true).
+struct LocalOp {
+ int rank;
+ int src_buf;
+ size_t src_off;
+ int dst_buf;
+ size_t dst_off;
+ size_t bytes;
+ bool add;
+};
+
+// One synchronization step of a Schedule: all `xfers` are considered to execute concurrently
+// (reading pre-step state), and only once every Xfer in the step has landed do the `post`
+// LocalOps run (also concurrently with each other, since by construction no two post ops of
+// the same step touch overlapping (rank,buf,offset) ranges).
+struct Step {
+ std::vector xfers;
+ std::vector post;
+};
+
+using Schedule = std::vector;
+
+// =====================================================================================
+// Torus topology helpers
+// =====================================================================================
+
+// Returns N = product(dims), the total rank count of the torus.
+inline int num_ranks(const std::vector& dims) {
+ int n = 1;
+ for (int d : dims) n *= d;
+ return n;
+}
+
+// Converts a rank id to its torus coordinates.
+//
+// Parameters
+// ----------
+// r : int -- rank id, 0 <= r < product(dims).
+// dims : const std::vector& -- torus extents, dims[K-1] is the fastest-varying axis.
+//
+// Returns
+// -------
+// std::vector -- coordinates, one per dimension, coords[j] in [0, dims[j]).
+inline std::vector coords_of(int r, const std::vector& dims) {
+ std::vector c(dims.size());
+ for (int j = (int)dims.size() - 1; j >= 0; --j) {
+ c[j] = r % dims[j];
+ r /= dims[j];
+ }
+ return c;
+}
+
+// Inverse of coords_of(): converts torus coordinates back to a rank id (row-major, last
+// dimension fastest).
+inline int rank_of(const std::vector& c, const std::vector& dims) {
+ int r = 0;
+ for (size_t j = 0; j < dims.size(); ++j) r = r * dims[j] + c[j];
+ return r;
+}
+
+// Returns the rank reached from `r` by moving one hop of `delta` (+1 or -1) along dimension
+// `dim`, wrapping around (mod dims[dim]). For an extent-2 dimension, delta=+1 and delta=-1
+// necessarily return the same rank -- this falls out of the modular arithmetic below with no
+// special-casing, matching the spec's requirement that extent-2 degeneracy not be hard-coded.
+inline int neighbor(int r, int dim, int delta, const std::vector& dims) {
+ std::vector c = coords_of(r, dims);
+ int e = dims[dim];
+ c[dim] = ((c[dim] + delta) % e + e) % e;
+ return rank_of(c, dims);
+}
+
+// Determines whether two ranks are torus neighbors: their coordinates must differ in exactly
+// one dimension, and in that dimension by +-1 modulo the dimension's extent.
+//
+// Notes
+// -----
+// For an extent-2 dimension this is trivially satisfied by any pair that differs there (both
+// possible non-zero differences, 1 and (extent-1)=1, coincide), which is the intended behavior.
+inline bool is_torus_neighbor(int a, int b, const std::vector& dims) {
+ std::vector ca = coords_of(a, dims), cb = coords_of(b, dims);
+ int diff_dim = -1, diff_count = 0;
+ for (size_t j = 0; j < dims.size(); ++j) {
+ if (ca[j] != cb[j]) { diff_dim = (int)j; ++diff_count; }
+ }
+ if (diff_count != 1) return false;
+ int e = dims[diff_dim];
+ int d = ((ca[diff_dim] - cb[diff_dim]) % e + e) % e;
+ return d == 1 || d == e - 1;
+}
+
+// Aborts the program with a file:line diagnostic if (src,dst) is not a torus edge, or if
+// src == dst. This is THE scientific guarantee of this benchmark (see file header) and must
+// remain active in both builds.
+//
+// Design decision: we use an explicit check + std::abort() rather than assert() from
+// , because assert() compiles to a no-op under -DNDEBUG and we do not control every
+// build environment this file might eventually be compiled in (e.g. a release-mode CI flag).
+// An explicit check is unconditionally active regardless of optimization/NDEBUG flags.
+inline void check_edge_or_abort(int src, int dst, const std::vector& dims) {
+ if (src == dst) {
+ std::fprintf(stderr,
+ "%s:%d: EDGE ASSERTION FAILED: Xfer has src==dst (rank %d); same-rank "
+ "movement must be expressed as a LocalOp, not an Xfer\n",
+ __FILE__, __LINE__, src);
+ std::abort();
+ }
+ if (!is_torus_neighbor(src, dst, dims)) {
+ std::fprintf(stderr,
+ "%s:%d: EDGE ASSERTION FAILED: rank %d -> rank %d is not a torus "
+ "neighbor for the given dims; this transfer would not exist on the "
+ "logical torus and must not be scheduled\n",
+ __FILE__, __LINE__, src, dst);
+ std::abort();
+ }
+}
+
+// Enumerates every combination of coordinate values across dimensions [0, d), invoking `cb`
+// once per combination with a coordinate vector `v` whose entries at indices >= d are left
+// exactly as passed in (typically already pinned to a specific rank's own coordinates).
+//
+// Parameters
+// ----------
+// d : int -- number of leading dimensions (0..d-1) to enumerate freely; if d == 0, `cb` is
+// invoked exactly once, with `v` unchanged (the "no free dimensions" case).
+// dims : const std::vector& -- torus extents.
+// v : std::vector -- base coordinate vector (taken by value since we mutate indices < d
+// during enumeration; entries at indices >= d are the caller's fixed values).
+// cb : const std::function&)>& -- invoked once per combination.
+//
+// Notes
+// -----
+// Enumeration order (dimension 0 varies slowest, in this implementation) is an arbitrary
+// choice: every caller in this file treats each combination as an independent, order-agnostic
+// unit of work (one Xfer/LocalOp pair per combination), so the traversal order used to reach
+// the same combination SET has no effect on correctness.
+inline void for_each_free_combo(int d, const std::vector& dims, std::vector v,
+ const std::function&)>& cb) {
+ std::function rec = [&](int j) {
+ if (j == d) { cb(v); return; }
+ for (int val = 0; val < dims[j]; ++val) {
+ v[j] = val;
+ rec(j + 1);
+ }
+ };
+ rec(0);
+}
+
+// =====================================================================================
+// Per-(collective,S) buffer sizing
+// =====================================================================================
+
+// Computes the per-rank byte size of each of the five logical buffers a given collective needs
+// for total-collective-size S. Both executors call this identically to allocate storage (heap
+// arrays for the simulator, cudaMalloc for the NCCL build) -- see file header for buffer roles.
+//
+// Parameters
+// ----------
+// c : Collective
+// dims : const std::vector& -- torus extents; N = product(dims).
+// S : size_t -- total collective size in bytes for this sweep point.
+//
+// Returns
+// -------
+// std::array -- indexed by BufId; unused buffers are size 0.
+//
+// Preconditions
+// -------------
+// S must be divisible by N and by 4*N*N (whole-float shard/chunk boundaries); callers are
+// expected to have already applied the sweep-level divisibility skip check (see main()) before
+// calling this.
+inline std::array buffer_sizes(Collective c, const std::vector& dims,
+ size_t S) {
+ int N = num_ranks(dims);
+ size_t m = S / (size_t)N; // per-rank shard, m = S/N, per the CLI spec's convention.
+ int K = (int)dims.size();
+ std::array sz{};
+ sz.fill(0);
+ switch (c) {
+ case Collective::SENDRECV:
+ sz[BUF_SEND] = m;
+ sz[BUF_RECV] = (size_t)K * m; // one m-byte region per dimension.
+ break;
+ case Collective::BROADCAST:
+ sz[BUF_SEND] = m; // only meaningful on root; allocated uniformly for simplicity.
+ sz[BUF_RECV] = m;
+ break;
+ case Collective::ALL_GATHER:
+ sz[BUF_SEND] = m;
+ sz[BUF_RECV] = (size_t)N * m; // N slots of m bytes each.
+ break;
+ case Collective::REDUCE_SCATTER:
+ sz[BUF_SEND] = (size_t)N * m; // N slots of m bytes each (destined for each rank).
+ sz[BUF_RECV] = m;
+ sz[BUF_WORK_A] = (size_t)N * m; // running accumulator, one slot per destination.
+ sz[BUF_TMP] = (size_t)N * m; // staging area for incoming adds.
+ break;
+ case Collective::ALL_REDUCE:
+ // Composed internally of reduce_scatter(shard=m/N) followed by all_gather
+ // (shard=m/N); see build_all_reduce() for the full derivation. Both sub-phases'
+ // internal buffer needs (N*(m/N) == m) collapse to a uniform m bytes here.
+ sz[BUF_SEND] = m;
+ sz[BUF_RECV] = m;
+ sz[BUF_WORK_A] = m;
+ sz[BUF_TMP] = m;
+ break;
+ case Collective::ALLTOALL:
+ sz[BUF_SEND] = m; // N chunks of c=m/N bytes each.
+ sz[BUF_RECV] = m;
+ sz[BUF_WORK_A] = 2 * m; // ping-pong in-flight staging; sized generously (2x) per
+ sz[BUF_WORK_B] = 2 * m; // spec, occupancy is asserted at schedule-build time.
+ break;
+ }
+ return sz;
+}
+
+// =====================================================================================
+// Schedule builders (pure host code -- no CUDA dependency anywhere in this section)
+// =====================================================================================
+
+// Builds the schedule for `sendrecv`: one step per torus dimension, every rank exchanges its
+// full BUF_SEND vector with its +1 neighbor in that dimension.
+//
+// Invariant per step d: after step d, every rank's BUF_RECV region [d*m, (d+1)*m) holds the
+// data that neighbor(r, d, -1) sent -- i.e. its BUF_SEND contents -- since neighbor(r,d,-1)'s
+// own send in this same step targets exactly rank r (send direction is always +1, and
+// neighbor(neighbor(r,d,-1), d, +1) == r by construction of neighbor()).
+//
+// Parameters
+// ----------
+// dims : const std::vector&
+// S : size_t -- total collective bytes; m = S/N is exchanged per dimension.
+//
+// Returns
+// -------
+// Schedule -- K steps, N Xfers each, no LocalOps.
+inline Schedule build_sendrecv(const std::vector& dims, size_t S) {
+ int N = num_ranks(dims);
+ int K = (int)dims.size();
+ size_t m = S / (size_t)N;
+ Schedule sched;
+ for (int d = 0; d < K; ++d) {
+ Step step;
+ for (int r = 0; r < N; ++r) {
+ int dst = neighbor(r, d, +1, dims);
+ check_edge_or_abort(r, dst, dims);
+ step.xfers.push_back({r, dst, BUF_SEND, BUF_RECV, 0, (size_t)d * m, m});
+ }
+ sched.push_back(std::move(step));
+ }
+ return sched;
+}
+
+// Builds the schedule for `broadcast` from root rank 0.
+//
+// Algorithm: dimension-by-dimension forward-chain propagation. A host-side has_data[N] tracks,
+// at schedule-BUILD time (not at run time), which ranks are already known to hold the
+// broadcast data after each step; this bookkeeping exists purely to decide which Xfers to
+// generate and is not part of the executed Schedule itself.
+//
+// Invariant: at the start of dimension d's phase, the set of ranks with has_data[r] == true is
+// exactly the set of ranks that agree with rank 0 on every coordinate j >= d (this holds for
+// d == 0 trivially: only rank 0 itself). Each of dimension d's (extent_d - 1) repeats extends
+// every already-seeded "line" one hop further in the +1 direction; after all extent_d - 1
+// repeats, every rank agreeing with rank 0 on coordinates j > d has data (regardless of its
+// coordinate d), establishing the invariant for phase d+1. After all K dimensions, every rank
+// has data.
+//
+// Design decision: root's own copy is written by a dedicated step-0 LocalOp (BUF_SEND ->
+// BUF_RECV) with NO concurrent Xfers in that same step. This lets every subsequent send (even
+// root's very first "real" send) source uniformly from BUF_RECV: had root's first send shared
+// step 0 with the LocalOp, it would need to special-case sourcing from BUF_SEND instead, since
+// a Step's post-LocalOps run strictly after that step's Xfers land.
+//
+// Parameters
+// ----------
+// dims : const std::vector&
+// S : size_t -- total collective bytes; m = S/N is the broadcast vector size.
+//
+// Returns
+// -------
+// Schedule -- 1 (LocalOp-only) + sum_d(extent_d - 1) steps.
+inline Schedule build_broadcast(const std::vector& dims, size_t S) {
+ int N = num_ranks(dims);
+ int K = (int)dims.size();
+ size_t m = S / (size_t)N;
+ Schedule sched;
+ {
+ Step step0;
+ step0.post.push_back({0, BUF_SEND, 0, BUF_RECV, 0, m, false});
+ sched.push_back(std::move(step0));
+ }
+ std::vector has_data(N, 0);
+ has_data[0] = 1;
+ for (int d = 0; d < K; ++d) {
+ int E = dims[d];
+ for (int s = 0; s < E - 1; ++s) {
+ Step step;
+ std::vector new_has = has_data; // conditions evaluated against pre-step state.
+ for (int r = 0; r < N; ++r) {
+ if (!has_data[r]) continue;
+ int dst = neighbor(r, d, +1, dims);
+ if (has_data[dst]) continue;
+ check_edge_or_abort(r, dst, dims);
+ step.xfers.push_back({r, dst, BUF_RECV, BUF_RECV, 0, 0, m});
+ new_has[dst] = 1;
+ }
+ has_data = new_has;
+ sched.push_back(std::move(step));
+ }
+ }
+ return sched;
+}
+
+// Builds the schedule for `all_gather`.
+//
+// Algorithm: standard ring all-gather, generalized to a mixed-radix torus by processing one
+// dimension at a time (ascending order). Step 0 seeds each rank's own slot via a LocalOp.
+//
+// Invariant: at the START of dimension d's phase, rank r holds exactly the slots
+// {u : u_j == r_j for all j >= d} (dimensions below d are already fully "free" -- rank r holds
+// every value there -- while dimensions >= d are still pinned to r's own coordinate). This
+// holds trivially at d == 0 (only slot r itself, u_j == r_j for ALL j) and, vacuously, means
+// every rank holds every slot once d reaches K (all dimensions processed).
+//
+// Each phase d performs (extent_d - 1) ring-relay steps in the +1 direction: at step s, rank r
+// forwards the slot batch it received on the PREVIOUS step (or, at s=1, the batch it started
+// the phase with) to its +1 neighbor. This is the classic "forward only what you just received"
+// ring relay, which avoids redundant retransmission and completes dimension d's expansion in
+// exactly extent_d - 1 hops.
+//
+// Parameters
+// ----------
+// dims : const std::vector&
+// S : size_t -- total collective bytes; m = S/N is each rank's own shard size.
+//
+// Returns
+// -------
+// Schedule -- 1 + sum_d(extent_d - 1) steps.
+inline Schedule build_all_gather(const std::vector& dims, size_t S) {
+ int N = num_ranks(dims);
+ int K = (int)dims.size();
+ size_t m = S / (size_t)N;
+ Schedule sched;
+ {
+ Step step0;
+ for (int r = 0; r < N; ++r)
+ step0.post.push_back({r, BUF_SEND, 0, BUF_RECV, (size_t)r * m, m, false});
+ sched.push_back(std::move(step0));
+ }
+ for (int d = 0; d < K; ++d) {
+ int E = dims[d];
+ for (int s = 1; s <= E - 1; ++s) {
+ Step step;
+ for (int r = 0; r < N; ++r) {
+ std::vector base = coords_of(r, dims);
+ int dst = neighbor(r, d, +1, dims);
+ check_edge_or_abort(r, dst, dims);
+ base[d] = ((base[d] - (s - 1)) % E + E) % E;
+ for_each_free_combo(d, dims, base, [&](const std::vector& v) {
+ int vslot = rank_of(v, dims);
+ size_t off = (size_t)vslot * m;
+ step.xfers.push_back({r, dst, BUF_RECV, BUF_RECV, off, off, m});
+ });
+ }
+ sched.push_back(std::move(step));
+ }
+ }
+ return sched;
+}
+
+// Builds the schedule for `reduce_scatter`.
+//
+// DERIVATION NOTE (see also file header): the planning spec's own tentative formula for this
+// collective was flagged mid-sentence as unreliable ("hold on, mirror the all_gather invariant
+// exactly"). This implementation instead derives the per-step slot set and DIRECTION from a
+// time-reversal argument against the (independently correct, spec-supplied) all_gather formula,
+// and that derivation was confirmed by an exhaustive brute-force simulation (dims=[8],[2,4],
+// [4,2],[2,2,2]) before being encoded here. Summary of the correction: the naive "mirror
+// all_gather with the same (r_d - (s-1)) mod E slot formula and +1 direction" places each fully
+// reduced slot ONE HOP SHORT of its destination rank (rank k's slot ends up fully summed at
+// rank k-1, not rank k). The fix is both a different slot formula AND a different direction:
+// sends go in the -1 direction, and at step t the sender's slot index is (r_d + t) mod E, not
+// (r_d - (s-1)) mod E.
+//
+// Algorithm: work happens in BUF_WORK_A (an N-slot accumulator seeded from BUF_SEND at step 0).
+// Dimensions are processed in DESCENDING order (K-1 down to 0), matching the spec.
+//
+// Invariant: at the END of dimension d's phase, rank r holds partial sums ONLY for slots
+// {v : v_j == r_j for all j >= d} (dimensions below d are not yet reduced -- rank r's held
+// slots still range over every value there -- while dimensions >= d have been fully reduced
+// down to r's own coordinate). This is exactly the all_gather invariant with the phase-transition
+// direction reversed: all_gather EXPANDS what a rank holds as d increases from 0 to K;
+// reduce_scatter CONTRACTS what a rank holds as d decreases from K-1 to 0, so by symmetry it is
+// stated as an END-of-phase (rather than start-of-phase) condition.
+//
+// Correctness of the per-step formula for a single dimension's ring (extent E, direction -1,
+// slot index v_d = (r_d + t) mod E at step t = 1..E-1): consider dimension d as an isolated ring
+// (the free dimensions jd ride along unchanged, in parallel, for
+// every value). For a fixed target slot-index k (a value of v_d), the ring of E nodes must sum
+// together each node's local contribution for chunk k, ending up entirely at node k. At reduce
+// step t, the node currently entrusted with chunk k's running partial sum is node (k - t) mod E
+// (t=1: node k-1 sends its OWN local value for chunk k to node k-2, which adds it in; t=2: node
+// k-2, now holding a 2-term partial sum, forwards it to node k-3; ...; t=E-1: node (k -
+// (E-1)) mod E == (k+1) mod E, holding an (E-1)-term partial sum -- every node except k itself
+// -- forwards it to node k, which adds in its own remaining term to complete the sum of all E
+// contributions). Restating "node (k - t) mod E sends chunk k to node (k - t - 1) mod E" from
+// the SENDER's own coordinate r = (k - t) mod E gives: sender r sends chunk k = (r + t) mod E to
+// receiver (r - 1) mod E == neighbor(r, d, -1). This is exactly the formula used below.
+//
+// Parameters
+// ----------
+// dims : const std::vector&
+// S : size_t -- total collective bytes; m = S/N, BUF_SEND holds N slots of m bytes each.
+//
+// Returns
+// -------
+// Schedule -- 1 (seed) + sum_d(extent_d - 1) (reduce) + 1 (final copy) steps.
+inline Schedule build_reduce_scatter(const std::vector& dims, size_t S) {
+ int N = num_ranks(dims);
+ int K = (int)dims.size();
+ size_t m = S / (size_t)N;
+ Schedule sched;
+ {
+ Step step0;
+ for (int r = 0; r < N; ++r)
+ step0.post.push_back({r, BUF_SEND, 0, BUF_WORK_A, 0, (size_t)N * m, false});
+ sched.push_back(std::move(step0));
+ }
+ for (int d = K - 1; d >= 0; --d) {
+ int E = dims[d];
+ for (int t = 1; t <= E - 1; ++t) {
+ Step step;
+ for (int r = 0; r < N; ++r) {
+ std::vector cr = coords_of(r, dims);
+ // Direction is -1 (see derivation above); this is the corrected direction,
+ // NOT the +1 direction all_gather uses.
+ int dst = neighbor(r, d, -1, dims);
+ check_edge_or_abort(r, dst, dims);
+ std::vector base = cr;
+ base[d] = ((cr[d] + t) % E + E) % E;
+ for_each_free_combo(d, dims, base, [&](const std::vector& v) {
+ int vslot = rank_of(v, dims);
+ size_t off = (size_t)vslot * m;
+ // Stage into the receiver's BUF_TMP at the same v*m offset (distinct
+ // offsets across combos within this step fall out automatically since each
+ // combo yields a distinct vslot), then a post LocalOp adds it into the
+ // receiver's running accumulator at that same slot.
+ step.xfers.push_back({r, dst, BUF_WORK_A, BUF_TMP, off, off, m});
+ step.post.push_back({dst, BUF_TMP, off, BUF_WORK_A, off, m, true});
+ });
+ }
+ sched.push_back(std::move(step));
+ }
+ }
+ {
+ Step stepf;
+ for (int r = 0; r < N; ++r)
+ stepf.post.push_back({r, BUF_WORK_A, (size_t)r * m, BUF_RECV, 0, m, false});
+ sched.push_back(std::move(stepf));
+ }
+ return sched;
+}
+
+// Builds the schedule for `all_reduce` by composing the (unmodified) reduce_scatter and
+// all_gather builders over sub-shards of size m/N, per the spec.
+//
+// Design: build_reduce_scatter(dims, m) is called with reduce_scatter's OWN "S" parameter set
+// to all_reduce's per-rank vector size m (not to S itself). Internally, reduce_scatter then
+// computes its own m_rs = m/N == the desired sub-shard size, and its own N*m_rs == m exactly
+// matches the size of all_reduce's per-rank input buffer -- so all_reduce's existing BUF_SEND
+// content can be fed to reduce_scatter completely unmodified, with no data rearrangement,
+// because a flat m-byte vector split into N contiguous m/N-byte pieces is precisely
+// reduce_scatter's own "N contiguous slots" input convention. Symmetrically,
+// build_all_gather(dims, m) produces an N*(m/N) == m byte output, exactly all_reduce's needed
+// result size.
+//
+// A single bridging LocalOp copies reduce_scatter's own (small, m/N-byte) output out of its
+// BUF_RECV into BUF_SEND, where the reused all_gather schedule's own step 0 expects to find its
+// input. This reuse is safe because reduce_scatter only ever READS BUF_SEND once, in its own
+// step 0; by the time its schedule finishes, BUF_SEND is dead and free to reuse as scratch.
+//
+// Parameters
+// ----------
+// dims : const std::vector&
+// S : size_t -- total collective bytes; m = S/N is each rank's full input/output vector size.
+//
+// Returns
+// -------
+// Schedule -- reduce_scatter(dims,m)'s steps, then 1 bridging step, then all_gather(dims,m)'s
+// steps.
+inline Schedule build_all_reduce(const std::vector& dims, size_t S) {
+ int N = num_ranks(dims);
+ size_t m = S / (size_t)N;
+ size_t shard = m / (size_t)N;
+
+ Schedule sched = build_reduce_scatter(dims, m);
+ {
+ Step bridge;
+ for (int r = 0; r < N; ++r)
+ bridge.post.push_back({r, BUF_RECV, 0, BUF_SEND, 0, shard, false});
+ sched.push_back(std::move(bridge));
+ }
+ Schedule ag = build_all_gather(dims, m);
+ for (auto& step : ag) sched.push_back(std::move(step));
+ return sched;
+}
+
+// Builds the schedule for `alltoall` using dimension-ordered minimal routing.
+//
+// Each of the N*(N-1) non-self chunks (u,v), u != v, starts at rank u (BUF_SEND offset v*c,
+// c = m/N) and must reach rank v (BUF_RECV offset u*c). Self chunks (u == v) never move; they
+// are resolved by a single step-0 LocalOp per rank.
+//
+// Routing: dimensions are processed in ascending order. Within dimension d's phase, every chunk
+// whose current holder disagrees with its target on coordinate d takes one hop per Step, in the
+// direction (+1 or -1) that minimizes remaining distance around that dimension's ring; the
+// phase repeats until no chunk needs to move in dimension d (bounded by extent_d - 1
+// iterations, the maximum possible remaining distance, with an explicit abort if that bound is
+// ever exceeded -- it should not be, since minimal-direction hops need at most
+// floor(extent_d/2) <= extent_d - 1 of them).
+//
+// In-flight buffer management: a chunk not currently at BUF_SEND or BUF_RECV lives in one of
+// two ping-pong work buffers, BUF_WORK_A/BUF_WORK_B, alternating buffers each time it takes a
+// hop (this guarantees a chunk's read-from buffer this step always differs from its
+// write-to buffer this step, so simple sequential-then-concurrent memory semantics --
+// read-all-then-write-all within a Step -- can never alias a chunk's own old and new copies of
+// itself).
+//
+// Slot allocation within a work buffer is a genuine per-step host-side bump allocator (a
+// closed-form offset formula is not safe in general: two chunks that share a source rank and
+// happen to have the same destination coordinate on every already-matched dimension travel
+// together and can simultaneously occupy the same intermediate rank, so slots must be assigned,
+// not computed). New allocations for step t are chosen to avoid every slot occupied at the
+// START of step t (including slots about to be vacated this same step, since those are still
+// being read from concurrently); old slots are freed only once the step's Xfers have been
+// fully built, so freed capacity becomes available starting with the NEXT step, never the
+// current one. Total occupancy per (rank, buffer) is asserted <= the buffer's slot capacity
+// (2*N slots of c bytes each, matching the 2*m-byte buffer size from buffer_sizes()); the spec
+// notes this should stay near N by symmetry for uniform all-to-all, which is also confirmed by
+// a max-occupancy check in the throwaway Python simulation used to validate this algorithm
+// before writing it here (observed max was well under half the budget for all tested dims).
+//
+// A chunk that arrives at its destination (all dimensions matched) is routed directly into
+// BUF_RECV instead of a work buffer and is marked done, removing it from further consideration.
+//
+// Parameters
+// ----------
+// dims : const std::vector&
+// S : size_t -- total collective bytes; m = S/N per rank, c = m/N per chunk.
+//
+// Returns
+// -------
+// Schedule -- 1 (self-chunk) step, followed by one step per routing hop actually taken.
+inline Schedule build_alltoall(const std::vector& dims, size_t S) {
+ int N = num_ranks(dims);
+ int K = (int)dims.size();
+ size_t m = S / (size_t)N;
+ size_t c = m / (size_t)N;
+ Schedule sched;
+ {
+ Step step0;
+ for (int r = 0; r < N; ++r)
+ step0.post.push_back({r, BUF_SEND, (size_t)r * c, BUF_RECV, (size_t)r * c, c, false});
+ sched.push_back(std::move(step0));
+ }
+
+ // Host-side (schedule-build-time only) bookkeeping of each in-flight chunk's current
+ // position. buf: -1 == still/again at BUF_SEND (never true after the first hop, but used
+ // as the initial state so the first hop's ping-pong toggle lands on BUF_WORK_A); 0 ==
+ // BUF_WORK_A; 1 == BUF_WORK_B. slot is meaningful only when buf >= 0.
+ struct A2AChunk {
+ int u, v, holder;
+ int buf;
+ int slot;
+ bool done;
+ };
+ std::vector chunks;
+ chunks.reserve((size_t)N * (N - 1));
+ for (int u = 0; u < N; ++u)
+ for (int v = 0; v < N; ++v)
+ if (u != v) chunks.push_back({u, v, u, -1, -1, false});
+
+ const int CAP_SLOTS = 2 * N; // 2*m bytes / c bytes-per-slot, matching buffer_sizes().
+ std::vector, 2>> occupied(N);
+ for (int r = 0; r < N; ++r) {
+ occupied[r][0].assign(CAP_SLOTS, 0);
+ occupied[r][1].assign(CAP_SLOTS, 0);
+ }
+ auto alloc_slot = [&](int rank, int buf) -> int {
+ for (int s = 0; s < CAP_SLOTS; ++s) {
+ if (!occupied[rank][buf][s]) {
+ occupied[rank][buf][s] = 1;
+ return s;
+ }
+ }
+ std::fprintf(stderr,
+ "%s:%d: alltoall work-buffer overflow at rank %d buf %d (occupancy "
+ "would exceed 2*m bytes)\n",
+ __FILE__, __LINE__, rank, buf);
+ std::abort();
+ return -1; // unreachable
+ };
+
+ for (int d = 0; d < K; ++d) {
+ int E = dims[d];
+ int iter = 0;
+ while (true) {
+ std::vector movers;
+ for (size_t i = 0; i < chunks.size(); ++i) {
+ if (chunks[i].done) continue;
+ std::vector cr = coords_of(chunks[i].holder, dims);
+ std::vector cv = coords_of(chunks[i].v, dims);
+ if (cr[d] != cv[d]) movers.push_back(i);
+ }
+ if (movers.empty()) break;
+ ++iter;
+ if (iter > E - 1) {
+ std::fprintf(stderr,
+ "%s:%d: alltoall routing failed to converge in dimension %d "
+ "within %d iterations\n",
+ __FILE__, __LINE__, d, E - 1);
+ std::abort();
+ }
+
+ struct Move {
+ size_t chunk_idx;
+ int dst_rank;
+ int new_buf, new_slot;
+ bool arrives;
+ };
+ std::vector moves;
+ moves.reserve(movers.size());
+ // Pass 1: decide direction/destination and allocate NEW slots against the
+ // occupancy snapshot as of the start of this step (see design note above: old
+ // slots are deliberately not freed until pass 3, so a slot being read from this
+ // step is never handed out as someone else's new landing spot this same step).
+ for (size_t i : movers) {
+ A2AChunk& ch = chunks[i];
+ std::vector cr = coords_of(ch.holder, dims);
+ std::vector cv = coords_of(ch.v, dims);
+ int distp = ((cv[d] - cr[d]) % E + E) % E;
+ int distm = ((cr[d] - cv[d]) % E + E) % E;
+ int dirn = (distp <= distm) ? +1 : -1;
+ int dst = neighbor(ch.holder, d, dirn, dims);
+ check_edge_or_abort(ch.holder, dst, dims);
+ std::vector cdst = coords_of(dst, dims);
+ bool arrives = true;
+ for (int j = 0; j < K; ++j)
+ if (cdst[j] != cv[j]) { arrives = false; break; }
+ int new_buf = -1, new_slot = -1;
+ if (!arrives) {
+ new_buf = (ch.buf == 0) ? 1 : 0; // ping-pong: SEND or B -> A; A -> B.
+ new_slot = alloc_slot(dst, new_buf);
+ }
+ moves.push_back({i, dst, new_buf, new_slot, arrives});
+ }
+ // Pass 2: emit Xfers from each chunk's OLD position to its NEW position.
+ Step step;
+ for (const Move& mv : moves) {
+ A2AChunk& ch = chunks[mv.chunk_idx];
+ int src_buf = (ch.buf < 0) ? (int)BUF_SEND : (ch.buf == 0 ? (int)BUF_WORK_A : (int)BUF_WORK_B);
+ size_t src_off = (ch.buf < 0) ? (size_t)ch.v * c : (size_t)ch.slot * c;
+ int dst_buf;
+ size_t dst_off;
+ if (mv.arrives) {
+ dst_buf = BUF_RECV;
+ dst_off = (size_t)ch.u * c;
+ } else {
+ dst_buf = (mv.new_buf == 0) ? (int)BUF_WORK_A : (int)BUF_WORK_B;
+ dst_off = (size_t)mv.new_slot * c;
+ }
+ step.xfers.push_back({ch.holder, mv.dst_rank, src_buf, dst_buf, src_off, dst_off, c});
+ }
+ // Pass 3: now that this step's Xfers are fully built, free vacated slots and
+ // update chunk bookkeeping for the next iteration.
+ for (const Move& mv : moves) {
+ A2AChunk& ch = chunks[mv.chunk_idx];
+ if (ch.buf >= 0) occupied[ch.holder][ch.buf][ch.slot] = 0;
+ ch.holder = mv.dst_rank;
+ if (mv.arrives) {
+ ch.done = true;
+ ch.buf = -1;
+ ch.slot = -1;
+ } else {
+ ch.buf = mv.new_buf;
+ ch.slot = mv.new_slot;
+ }
+ }
+ sched.push_back(std::move(step));
+ }
+ }
+
+ for (const A2AChunk& ch : chunks) {
+ if (!ch.done) {
+ std::fprintf(stderr,
+ "%s:%d: alltoall chunk (u=%d,v=%d) failed to reach its destination\n",
+ __FILE__, __LINE__, ch.u, ch.v);
+ std::abort();
+ }
+ }
+ return sched;
+}
+
+// =====================================================================================
+// Intra-step aliasing validation (shared, pure host code -- runs as part of build_schedule()'s
+// final validation pass, so it applies identically to whichever executor -- simulator or
+// NCCL/CUDA -- ends up running the assembled Schedule).
+// =====================================================================================
+//
+// Design/WHY (see also run_schedule()'s Notes, in both the NCCL and simulator branches below,
+// for the two executors' actual ordering guarantees this assertion is checked against): the
+// host-memory simulator's run_schedule() gives every Step's Xfers STRONGER semantics than the
+// NCCL/CUDA executor actually provides -- it snapshots every Xfer's source data before writing
+// ANY Xfer's destination (see that function's own "Design decision" comment), so a schedule with
+// overlapping src/dst ranges within one Step would silently produce a correct result there even
+// though it would NOT on real hardware. The NCCL executor instead relies only on: (1) all of a
+// Step's Xfers issued together inside one ncclGroupStart/End, with no ordering guarantee among
+// DIFFERENT Xfers of that group beyond NCCL's own send/recv rendezvous (which pairs a specific
+// send with its specific matching recv -- it says nothing about two unrelated Xfers of the same
+// Step racing each other), and (2) per-rank CUDA stream order, which guarantees only that ops
+// enqueued LATER on one rank's stream see the effects of ops enqueued EARLIER on that SAME
+// stream -- in particular, a Step's post LocalOps (always enqueued, on every rank's stream,
+// strictly after that Step's Xfers -- see run_schedule()) are guaranteed to see every effect of
+// that Step's own Xfers, but no other cross-operation ordering is guaranteed by construction.
+//
+// This function makes that gap an assertion instead of a latent, hard-to-reproduce bug: every
+// schedule this file ever builds is validated here to never depend on any same-step read/write
+// ordering stronger than "a post-op reads a range a same-step Xfer just wrote" -- the one
+// relationship the NCCL executor's fixed enqueue order (a Step's Xfers, then that Step's
+// post-ops) always provides "for free", with no host synchronization required. Any OTHER
+// same-step read/write overlap on the same (rank, buffer) -- e.g. one Xfer's source overlapping
+// another Xfer's destination, or a post-op's destination overlapping another post-op's source --
+// would only be safe under the simulator's stronger snapshot semantics, and must never occur.
+
+// One (rank, buffer) byte range read or written by one Xfer or LocalOp within a single Step, used
+// only by check_step_aliasing_or_abort() below.
+struct _ByteRange {
+ int rank;
+ int buf;
+ size_t begin, end; // half-open [begin, end), in bytes.
+ bool is_post; // true if this range comes from a LocalOp (post); false if from an Xfer.
+};
+
+// Returns whether two half-open byte ranges overlap.
+inline bool _ranges_overlap(size_t a_begin, size_t a_end, size_t b_begin, size_t b_end) {
+ return a_begin < b_end && b_begin < a_end;
+}
+
+// Validates one Step against the aliasing rule described above; aborts with a diagnostic naming
+// the step index and both offending ranges on the first violation found.
+//
+// Parameters
+// ----------
+// step_idx : size_t -- index of `step` within its Schedule (diagnostics only).
+// step : const Step&
+//
+// Notes
+// -----
+// The single allowed exception -- a post-op's SOURCE range overlapping an Xfer's DESTINATION
+// range, both of the same Step -- is exactly the reduce_scatter reduce-phase pattern (an Xfer
+// lands a value in BUF_TMP, and that same Step's post-op immediately adds it out of BUF_TMP);
+// see run_schedule()'s Notes for why that ordering (post-ops always after that Step's Xfers) is
+// safe on both executors. Every other overlap combination aborts.
+inline void check_step_aliasing_or_abort(size_t step_idx, const Step& step) {
+ std::vector<_ByteRange> reads, writes;
+ for (const Xfer& x : step.xfers) {
+ reads.push_back({x.src, x.src_buf, x.src_off, x.src_off + x.bytes, false});
+ writes.push_back({x.dst, x.dst_buf, x.dst_off, x.dst_off + x.bytes, false});
+ }
+ for (const LocalOp& op : step.post) {
+ reads.push_back({op.rank, op.src_buf, op.src_off, op.src_off + op.bytes, true});
+ writes.push_back({op.rank, op.dst_buf, op.dst_off, op.dst_off + op.bytes, true});
+ }
+ for (const _ByteRange& r : reads) {
+ for (const _ByteRange& w : writes) {
+ if (r.rank != w.rank || r.buf != w.buf) continue;
+ if (!_ranges_overlap(r.begin, r.end, w.begin, w.end)) continue;
+ // The one guaranteed-safe relationship: a post-op reading exactly what a same-step
+ // Xfer just wrote (see this function's Notes above and run_schedule()'s Notes).
+ if (r.is_post && !w.is_post) continue;
+ std::fprintf(
+ stderr,
+ "%s:%d: INTRA-STEP ALIAS ASSERTION FAILED at step %zu: rank %d buf %s read "
+ "range [%zu,%zu) (from %s) overlaps write range [%zu,%zu) (from %s); this "
+ "schedule relies on same-step read/write ordering the NCCL executor does not "
+ "guarantee (see check_step_aliasing_or_abort()'s Notes)\n",
+ __FILE__, __LINE__, step_idx, r.rank, buf_name(r.buf), r.begin, r.end,
+ r.is_post ? "a post-op" : "an Xfer", w.begin, w.end,
+ w.is_post ? "a post-op" : "an Xfer");
+ std::abort();
+ }
+ }
+}
+
+// Dispatches to the appropriate per-collective builder and then re-validates every Xfer in the
+// assembled schedule against the torus-edge guarantee, and every Step against the intra-step
+// aliasing rule above.
+//
+// Notes
+// -----
+// Every builder above already calls check_edge_or_abort() at each Xfer's construction site,
+// which fails fastest and with the most local context. The bulk re-scan here is a deliberate
+// belt-and-suspenders duplication: it makes build_schedule() itself -- not just its helpers --
+// the literal authority for the scientific guarantee described in the file header, matching the
+// spec's requirement that "build_schedule asserts is_torus_neighbor(...) for every Xfer". The
+// aliasing pass (check_step_aliasing_or_abort()) has no earlier per-Xfer equivalent -- it is
+// inherently a whole-Step check -- so build_schedule() is the only place it can run.
+//
+// Parameters
+// ----------
+// c : Collective
+// dims : const std::vector& -- torus extents.
+// S : size_t -- total collective bytes for this sweep point.
+//
+// Returns
+// -------
+// Schedule -- fully built, edge-validated, and aliasing-validated.
+inline Schedule build_schedule(Collective c, const std::vector& dims, size_t S) {
+ Schedule sched;
+ switch (c) {
+ case Collective::SENDRECV: sched = build_sendrecv(dims, S); break;
+ case Collective::BROADCAST: sched = build_broadcast(dims, S); break;
+ case Collective::ALL_GATHER: sched = build_all_gather(dims, S); break;
+ case Collective::REDUCE_SCATTER: sched = build_reduce_scatter(dims, S); break;
+ case Collective::ALL_REDUCE: sched = build_all_reduce(dims, S); break;
+ case Collective::ALLTOALL: sched = build_alltoall(dims, S); break;
+ }
+ for (size_t i = 0; i < sched.size(); ++i) {
+ for (const Xfer& x : sched[i].xfers) check_edge_or_abort(x.src, x.dst, dims);
+ check_step_aliasing_or_abort(i, sched[i]);
+ }
+ return sched;
+}
+
+// =====================================================================================
+// --check data pattern: shared, pure-host generation and verification logic
+// =====================================================================================
+//
+// The --check input convention is uniform across all six collectives: every rank's BUF_SEND
+// buffer (whatever its collective-specific size happens to be) is filled with f(rank, i) at
+// flat float index i = 0 .. (buffer_bytes/4 - 1). This single convention subsumes every
+// per-collective fill rule described in the spec (e.g. reduce_scatter's "rank u's BUF_SEND slot
+// v holds f(u, v*elems_per_slot + i)" and alltoall's analogous per-chunk rule are both just this
+// same flat-buffer fill, since a slot/chunk is by construction a contiguous sub-range of the
+// flat buffer and v*elems_per_slot+i (or v*celems+i) is exactly that sub-range's flat index).
+//
+// Because f() is a pure, deterministic function of (source rank, index), verification never
+// needs to read another rank's actual buffer contents: for every element of a rank's BUF_RECV,
+// we know analytically which (source rank, source flat index) it is supposed to equal (or, for
+// the reducing collectives, which SET of them it is supposed to sum) and simply recompute f()
+// directly. This is implemented once, in verify_collective() below, and reused unmodified by
+// both executors.
+
+// Deterministic --check fill value.
+//
+// Parameters
+// ----------
+// r : int -- source rank.
+// i : long long -- flat float index within that rank's buffer; must be >= 0 in practice (the
+// modulo-97 reduction below is defensive against negative input but this file never
+// constructs a negative index).
+//
+// Returns
+// -------
+// float -- r*100 + (i mod 97). Exact in fp32 (see verify_collective() notes on why summing up
+// to N such values remains exactly representable for the N, dims this benchmark targets).
+inline float f_pattern(int r, long long i) {
+ long long im = i % 97;
+ if (im < 0) im += 97;
+ return (float)(r * 100 + (int)im);
+}
+
+// Generates the full --check fill pattern for one rank's BUF_SEND buffer.
+//
+// Parameters
+// ----------
+// rank : int
+// nbytes : size_t -- BUF_SEND size for this rank (from buffer_sizes()); must be a multiple of
+// sizeof(float) (guaranteed by the S % (4*N*N) == 0 sweep-level precondition).
+//
+// Returns
+// -------
+// std::vector -- nbytes/4 floats, out[i] == f_pattern(rank, i).
+inline std::vector gen_pattern(int rank, size_t nbytes) {
+ size_t nf = nbytes / sizeof(float);
+ std::vector out(nf);
+ for (size_t i = 0; i < nf; ++i) out[i] = f_pattern(rank, (long long)i);
+ return out;
+}
+
+// Diagnostic record for the first mismatch found by verify_collective(), used to format the
+// "# CHECK ... FAIL (first mismatch: ...)" line.
+struct MismatchInfo {
+ int rank = -1;
+ int buf = BUF_RECV;
+ long long idx = -1;
+ float want = 0.f;
+ float got = 0.f;
+};
+
+// Verifies one rank's BUF_RECV contents against the analytically-known-correct result for the
+// given collective, per the check rules in the spec (all reducible, per the discussion above,
+// to recomputing f_pattern() directly rather than reading other ranks' data).
+//
+// Parameters
+// ----------
+// c : Collective
+// dims : const std::vector&
+// S : size_t -- total collective bytes for this sweep point.
+// rank : int -- the rank whose BUF_RECV is being checked.
+// recv : const std::vector& -- that rank's full BUF_RECV contents, already copied to
+// host memory by the caller (trivial for the simulator; a cudaMemcpy D2H for the NCCL
+// build).
+// mm : MismatchInfo& -- filled in on the first mismatch found; unmodified if this returns true.
+//
+// Returns
+// -------
+// bool -- true iff every checked element matched exactly.
+//
+// Notes
+// -----
+// Comparisons use exact float equality. This is intentional and safe here, not a bug: every
+// f_pattern() value is a small non-negative integer (r*100 + i%97, comfortably under 2^24 for
+// the rank counts and indices this benchmark exercises), and summing up to N such values (N <=
+// a few hundred in any realistic torus) never leaves the range of exactly-representable
+// integers in fp32 at any intermediate step -- so no rounding ever occurs, in the executor's
+// float accumulation OR in this function's own reference summation, and exact comparison is
+// mathematically justified rather than merely convenient.
+inline bool verify_collective(Collective c, const std::vector& dims, size_t S, int rank,
+ const std::vector& recv, MismatchInfo& mm) {
+ int N = num_ranks(dims);
+ int K = (int)dims.size();
+ size_t m = S / (size_t)N;
+
+ auto check_eq = [&](size_t idx, float want, float got) -> bool {
+ if (got != want) {
+ mm.rank = rank;
+ mm.buf = BUF_RECV;
+ mm.idx = (long long)idx;
+ mm.want = want;
+ mm.got = got;
+ return false;
+ }
+ return true;
+ };
+
+ switch (c) {
+ case Collective::SENDRECV: {
+ size_t nf = m / sizeof(float);
+ for (int d = 0; d < K; ++d) {
+ int src = neighbor(rank, d, -1, dims);
+ for (size_t i = 0; i < nf; ++i) {
+ size_t idx = (size_t)d * nf + i;
+ if (!check_eq(idx, f_pattern(src, (long long)i), recv[idx])) return false;
+ }
+ }
+ return true;
+ }
+ case Collective::BROADCAST: {
+ size_t nf = m / sizeof(float);
+ for (size_t i = 0; i < nf; ++i)
+ if (!check_eq(i, f_pattern(0, (long long)i), recv[i])) return false;
+ return true;
+ }
+ case Collective::ALL_GATHER: {
+ size_t nf = m / sizeof(float);
+ for (int u = 0; u < N; ++u) {
+ for (size_t i = 0; i < nf; ++i) {
+ size_t idx = (size_t)u * nf + i;
+ if (!check_eq(idx, f_pattern(u, (long long)i), recv[idx])) return false;
+ }
+ }
+ return true;
+ }
+ case Collective::REDUCE_SCATTER: {
+ size_t elems_per_slot = m / sizeof(float);
+ for (size_t i = 0; i < elems_per_slot; ++i) {
+ float want = 0.f;
+ for (int u = 0; u < N; ++u)
+ want += f_pattern(u, (long long)((size_t)rank * elems_per_slot + i));
+ if (!check_eq(i, want, recv[i])) return false;
+ }
+ return true;
+ }
+ case Collective::ALL_REDUCE: {
+ size_t nf = m / sizeof(float);
+ for (size_t i = 0; i < nf; ++i) {
+ float want = 0.f;
+ for (int u = 0; u < N; ++u) want += f_pattern(u, (long long)i);
+ if (!check_eq(i, want, recv[i])) return false;
+ }
+ return true;
+ }
+ case Collective::ALLTOALL: {
+ size_t c_bytes = m / (size_t)N;
+ size_t celems = c_bytes / sizeof(float);
+ for (int u = 0; u < N; ++u) {
+ for (size_t i = 0; i < celems; ++i) {
+ size_t idx = (size_t)u * celems + i;
+ float want = f_pattern(u, (long long)((size_t)rank * celems + i));
+ if (!check_eq(idx, want, recv[idx])) return false;
+ }
+ }
+ return true;
+ }
+ }
+ return true;
+}
+
+// Prints the single required "# CHECK ..." line for one sweep size, per the exact output
+// format in the spec.
+inline void print_check_line(Collective c, size_t S, bool pass, const MismatchInfo& mm) {
+ if (pass) {
+ std::printf("# CHECK collective=%s S=%zu: PASS\n", collective_name(c), S);
+ } else {
+ std::printf(
+ "# CHECK collective=%s S=%zu: FAIL (first mismatch: rank=%d buf=%s idx=%lld "
+ "want=%g got=%g)\n",
+ collective_name(c), S, mm.rank, buf_name(mm.buf), mm.idx, (double)mm.want,
+ (double)mm.got);
+ }
+}
+
+// =====================================================================================
+// CLI helpers
+// =====================================================================================
+
+// Parses a --dims argument of the form "2x2x2", "8", or "2x4" into per-dimension extents.
+//
+// Parameters
+// ----------
+// s : const std::string& -- 'x'-separated positive integers.
+//
+// Returns
+// -------
+// std::vector -- one entry per dimension, each > 0. Empty on any parse failure (missing
+// token, non-digit character, or non-positive value), which callers treat as a CLI error.
+inline std::vector parse_dims(const std::string& s) {
+ std::vector out;
+ size_t pos = 0;
+ while (pos <= s.size()) {
+ size_t next = s.find('x', pos);
+ std::string tok = (next == std::string::npos) ? s.substr(pos) : s.substr(pos, next - pos);
+ if (tok.empty()) return {};
+ for (char ch : tok)
+ if (!std::isdigit((unsigned char)ch)) return {};
+ int val = std::atoi(tok.c_str());
+ if (val <= 0) return {};
+ out.push_back(val);
+ if (next == std::string::npos) break;
+ pos = next + 1;
+ }
+ return out;
+}
+
+// =====================================================================================
+// Executors: identical function signatures, divergent implementations.
+//
+// Everything above this point is pure host C++ with no CUDA dependency and is shared,
+// unmodified, by both builds. Everything below is fenced per the spec: CUDA/NCCL code lives
+// under `#ifndef TORUS_SIM`; the host-memory simulator (this environment's test vehicle, since
+// no nvcc/GPU is available here) lives in the `#else` branch. Both branches implement the exact
+// same set of function names and signatures -- alloc_buffers, free_buffers, sync_all_devices,
+// fill_input_pattern, run_schedule, run_check, global_teardown -- so main() below is written
+// once and never itself needs an #ifdef.
+// =====================================================================================
+
+#ifndef TORUS_SIM
+// --------------------------------------------------------------------------------
+// NCCL/CUDA executor. Requires nvcc + a CUDA toolkit + NCCL; NOT buildable or runnable in this
+// development environment (no nvcc/CUDA installed here -- see project report). Written to
+// mirror every simulator-side function signature exactly, per the spec, so this branch never
+// references anything sim-only.
+// --------------------------------------------------------------------------------
+#include
+#include
+
+// CUDA error-check macro: prints file:line and aborts on any non-success cudaError_t.
+#define CUDA_CHECK(call) \
+ do { \
+ cudaError_t _e = (call); \
+ if (_e != cudaSuccess) { \
+ std::fprintf(stderr, "%s:%d: CUDA error: %s\n", __FILE__, __LINE__, \
+ cudaGetErrorString(_e)); \
+ std::abort(); \
+ } \
+ } while (0)
+
+// NCCL error-check macro: prints file:line and aborts on any non-success ncclResult_t.
+#define NCCL_CHECK(call) \
+ do { \
+ ncclResult_t _r = (call); \
+ if (_r != ncclSuccess) { \
+ std::fprintf(stderr, "%s:%d: NCCL error: %s\n", __FILE__, __LINE__, \
+ ncclGetErrorString(_r)); \
+ std::abort(); \
+ } \
+ } while (0)
+
+// Elementwise in-place float accumulate: dst[i] += src[i] for i in [0, n). Used to implement
+// LocalOp with add == true on the device (the host-side simulator does the equivalent with a
+// plain loop; see the #else branch below).
+__global__ void add_inplace(float* dst, const float* src, size_t n) {
+ size_t i = (size_t)blockIdx.x * blockDim.x + threadIdx.x;
+ if (i < n) dst[i] += src[i];
+}
+
+static int g_N = 0;
+static std::vector> g_dev_bufs;
+static std::array g_buf_sizes{};
+static ncclComm_t* g_comms = nullptr;
+static cudaStream_t* g_streams = nullptr;
+
+// Allocates per-(rank,buffer) device storage for the upcoming sweep size, and lazily
+// initializes the (persistent, reused-across-sizes) NCCL communicators and per-device streams
+// on the first call.
+//
+// Parameters
+// ----------
+// N : int -- device/rank count (must match every subsequent call until the matching
+// free_buffers()).
+// sizes : const size_t[NUM_BUFS] -- from buffer_sizes(); a 0 entry allocates nothing (null ptr).
+inline void alloc_buffers(int N, const size_t sizes[NUM_BUFS]) {
+ g_N = N;
+ for (int b = 0; b < NUM_BUFS; ++b) g_buf_sizes[b] = sizes[b];
+ if (!g_comms) {
+ g_comms = new ncclComm_t[N];
+ std::vector devs(N);
+ for (int i = 0; i < N; ++i) devs[i] = i;
+ NCCL_CHECK(ncclCommInitAll(g_comms, N, devs.data()));
+ g_streams = new cudaStream_t[N];
+ for (int r = 0; r < N; ++r) {
+ CUDA_CHECK(cudaSetDevice(r));
+ CUDA_CHECK(cudaStreamCreate(&g_streams[r]));
+ }
+ }
+ g_dev_bufs.assign(N, {nullptr, nullptr, nullptr, nullptr, nullptr});
+ for (int r = 0; r < N; ++r) {
+ CUDA_CHECK(cudaSetDevice(r));
+ for (int b = 0; b < NUM_BUFS; ++b) {
+ if (sizes[b] > 0) CUDA_CHECK(cudaMalloc(&g_dev_bufs[r][b], sizes[b]));
+ }
+ }
+}
+
+// Frees the per-(rank,buffer) device storage allocated by the matching alloc_buffers() call.
+// Communicators/streams are intentionally NOT destroyed here (they are reused across sweep
+// sizes); see global_teardown() for final cleanup.
+inline void free_buffers(int N) {
+ for (int r = 0; r < N; ++r) {
+ CUDA_CHECK(cudaSetDevice(r));
+ for (int b = 0; b < NUM_BUFS; ++b)
+ if (g_dev_bufs[r][b]) CUDA_CHECK(cudaFree(g_dev_bufs[r][b]));
+ }
+ g_dev_bufs.clear();
+}
+
+// Blocks the host until every device has completed all work previously enqueued on its stream.
+inline void sync_all_devices() {
+ for (int r = 0; r < g_N; ++r) {
+ CUDA_CHECK(cudaSetDevice(r));
+ CUDA_CHECK(cudaDeviceSynchronize());
+ }
+}
+
+// Fills every rank's BUF_SEND with the deterministic --check pattern (see gen_pattern()),
+// H2D-copying a host-generated array rather than launching a fill kernel (simplest correct
+// option given the CUDA path cannot be tested locally; avoids a second untestable kernel).
+inline void fill_input_pattern(int N) {
+ size_t send_bytes = g_buf_sizes[BUF_SEND];
+ for (int r = 0; r < N; ++r) {
+ std::vector pat = gen_pattern(r, send_bytes);
+ CUDA_CHECK(cudaSetDevice(r));
+ CUDA_CHECK(cudaMemcpy(g_dev_bufs[r][BUF_SEND], pat.data(), send_bytes,
+ cudaMemcpyHostToDevice));
+ }
+}
+
+// Executes one full pass of `sched` across all N devices.
+//
+// Notes
+// -----
+// Correctness argument for the (deliberate) ABSENCE of any per-step host synchronization here
+// (see the file header's SCHEDULE-AS-DATA section for the two-executor contract this satisfies):
+//
+// - RANK-LOCAL ordering is preserved by CUDA stream order alone, with no host round-trip
+// needed. Every one of a given rank r's operations -- its ncclSend/ncclRecv calls (issued on
+// g_streams[r] via g_comms[r]), its cudaMemcpyAsync D2D post-ops, and its add_inplace kernel
+// launches -- are ALL enqueued on that same single stream g_streams[r], in the exact order
+// this function enqueues them (every Step's Xfers, in loop order, followed by that Step's
+// post LocalOps, in loop order, before moving to the next Step). CUDA guarantees operations
+// enqueued on one stream execute in that enqueue order; a later-enqueued op on a stream is
+// therefore guaranteed to see the effects of every earlier op on that SAME stream without
+// any explicit sync between them. This is exactly what build_schedule()'s intra-step alias
+// assertion (see that function) verifies is sufficient: it never lets a schedule reach this
+// executor if it would require a per-rank read/write ordering stronger than "post-op reads
+// what an Xfer of the same step already wrote" -- precisely the one relationship stream
+// order already provides here, for free.
+// - CROSS-RANK ordering (rank A's send must be matched by rank B's matching recv before either
+// side's dependent work proceeds) is enforced by NCCL itself, not by any host synchronization
+// this function performs: every Step's Xfers are issued inside one ncclGroupStart/End, and
+// NCCL's own send/recv rendezvous protocol is what guarantees a recv only completes (on its
+// own stream) once its matching send has actually transferred the data -- that handshake is
+// GPU-side, asynchronous, and requires no host cudaStreamSynchronize call to be correct.
+//
+// Given both of the above, a host round-trip is not needed between every Step, nor between a
+// Step's Xfers and its post LocalOps -- it is needed exactly ONCE per call to this function, to
+// give the HOST a defined point at which every device's work for this entire schedule execution
+// is known to have completed (callers that need to read results back, e.g. run_check() via
+// cudaMemcpy, or that are timing this call, already provide that host sync themselves --
+// sync_all_devices() in main()'s check/warmup/timed-loop call sites -- but performing it once
+// here too keeps this function's own postcondition self-contained rather than relying on every
+// caller to remember to do it). This drops per-iteration host-sync round-trips from ~2x the
+// number of Steps (the pre-fix per-step synchronize-after-Xfers +
+// synchronize-after-post-LocalOps pattern) to exactly 1, with no change to the timed loop's
+// measured semantics (main()'s warmup/timed loops already sync once before/after the whole w+n
+// iteration count, not between individual run_schedule() calls).
+inline void run_schedule(const Schedule& sched, int N) {
+ for (const Step& step : sched) {
+ ncclGroupStart();
+ for (const Xfer& x : step.xfers) {
+ char* sptr = (char*)g_dev_bufs[x.src][x.src_buf] + x.src_off;
+ char* dptr = (char*)g_dev_bufs[x.dst][x.dst_buf] + x.dst_off;
+ NCCL_CHECK(ncclSend(sptr, x.bytes, ncclChar, x.dst, g_comms[x.src], g_streams[x.src]));
+ NCCL_CHECK(ncclRecv(dptr, x.bytes, ncclChar, x.src, g_comms[x.dst], g_streams[x.dst]));
+ }
+ ncclGroupEnd();
+ for (const LocalOp& op : step.post) {
+ CUDA_CHECK(cudaSetDevice(op.rank));
+ char* sptr = (char*)g_dev_bufs[op.rank][op.src_buf] + op.src_off;
+ char* dptr = (char*)g_dev_bufs[op.rank][op.dst_buf] + op.dst_off;
+ if (!op.add) {
+ CUDA_CHECK(cudaMemcpyAsync(dptr, sptr, op.bytes, cudaMemcpyDeviceToDevice,
+ g_streams[op.rank]));
+ } else {
+ size_t nf = op.bytes / sizeof(float);
+ int threads = 256;
+ int blocks = (int)((nf + (size_t)threads - 1) / (size_t)threads);
+ add_inplace<<>>((float*)dptr,
+ (const float*)sptr, nf);
+ }
+ }
+ }
+ // Host synchronization happens ONCE per schedule execution, here, after the last Step --
+ // see the Notes above for why nothing between Steps (or between a Step's Xfers and its
+ // post-ops) needs it.
+ for (int r = 0; r < N; ++r) {
+ CUDA_CHECK(cudaSetDevice(r));
+ CUDA_CHECK(cudaStreamSynchronize(g_streams[r]));
+ }
+}
+
+// Runs the --check verification for every rank, D2H-copying each rank's BUF_RECV before
+// delegating to the shared verify_collective(). Prints the single required "# CHECK ..." line.
+//
+// Returns
+// -------
+// bool -- true iff every rank's BUF_RECV matched the expected result.
+inline bool run_check(Collective c, const std::vector& dims, size_t S, int N) {
+ bool overall_pass = true;
+ MismatchInfo first_mm;
+ for (int r = 0; r < N; ++r) {
+ size_t nbytes = g_buf_sizes[BUF_RECV];
+ size_t nf = nbytes / sizeof(float);
+ std::vector recv(nf);
+ CUDA_CHECK(cudaSetDevice(r));
+ CUDA_CHECK(cudaMemcpy(recv.data(), g_dev_bufs[r][BUF_RECV], nbytes, cudaMemcpyDeviceToHost));
+ MismatchInfo mm;
+ bool ok = verify_collective(c, dims, S, r, recv, mm);
+ if (!ok && overall_pass) {
+ overall_pass = false;
+ first_mm = mm;
+ }
+ }
+ print_check_line(c, S, overall_pass, first_mm);
+ return overall_pass;
+}
+
+// Destroys the persistent NCCL communicators and CUDA streams created lazily by the first
+// alloc_buffers() call. Safe to call even if alloc_buffers() was never called.
+inline void global_teardown() {
+ if (g_comms) {
+ for (int r = 0; r < g_N; ++r) ncclCommDestroy(g_comms[r]);
+ delete[] g_comms;
+ g_comms = nullptr;
+ }
+ if (g_streams) {
+ for (int r = 0; r < g_N; ++r) {
+ cudaSetDevice(r);
+ cudaStreamDestroy(g_streams[r]);
+ }
+ delete[] g_streams;
+ g_streams = nullptr;
+ }
+}
+
+#else
+// --------------------------------------------------------------------------------
+// Host-memory simulator executor (TORUS_SIM). No CUDA/NCCL dependency whatsoever -- this is
+// the local, GPU-free test vehicle exercised by the acceptance criteria in this environment.
+// Timings produced here are, per the spec, not scientifically meaningful (there is no real
+// interconnect being modeled); the schedule-building, edge-assertion, and per-collective
+// correctness logic under test is identical to what the NCCL branch would execute.
+// --------------------------------------------------------------------------------
+
+static std::vector, NUM_BUFS>> g_host_bufs;
+
+// Allocates zero-initialized per-(rank,buffer) host storage for the upcoming sweep size.
+//
+// Parameters
+// ----------
+// N : int -- rank count.
+// sizes : const size_t[NUM_BUFS] -- from buffer_sizes().
+inline void alloc_buffers(int N, const size_t sizes[NUM_BUFS]) {
+ g_host_bufs.assign((size_t)N, std::array, NUM_BUFS>{});
+ for (int r = 0; r < N; ++r)
+ for (int b = 0; b < NUM_BUFS; ++b) g_host_bufs[r][b].assign(sizes[b], 0);
+}
+
+// Releases the host storage allocated by the matching alloc_buffers() call.
+inline void free_buffers(int N) {
+ (void)N;
+ g_host_bufs.clear();
+}
+
+// No-op: the simulator is single-threaded host code, so every operation is already
+// synchronous by construction. Present only so main()'s driver loop can call the same function
+// name in both builds.
+inline void sync_all_devices() {}
+
+// Fills every rank's BUF_SEND with the deterministic --check pattern (see gen_pattern()).
+inline void fill_input_pattern(int N) {
+ for (int r = 0; r < N; ++r) {
+ std::vector pat = gen_pattern(r, g_host_bufs[r][BUF_SEND].size());
+ std::memcpy(g_host_bufs[r][BUF_SEND].data(), pat.data(), pat.size() * sizeof(float));
+ }
+}
+
+// Executes one full pass of `sched` over host memory.
+//
+// Design decision: within a Step, ALL Xfer sources are first snapshotted into temporary
+// buffers, and only then are all destinations written. This makes the "all xfers execute
+// concurrently, reading pre-step state" semantics of Step literally true regardless of
+// iteration order or any potential (believed absent, but not asserted) offset aliasing between
+// a step's reads and writes -- a small, cheap robustness margin given how load-bearing exact
+// schedule semantics are for this benchmark's scientific validity.
+inline void run_schedule(const Schedule& sched, int N) {
+ (void)N;
+ for (const Step& step : sched) {
+ std::vector> staged(step.xfers.size());
+ for (size_t i = 0; i < step.xfers.size(); ++i) {
+ const Xfer& x = step.xfers[i];
+ const std::vector& src = g_host_bufs[x.src][x.src_buf];
+ staged[i].assign(src.begin() + (long)x.src_off, src.begin() + (long)(x.src_off + x.bytes));
+ }
+ for (size_t i = 0; i < step.xfers.size(); ++i) {
+ const Xfer& x = step.xfers[i];
+ std::memcpy(g_host_bufs[x.dst][x.dst_buf].data() + x.dst_off, staged[i].data(), x.bytes);
+ }
+ for (const LocalOp& op : step.post) {
+ if (!op.add) {
+ std::memcpy(g_host_bufs[op.rank][op.dst_buf].data() + op.dst_off,
+ g_host_bufs[op.rank][op.src_buf].data() + op.src_off, op.bytes);
+ } else {
+ float* dst = reinterpret_cast(g_host_bufs[op.rank][op.dst_buf].data() + op.dst_off);
+ const float* src = reinterpret_cast(g_host_bufs[op.rank][op.src_buf].data() + op.src_off);
+ size_t nf = op.bytes / sizeof(float);
+ for (size_t i = 0; i < nf; ++i) dst[i] += src[i];
+ }
+ }
+ }
+}
+
+// Runs the --check verification for every rank directly against host memory (no device copy
+// needed) and prints the single required "# CHECK ..." line.
+//
+// Returns
+// -------
+// bool -- true iff every rank's BUF_RECV matched the expected result.
+inline bool run_check(Collective c, const std::vector& dims, size_t S, int N) {
+ bool overall_pass = true;
+ MismatchInfo first_mm;
+ for (int r = 0; r < N; ++r) {
+ const std::vector& raw = g_host_bufs[r][BUF_RECV];
+ size_t nf = raw.size() / sizeof(float);
+ std::vector recv(nf);
+ std::memcpy(recv.data(), raw.data(), raw.size());
+ MismatchInfo mm;
+ bool ok = verify_collective(c, dims, S, r, recv, mm);
+ if (!ok && overall_pass) {
+ overall_pass = false;
+ first_mm = mm;
+ }
+ }
+ print_check_line(c, S, overall_pass, first_mm);
+ return overall_pass;
+}
+
+// No-op in the simulator: there is no persistent device/communicator state to tear down.
+inline void global_teardown() {}
+
+#endif // TORUS_SIM
+
+// =====================================================================================
+// main(): CLI parsing, sweep loop, and output. Shared verbatim by both builds -- everything it
+// calls (alloc_buffers, free_buffers, fill_input_pattern, run_schedule, sync_all_devices,
+// run_check, global_teardown, build_schedule, buffer_sizes) has an identical signature in both
+// the CUDA and simulator branches above, so no #ifdef is needed here at all.
+// =====================================================================================
+
+int main(int argc, char** argv) {
+ std::string collective_str;
+ std::string dims_str;
+ long long b = -1, e = -1;
+ long long f = 2;
+ long long w = 5;
+ long long n = 20;
+ bool check = false;
+
+ for (int i = 1; i < argc; ++i) {
+ std::string a = argv[i];
+ auto need_val = [&](const char* flag) -> std::string {
+ if (i + 1 >= argc) {
+ std::fprintf(stderr, "torus_bench: missing value for %s\n", flag);
+ std::exit(2);
+ }
+ return std::string(argv[++i]);
+ };
+ if (a == "--collective") collective_str = need_val("--collective");
+ else if (a == "--dims") dims_str = need_val("--dims");
+ else if (a == "-b") b = std::atoll(need_val("-b").c_str());
+ else if (a == "-e") e = std::atoll(need_val("-e").c_str());
+ else if (a == "-f") f = std::atoll(need_val("-f").c_str());
+ else if (a == "-w") w = std::atoll(need_val("-w").c_str());
+ else if (a == "-n") n = std::atoll(need_val("-n").c_str());
+ else if (a == "--check") check = true;
+ else {
+ std::fprintf(stderr, "torus_bench: unknown argument '%s'\n", a.c_str());
+ return 2;
+ }
+ }
+
+ if (collective_str.empty() || dims_str.empty() || b <= 0 || e <= 0) {
+ std::fprintf(stderr,
+ "usage: torus_bench --collective {all_reduce,all_gather,reduce_scatter,"
+ "alltoall,broadcast,sendrecv} --dims -b "
+ "-e [-f 2] [-w 5] [-n 20] [--check]\n");
+ return 2;
+ }
+ if (f <= 1) {
+ std::fprintf(stderr, "torus_bench: -f must be >= 2 (got %lld)\n", f);
+ return 2;
+ }
+
+ Collective c;
+ if (!parse_collective(collective_str, c)) {
+ std::fprintf(stderr, "torus_bench: unknown --collective '%s'\n", collective_str.c_str());
+ return 2;
+ }
+ std::vector dims = parse_dims(dims_str);
+ if (dims.empty()) {
+ std::fprintf(stderr, "torus_bench: invalid --dims '%s'\n", dims_str.c_str());
+ return 2;
+ }
+ int N = num_ranks(dims);
+
+#ifdef TORUS_SIM
+ std::printf("# torus_bench [SIMULATOR build -- host memory only, no GPU/NCCL]\n");
+#else
+ std::printf("# torus_bench [NCCL/CUDA build]\n");
+#endif
+ std::printf(
+ "# collective=%s dims=%s N=%d range=[%lld,%lld] factor=%lld warmup=%lld iters=%lld "
+ "check=%d\n",
+ collective_name(c), dims_str.c_str(), N, b, e, f, w, n, check ? 1 : 0);
+
+ bool any_check_failure = false;
+ for (long long S = b; S <= e; S *= f) {
+ size_t divisor = 4ull * (size_t)N * (size_t)N;
+ if ((size_t)S % divisor != 0) {
+ std::printf(
+ "# SKIP S=%lld not divisible by 4*N*N=%zu (N=%d): shard/chunk boundaries "
+ "would not align to whole floats\n",
+ S, divisor, N);
+ continue;
+ }
+
+ Schedule sched = build_schedule(c, dims, (size_t)S);
+ std::array sizes = buffer_sizes(c, dims, (size_t)S);
+ alloc_buffers(N, sizes.data());
+
+ char check_char = '-';
+ if (check) {
+ fill_input_pattern(N);
+ run_schedule(sched, N);
+ sync_all_devices();
+ bool pass = run_check(c, dims, (size_t)S, N);
+ if (!pass) any_check_failure = true;
+ check_char = pass ? '1' : '0';
+ }
+
+ for (long long it = 0; it < w; ++it) run_schedule(sched, N);
+ sync_all_devices();
+ auto t0 = std::chrono::steady_clock::now();
+ for (long long it = 0; it < n; ++it) run_schedule(sched, N);
+ sync_all_devices();
+ auto t1 = std::chrono::steady_clock::now();
+ double total_us = std::chrono::duration(t1 - t0).count();
+ double avg_us = (n > 0) ? (total_us / (double)n) : 0.0;
+
+ free_buffers(N);
+
+ std::printf("TORUSBENCH,%s,%s,%lld,%.2f,%c\n", collective_name(c), dims_str.c_str(), S,
+ avg_us, check_char);
+ std::fflush(stdout);
+ }
+
+ global_teardown();
+ return any_check_failure ? 1 : 0;
+}
+
diff --git a/notebooks/astrasim2_correlation/distributed_models_demo.ipynb b/notebooks/astrasim2_correlation/distributed_models_demo.ipynb
new file mode 100644
index 00000000..e107ea4b
--- /dev/null
+++ b/notebooks/astrasim2_correlation/distributed_models_demo.ipynb
@@ -0,0 +1,906 @@
+{
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "id": "c733530b",
+ "metadata": {},
+ "source": [
+ "# Distributed Transfer Models: Fully-Connected, Star, and XY Routing --- Demo & Review Guide\n",
+ "\n",
+ "This notebook demonstrates the distributed ISL transfer models in\n",
+ "`accelforge/model/_looptree/reuse/isl/distributed/distributed_buffers.py`:\n",
+ "`FullyConnectedMulticastModel`, `StarMulticastModel`, `XYRoutingMulticastModel`,\n",
+ "and the `EdgePressure` per-link load abstraction they share with the older\n",
+ "`HypercubeMulticastModel`. It is a companion to\n",
+ "`notebooks/astrasim2_correlation/correlation.ipynb` (same conventions: one-hot\n",
+ "GPU coordinates, `isl.Map.read_from_str`, `model.apply(0, Fill(...), Occupancy(...))`).\n",
+ "\n",
+ "## What's new\n",
+ "\n",
+ "- **`FullyConnectedMulticastModel`** --- 1 hop per fabric-crossing delivery,\n",
+ " regardless of distance (the \"flat full-mesh\" view of a fully-connected fabric).\n",
+ "- **`StarMulticastModel`** --- the *spokes* realization of the same fabric\n",
+ " (`src -> switch -> dst`); exposes per-spoke `EdgePressure` where\n",
+ " `FullyConnectedMulticastModel` sees only a flat crossing count.\n",
+ "- **`XYRoutingMulticastModel`** --- dimension-order routing on a 2-D mesh, with\n",
+ " a full per-directed-mesh-link `EdgePressure` decomposition.\n",
+ "- **`EdgePressure`** --- `{ edge -> #multicast-trees-crossing-it }`, with\n",
+ " `.total()`, `.bottleneck()`, and `.eval_edge(name, coords)`, threaded through\n",
+ " `TransferInfo.edge_pressure` (`accelforge/model/_looptree/reuse/isl/spatial.py`).\n",
+ "- The `fulfilled_fill` / `unfulfilled_fill` partition (`_covered_fills`) and the\n",
+ " generic node-tuple lookup (`_mesh_node_tuple`) that every model above now\n",
+ " shares, replacing ad hoc per-model logic.\n",
+ "\n",
+ "## The two commits under review\n",
+ "\n",
+ "The code this notebook exercises arrived across two commits on `rengz-correl`:\n",
+ "\n",
+ "- **`142722f5`** --- *\"LLM implemented code from pseudocode but it needs\n",
+ " verification\"*: introduces `EdgePressure`, `StarMulticastModel`, and the\n",
+ " directed-mesh-link decomposition inside\n",
+ " `XYRoutingMulticastModel._directed_mesh_links`.\n",
+ "- **`81be62a4`** --- *\"isl distribuffers refactor\"*: consolidates the\n",
+ " `_eval_const` / `_const_pwq` / `_covered_fills` / `_mesh_node_tuple` helpers,\n",
+ " fixes a `fulfilled_fill`/`unfulfilled_fill` bug (every fill used to be\n",
+ " reported fulfilled unconditionally, even when no source held its datum), and\n",
+ " threads `EdgePressure` through `TransferInfo.edge_pressure` so `hops` and the\n",
+ " per-link decomposition come from one `identify_mesh_casts` call instead of\n",
+ " two independently computed paths.\n",
+ "\n",
+ "`HypercubeMulticastModel`, `identify_mesh_casts`, and `calculate_extents_per_dim`\n",
+ "predate both commits (they are the pre-existing, already-reviewed baseline) and\n",
+ "are not exercised here for that reason.\n",
+ "\n",
+ "## How to use this notebook as a review guide\n",
+ "\n",
+ "Each numbered section states which functions it exercises and cites their\n",
+ "current `file:line` ranges, recomputed against this working tree:\n",
+ "`EdgePressure` and its scalar-extraction helpers live in `edge_pressure.py`,\n",
+ "`identify_mesh_casts` and its helpers in `mesh_casts.py`, and the\n",
+ "`MulticastModel` base class (the shared `__init__`/`apply` shape every model\n",
+ "below uses) with each model's own cost kernel in `distributed_buffers.py`.\n",
+ "The citations reflect current file locations and line numbers, not the\n",
+ "original commit diffs. Every number asserted in a code cell below (`56`, `64`,\n",
+ "`448`, ...) was independently re-verified against this exact code before\n",
+ "being written here.\n",
+ "Section 6 collects the section -> function -> file:line-range -> commit\n",
+ "mapping into one table so a reviewer can walk the diff systematically.\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 1,
+ "id": "e3d6abf5",
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-07-06T22:33:22.330405Z",
+ "iopub.status.busy": "2026-07-06T22:33:22.330272Z",
+ "iopub.status.idle": "2026-07-06T22:33:23.645047Z",
+ "shell.execute_reply": "2026-07-06T22:33:23.644422Z"
+ }
+ },
+ "outputs": [],
+ "source": [
+ "import islpy as isl\n",
+ "import pandas as pd\n",
+ "import matplotlib.pyplot as plt\n",
+ "\n",
+ "from accelforge.model._looptree.reuse.isl.distributed.distributed_buffers import (\n",
+ " FullyConnectedMulticastModel,\n",
+ " StarMulticastModel,\n",
+ " XYRoutingMulticastModel,\n",
+ ")\n",
+ "from accelforge.model._looptree.reuse.isl.distributed.edge_pressure import (\n",
+ " _eval_const,\n",
+ ")\n",
+ "from accelforge.model._looptree.reuse.isl.mapping_to_isl.types import (\n",
+ " Fill,\n",
+ " Occupancy,\n",
+ " SpatialTag,\n",
+ ")\n",
+ "\n",
+ "CTX = isl.DEFAULT_CONTEXT\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "031ba8c0",
+ "metadata": {},
+ "source": [
+ "## 2. Fully-connected vs. star: 8-GPU one-hot all-to-all\n",
+ "\n",
+ "GPU $i$ sits at one-hot coordinate $e_i$ of an 8-dimensional `noc` space (the\n",
+ "same encoding as `correlation.ipynb`). Two workloads distinguish the two models:\n",
+ "\n",
+ "- **all-to-all (unicast)**: `data[s, d]` is a unique chunk from GPU $s$ to GPU\n",
+ " $d$ --- every delivery has exactly one source and one destination.\n",
+ "- **broadcast**: GPU $g$ holds one private chunk `data[g]`, requested by every\n",
+ " other GPU --- a genuine multicast (one producer, $N-1$ consumers, fan-out at\n",
+ " the switch).\n",
+ "\n",
+ "`FullyConnectedMulticastModel` costs every fabric-crossing delivery at 1 hop\n",
+ "regardless of the payload pattern, so **both** workloads give\n",
+ "$N(N-1) = 56$ hops. `StarMulticastModel` is the *spokes* realization of the\n",
+ "same fabric --- a delivery routes `src -> switch -> dst` --- and its per-spoke\n",
+ "`EdgePressure` **is** pattern-sensitive: only the broadcast pattern exposes the\n",
+ "switch's fan-out (egress load $1$, not $N-1$), because the all-to-all pattern\n",
+ "has no shared payload to fan out in the first place.\n",
+ "\n",
+ "**Code exercised**: `MulticastModel.apply` (the shared `apply` every model\n",
+ "below now uses; `distributed_buffers.py:97-147`) together with\n",
+ "`FullyConnectedMulticastModel._transfer_cost`/`_cost_fully_connected`\n",
+ "(`distributed_buffers.py:265-284`, commit `81be62a4`); `StarMulticastModel`\n",
+ "(`distributed_buffers.py:523-599`, commit `142722f5`; its `hops`/\n",
+ "`edge_pressure` aggregation is now the shared `MulticastModel.apply` above,\n",
+ "fed by `StarMulticastModel._transfer_cost`, `distributed_buffers.py:546-547`);\n",
+ "`EdgePressure` (`edge_pressure.py:21-133`, commit `142722f5`).\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 2,
+ "id": "cd3dd43b",
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-07-06T22:33:23.647609Z",
+ "iopub.status.busy": "2026-07-06T22:33:23.647337Z",
+ "iopub.status.idle": "2026-07-06T22:33:23.655758Z",
+ "shell.execute_reply": "2026-07-06T22:33:23.655106Z"
+ }
+ },
+ "outputs": [],
+ "source": [
+ "def onehot_constraints(prefix: str, n: int) -> str:\n",
+ " \"\"\"One-hot constraints over dims `{prefix}0..{prefix}{n-1}`. Reused from correlation.ipynb.\"\"\"\n",
+ " bounds = \" and \".join(f\"0 <= {prefix}{i} <= 1\" for i in range(n))\n",
+ " hot = \" + \".join(f\"{prefix}{i}\" for i in range(n)) + \" = 1\"\n",
+ " return f\"{bounds} and {hot}\"\n",
+ "\n",
+ "\n",
+ "def linear_id(prefix: str, n: int) -> str:\n",
+ " \"\"\"Affine recovery of the GPU id from a one-hot vector: id = sum i*g_i. Reused from correlation.ipynb.\"\"\"\n",
+ " return \" + \".join(f\"{i}*{prefix}{i}\" for i in range(1, n))\n",
+ "\n",
+ "\n",
+ "def onehot_dist_fn(n: int) -> isl.Map:\n",
+ " \"\"\"Unit-cost distance function on the one-hot `noc` space: 0 hops if same\n",
+ " GPU, 1 hop otherwise. Same shape as correlation.ipynb's `all_to_all_maps`.\n",
+ " \"\"\"\n",
+ " xd = \", \".join(f\"xd{i}\" for i in range(n))\n",
+ " xs = \", \".join(f\"xs{i}\" for i in range(n))\n",
+ " same = \" and \".join(f\"xd{i} = xs{i}\" for i in range(n))\n",
+ " diff = \" or \".join(f\"(xd{i} < xs{i}) or (xd{i} > xs{i})\" for i in range(n))\n",
+ " return isl.Map.read_from_str(\n",
+ " CTX,\n",
+ " f\"{{ [noc[{xd}] -> noc[{xs}]] -> hops[0] : {same}; \"\n",
+ " f\" [noc[{xd}] -> noc[{xs}]] -> hops[1] : {diff} }}\",\n",
+ " )\n",
+ "\n",
+ "\n",
+ "def all_to_all_maps(n: int) -> tuple[isl.Map, isl.Map, isl.Map]:\n",
+ " \"\"\"Build (occupancy, fill, dist_fn) for the N-GPU unicast all-to-all:\n",
+ " data[s, d] is the unique chunk sent by GPU s to GPU d. Verbatim from\n",
+ " notebooks/astrasim2_correlation/correlation.ipynb.\n",
+ " \"\"\"\n",
+ " gs = \", \".join(f\"gs{i}\" for i in range(n))\n",
+ " gd = \", \".join(f\"gd{i}\" for i in range(n))\n",
+ " occ = isl.Map.read_from_str(\n",
+ " CTX,\n",
+ " f\"{{ noc[{gs}] -> data[s, d] : {onehot_constraints('gs', n)} \"\n",
+ " f\"and s = {linear_id('gs', n)} and 0 <= d < {n} }}\",\n",
+ " )\n",
+ " fill = isl.Map.read_from_str(\n",
+ " CTX,\n",
+ " f\"{{ noc[{gd}] -> data[s, d] : {onehot_constraints('gd', n)} \"\n",
+ " f\"and d = {linear_id('gd', n)} and 0 <= s < {n} }}\",\n",
+ " )\n",
+ " return occ, fill, onehot_dist_fn(n)\n",
+ "\n",
+ "\n",
+ "def broadcast_maps(n: int) -> tuple[isl.Map, isl.Map, isl.Map]:\n",
+ " \"\"\"Build (occupancy, fill, dist_fn) for the N-GPU one-hot *broadcast*\n",
+ " pattern: GPU g holds one private chunk data[g], requested by every other\n",
+ " GPU (one producer, N-1 consumers, fan-out at the switch). Same one-hot\n",
+ " node encoding as `all_to_all_maps`, different fill/occ.\n",
+ " \"\"\"\n",
+ " gs = \", \".join(f\"gs{i}\" for i in range(n))\n",
+ " gd = \", \".join(f\"gd{i}\" for i in range(n))\n",
+ " occ = isl.Map.read_from_str(\n",
+ " CTX,\n",
+ " f\"{{ noc[{gs}] -> data[d] : {onehot_constraints('gs', n)} \"\n",
+ " f\"and d = {linear_id('gs', n)} }}\",\n",
+ " )\n",
+ " fill = isl.Map.read_from_str(\n",
+ " CTX,\n",
+ " f\"{{ noc[{gd}] -> data[d] : {onehot_constraints('gd', n)} \"\n",
+ " f\"and 0 <= d < {n} and d != {linear_id('gd', n)} }}\",\n",
+ " )\n",
+ " return occ, fill, onehot_dist_fn(n)\n",
+ "\n",
+ "\n",
+ "N = 8\n",
+ "tags = [SpatialTag(i, 0) for i in range(N)]\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 3,
+ "id": "99287f8f",
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-07-06T22:33:23.658028Z",
+ "iopub.status.busy": "2026-07-06T22:33:23.657872Z",
+ "iopub.status.idle": "2026-07-06T22:33:24.595478Z",
+ "shell.execute_reply": "2026-07-06T22:33:24.594519Z"
+ }
+ },
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "FullyConnectedMulticastModel (all-to-all unicast): 56 hops\n"
+ ]
+ }
+ ],
+ "source": [
+ "# All-to-all (unicast) leg: reproduces correlation.ipynb's FullyConnectedMulticastModel\n",
+ "# result (56 hops = 8x7 fabric-crossing chunks) -- pattern-invariant, since FC only\n",
+ "# counts (dst, data) deliveries that cross the fabric, not whether a payload is shared.\n",
+ "occ_map, fill_map, dist_fn = all_to_all_maps(N)\n",
+ "fc_info = FullyConnectedMulticastModel(dist_fn).apply(\n",
+ " 0, Fill(tags, fill_map), Occupancy(tags, occ_map)\n",
+ ")\n",
+ "fc_hops = _eval_const(fc_info.hops)\n",
+ "assert fc_hops == N * (N - 1) == 56, fc_hops\n",
+ "print(f\"FullyConnectedMulticastModel (all-to-all unicast): {fc_hops} hops\")\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 4,
+ "id": "da1b83a5",
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-07-06T22:33:24.598198Z",
+ "iopub.status.busy": "2026-07-06T22:33:24.597829Z",
+ "iopub.status.idle": "2026-07-06T22:33:26.845318Z",
+ "shell.execute_reply": "2026-07-06T22:33:26.844213Z"
+ }
+ },
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "StarMulticastModel (broadcast): 64 hops = 8 injections + 56 deliveries\n",
+ " spoke_in per node : 7 (every node receives N-1 = 7)\n",
+ " spoke_out per node: 1 (every node sources 1 -- fan-out happens at the switch)\n",
+ " bottleneck spoke : 7\n",
+ " invariant: sum(spoke_in) = 56 == FullyConnected crossings = 56\n"
+ ]
+ }
+ ],
+ "source": [
+ "# Broadcast leg: each GPU's chunk is requested by every other GPU. This is what\n",
+ "# exposes StarMulticastModel's spoke asymmetry -- the switch fans a shared\n",
+ "# chunk out from a single egress -- which the unicast pattern above cannot show\n",
+ "# (there, every delivery is already a distinct chunk with no sharing to fan out).\n",
+ "def onehot_coords(g: int, n: int) -> list[int]:\n",
+ " \"\"\"One-hot coordinate vector for GPU `g` among `n` GPUs.\"\"\"\n",
+ " return [1 if i == g else 0 for i in range(n)]\n",
+ "\n",
+ "\n",
+ "occ_b, fill_b, dist_fn_b = broadcast_maps(N)\n",
+ "fill_bc = Fill(tags, fill_b)\n",
+ "occ_bc = Occupancy(tags, occ_b)\n",
+ "\n",
+ "fc_info_b = FullyConnectedMulticastModel(dist_fn_b).apply(0, fill_bc, occ_bc)\n",
+ "fc_hops_b = _eval_const(fc_info_b.hops)\n",
+ "assert fc_hops_b == 56, fc_hops_b\n",
+ "\n",
+ "star_info = StarMulticastModel(dist_fn_b).apply(0, fill_bc, occ_bc)\n",
+ "star_pressure = star_info.edge_pressure\n",
+ "star_hops = _eval_const(star_info.hops)\n",
+ "\n",
+ "spoke_in = [star_pressure.eval_edge(\"spoke_in\", onehot_coords(g, N)) for g in range(N)]\n",
+ "spoke_out = [star_pressure.eval_edge(\"spoke_out\", onehot_coords(g, N)) for g in range(N)]\n",
+ "\n",
+ "assert all(v == N - 1 for v in spoke_in), spoke_in\n",
+ "assert all(v == 1 for v in spoke_out), spoke_out\n",
+ "assert star_pressure.bottleneck() == N - 1 == 7\n",
+ "assert star_hops == N + N * (N - 1) == 64\n",
+ "\n",
+ "total_ingress = sum(spoke_in)\n",
+ "assert total_ingress == fc_hops_b == 56, (total_ingress, fc_hops_b)\n",
+ "\n",
+ "print(f\"StarMulticastModel (broadcast): {star_hops} hops = {N} injections + {N * (N - 1)} deliveries\")\n",
+ "print(f\" spoke_in per node : {spoke_in[0]} (every node receives N-1 = {N - 1})\")\n",
+ "print(f\" spoke_out per node: {spoke_out[0]} (every node sources 1 -- fan-out happens at the switch)\")\n",
+ "print(f\" bottleneck spoke : {star_pressure.bottleneck()}\")\n",
+ "print(f\" invariant: sum(spoke_in) = {total_ingress} == FullyConnected crossings = {fc_hops_b}\")\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "4524391d",
+ "metadata": {},
+ "source": [
+ "## 3. XY (dimension-order) routing on a 2-D mesh\n",
+ "\n",
+ "XY routing routes every packet along $x$ first, then $y$, so a multicast from\n",
+ "one source is a rigid tree: an X segment along the source row out to every\n",
+ "destination column, then an independent Y segment down each column from the\n",
+ "source row. Two geometries from\n",
+ "`tests/isl/distributed/xy_routing/test_cases.yaml` (already\n",
+ "oracle-verified there; re-verified here):\n",
+ "\n",
+ "- **Case B** --- source $(1,0)$ casting to $(0,2)$ and $(2,2)$: a single tree,\n",
+ " $6$ hops, every used link carries load $1$ (no sharing).\n",
+ "- **Case F** --- an $8 \\times 8$ *column flood*: datum $(d_0, d_1)$ lives at\n",
+ " node $(d_0, d_1)$ and is requested by every node in column $x = d_0$. Each\n",
+ " column's vertical links are shared by up to 7 overlapping trees: $448$ total\n",
+ " hops, bottleneck $7$, `yedge_u[0,6] = 7`, `yedge_d[0,1] = 6`.\n",
+ "\n",
+ "**Code exercised**: `MulticastModel.apply` (shared `apply`;\n",
+ "`distributed_buffers.py:97-147`) together with\n",
+ "`XYRoutingMulticastModel._transfer_cost`\n",
+ "(`distributed_buffers.py:414-415`, commit `81be62a4`); `_directed_mesh_links`\n",
+ "(`distributed_buffers.py:417-515`, commit `142722f5`, touched again in\n",
+ "`81be62a4`); `EdgePressure.bottleneck` / `EdgePressure.eval_edge`\n",
+ "(`edge_pressure.py:89-133`, commit `142722f5`)."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 5,
+ "id": "d6c4b869",
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-07-06T22:33:26.847779Z",
+ "iopub.status.busy": "2026-07-06T22:33:26.847607Z",
+ "iopub.status.idle": "2026-07-06T22:33:26.864876Z",
+ "shell.execute_reply": "2026-07-06T22:33:26.864027Z"
+ }
+ },
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Case B: hops = 6, every used edge load 1 (single tree, no sharing)\n"
+ ]
+ },
+ {
+ "data": {
+ "text/html": [
+ "