diff --git a/tests/utils/mesh_test.py b/tests/utils/mesh_test.py index b988ef586..f1c837943 100644 --- a/tests/utils/mesh_test.py +++ b/tests/utils/mesh_test.py @@ -480,6 +480,27 @@ def __init__(self, device_id, coords, process_index): [device.id for device in allocated], [0, 1, 4, 5, 8, 9, 12, 13] # pyrefly: ignore[not-iterable] ) + def test_allocate_devices_by_coords_supports_edge_family_2d_subslice(self): + class FakeDevice: + + def __init__(self, device_id, coords): + self.id = device_id + self.coords = coords + self.process_index = 0 + self.slice_index = 0 + self.device_kind = "TPU v5e" + + fake_devices = [ + FakeDevice(0, (0, 0)), + FakeDevice(1, (1, 0)), + FakeDevice(2, (0, 1)), + FakeDevice(3, (1, 1)), + ] + + allocated, _ = mesh._allocate_devices_by_coords(fake_devices, 2) + + self.assertEqual([device.id for device in allocated], [0, 1]) + def test_allocate_named_mesh_device_slices_prefers_coord_boxes(self): class FakeDevice: @@ -552,6 +573,55 @@ def __init__(self, process_index): (None, None), ) + def test_find_best_candidate_coords_does_not_fallback_when_candidate_shapes_empty( + self, + ): + class FakeDevice: + + def __init__(self, device_id, coords): + self.id = device_id + self.coords = coords + + fake_devices = [ + FakeDevice(0, (0, 0, 0)), + FakeDevice(1, (1, 0, 0)), + FakeDevice(2, (0, 1, 0)), + FakeDevice(3, (1, 1, 0)), + ] + + coord_topology = mesh.get_coord_topology(fake_devices) + + self.assertIsNotNone(coord_topology) + self.assertIsNone( + mesh.find_best_candidate_coords( + coord_topology, + 4, + candidate_shapes=[], + ) + ) + + def test_allocate_devices_by_coords_rejects_unsupported_generic_fish_box( + self, + ): + class FakeDevice: + + def __init__(self, device_id, coords): + self.id = device_id + self.coords = coords + self.device_kind = "TPU v7" + + fake_devices = [] + device_id = 0 + for x in range(4): + for y in range(4): + for z in range(8): + fake_devices.append(FakeDevice(device_id, (x, y, z))) + device_id += 1 + + allocated, _ = mesh._allocate_devices_by_coords(fake_devices, 48) + + self.assertIsNone(allocated) + def test_allocate_devices_by_coords_returns_best_contiguous_box(self): class FakeDevice: @@ -607,6 +677,42 @@ def __init__(self, device_id, coords): self.assertEqual(mins, (0, 0, 0)) self.assertEqual(maxs, (3, 7, 7)) + def test_allocate_devices_by_coords_compact_policy_prefers_simpler_remainder( + self, + ): + class FakeDevice: + + def __init__(self, device_id, coords): + self.id = device_id + self.coords = coords + self.device_kind = "TPU v7" + + fake_devices = [] + device_id = 0 + for x in range(16): + for y in range(16): + for z in range(16): + fake_devices.append(FakeDevice(device_id, (x, y, z))) + device_id += 1 + + allocated, _ = mesh._allocate_devices_by_coords( + fake_devices, + 512, + allocation_policy="COMPACT", + ) + + allocated_coords = [device.coords for device in allocated] + mins = tuple( + min(coords[dim] for coords in allocated_coords) for dim in range(3) + ) + maxs = tuple( + max(coords[dim] for coords in allocated_coords) for dim in range(3) + ) + + self.assertLen(allocated, 512) + self.assertEqual(mins, (0, 0, 0)) + self.assertEqual(maxs, (3, 7, 15)) + def test_allocate_devices_tracks_remaining_coord_regions(self): class FakeDevice: @@ -1233,6 +1339,18 @@ def __init__(self, devices, axis_names, axis_types=None): ) self.assertEqual(created_mesh.axis_names, ("x", "y")) + def test_create_mesh_raises_assigned_device_count_mismatch(self): + with self.assertRaises(ValueError) as exc: + mesh.create_mesh( + (2, 2), + ("x", "y"), + devices=["d0", "d1", "d2", "d3", "d4"], + ) + + self.assertIn("but was assigned 5", str(exc.exception)) + self.assertIn("axis_names=('x', 'y')", str(exc.exception)) + self.assertIn("assigned_device_sample", str(exc.exception)) + def test_allocate_named_mesh_device_slices_uses_jax_devices_by_default(self): class FakeDevice: diff --git a/tests/utils/topology_test.py b/tests/utils/topology_test.py index 11cfe8dc2..42cee3a96 100644 --- a/tests/utils/topology_test.py +++ b/tests/utils/topology_test.py @@ -38,6 +38,22 @@ def test_best_topology_shapes_for_chip_count_returns_unique_edge_shape(self): [(2, 4, 1)], ) + def test_supported_topology_shapes_for_chip_count_returns_edge_subslice_shapes( + self, + ): + self.assertEqual( + topology.supported_topology_shapes_for_chip_count( + "TPU v6e", 2, chip_rank=2 + ), + [(2, 1), (1, 2)], + ) + self.assertEqual( + topology.best_topology_shapes_for_chip_count( + "TPU v6e", 2, chip_rank=2 + ), + [(2, 1)], + ) + def test_best_topology_shapes_for_chip_count_returns_empty_for_unsupported_edge_count( self, ): @@ -55,6 +71,14 @@ def test_best_topology_shapes_for_chip_count_prefers_most_cubical_fish_shape( [(4, 8, 8)], ) + def test_supported_topology_shapes_for_chip_count_returns_all_fish_shapes( + self, + ): + self.assertEqual( + topology.supported_topology_shapes_for_chip_count("TPU v7", 512), + [(8, 8, 8), (4, 8, 16), (4, 4, 32)], + ) + def test_best_topology_shapes_for_chip_count_supports_single_host_fish_subslice( self, ): diff --git a/tunix/utils/mesh.py b/tunix/utils/mesh.py index a1bb5687b..55012e828 100644 --- a/tunix/utils/mesh.py +++ b/tunix/utils/mesh.py @@ -81,21 +81,32 @@ def create_mesh( if len(axis_shapes) != len(axis_names): raise ValueError( f"mesh.shape {axis_shapes} and mesh.axis_names {axis_names} " - "must have the same length." + "must have the same length " + f"(got {len(axis_shapes)} shapes and {len(axis_names)} names)." ) num_devices = len(devices) if devices is not None else jax.device_count() required_devices = int(np.prod(axis_shapes)) if required_devices > num_devices: + assigned_device_sample = ( + summarize_devices_for_logging(list(devices), limit=8) + if devices is not None + else None + ) raise ValueError( f"Mesh shape {axis_shapes} requires {required_devices} devices, " - f"but found {num_devices}." + f"but found {num_devices}. axis_names={axis_names}. " + f"assigned_device_sample={assigned_device_sample}." ) if devices is not None: if required_devices != num_devices: + assigned_device_sample = summarize_devices_for_logging( + list(devices), limit=8 + ) raise ValueError( f"Mesh shape {axis_shapes} requires {required_devices} devices, " - f"but was assigned {num_devices}." + f"but was assigned {num_devices}. axis_names={axis_names}. " + f"assigned_device_sample={assigned_device_sample}." ) return jax.sharding.Mesh( np.array(list(devices)).reshape(axis_shapes), @@ -654,6 +665,8 @@ def _supported_coord_box_shapes( devices: Sequence[Any], coord_topology: CoordTopology, required_devices: int, + *, + allocation_policy: str = _COMPACT_ALLOCATION_POLICY, available_coord_shape: Sequence[int] | None = None, ) -> list[tuple[int, ...]] | None: """Returns topology-valid physical box shapes for the current device pool. @@ -667,6 +680,8 @@ def _supported_coord_box_shapes( devices: The candidate device pool. coord_topology: The coordinate topology containing device shape boundaries. required_devices: Requested device count. + allocation_policy: Whether to keep all legal topology shapes (compact) or + only the most cubical one (performance). available_coord_shape: Optional remaining region boundaries. Returns: @@ -675,6 +690,14 @@ def _supported_coord_box_shapes( """ family = topology._device_family(devices) # pylint: disable=protected-access if family is None: + logging.info( + "Coord allocation is using generic contiguous-box enumeration because " + "the accelerator family could not be resolved: required_devices=%d " + "topology_shape=%s device_sample=%s.", + required_devices, + coord_topology.max_shape, + summarize_devices_for_logging(devices, debug=True, limit=8), + ) return None # The topology tables reason in physical chips, but the request is in logical @@ -682,6 +705,15 @@ def _supported_coord_box_shapes( # chips, and the chip count drives the topology-shape search below. core_count = infer_core_on_chip_count(devices) or 1 if required_devices % core_count != 0: + logging.info( + "Coord allocation has no supported topology shape because the request " + "does not align to whole chips: required_devices=%d cores_per_chip=%d " + "family=%s topology_shape=%s.", + required_devices, + core_count, + family, + coord_topology.max_shape, + ) return [] # Drop the trailing core axis (if present) so we search in chip space. @@ -689,19 +721,43 @@ def _supported_coord_box_shapes( 1 if coord_topology.has_core_on_chip_dimension else 0 ) if chip_rank <= 0: + logging.info( + "Coord allocation has no supported topology shape because the derived " + "chip rank is invalid: required_devices=%d family=%s num_dims=%d " + "has_core_axis=%s topology_shape=%s.", + required_devices, + family, + coord_topology.num_dims, + coord_topology.has_core_on_chip_dimension, + coord_topology.max_shape, + ) return [] available_shape = tuple(available_coord_shape or coord_topology.max_shape) available_chip_shape = available_shape[:chip_rank] required_chips = required_devices // core_count - candidate_chip_shapes = topology.best_topology_shapes_for_chip_count( + candidate_chip_shapes = topology.supported_topology_shapes_for_chip_count( family, required_chips, chip_rank=chip_rank, available_chip_shape=available_chip_shape, ) + if allocation_policy == _PERFORMANCE_ALLOCATION_POLICY: + candidate_chip_shapes = candidate_chip_shapes[:1] if not candidate_chip_shapes: + logging.info( + "Coord allocation found no supported topology shapes: " + "required_devices=%d required_chips=%d family=%s chip_rank=%d " + "available_chip_shape=%s allocation_policy=%s topology_shape=%s.", + required_devices, + required_chips, + family, + chip_rank, + available_chip_shape, + allocation_policy, + coord_topology.max_shape, + ) return [] # Re-attach the core axis so the returned shape is in device space again, # e.g. a (2, 1, 1) chip box on a 2-core chip becomes a (2, 1, 1, 2) box. @@ -960,38 +1016,90 @@ def _coord_box_score( ) -def _order_candidate_regions( - candidate_regions: Sequence[tuple[CoordRegion, Sequence[tuple[int, ...]]]], +def _compact_coord_box_score( + region: CoordRegion, + start: tuple[int, ...], + shape: tuple[int, ...], +) -> tuple[Any, ...]: + """Ranks compact allocations by how cleanly they preserve the remainder. + + Compact mode should prefer allocations that keep the remaining free space as + simple as possible: a single leftover slab is better than several fragmented + remainders, and a tighter leftover slab is better than a stretched one. + + Args: + region: Remaining coord region the box is carved from. + start: Candidate box origin. + shape: Candidate box shape. + + Returns: + A lexicographic sort key. Lower sorts better. + """ + remainder_regions = _split_coord_region(region, start, shape) + return ( + len(remainder_regions), + _coord_box_score(start, shape), + tuple( + sorted( + _coord_box_score(remainder.start, remainder.shape) + for remainder in remainder_regions + ) + ), + ) + + +def _allocation_box_score( + region: CoordRegion, + start: tuple[int, ...], + shape: tuple[int, ...], allocation_policy: str, -) -> list[tuple[CoordRegion, Sequence[tuple[int, ...]]]]: - """Orders candidate regions according to the requested allocation policy. +) -> tuple[Any, ...]: + """Returns the candidate-box score for the requested allocation policy.""" + if allocation_policy == _COMPACT_ALLOCATION_POLICY: + return _compact_coord_box_score(region, start, shape) + return _coord_box_score(start, shape) + + +def _order_region_candidates( + candidate_regions: Sequence[ + tuple[ + CoordRegion, + tuple[int, ...], + tuple[int, ...], + list[tuple[int, ...]], + ] + ], + allocation_policy: str, +) -> list[ + tuple[CoordRegion, tuple[int, ...], tuple[int, ...], list[tuple[int, ...]]] +]: + """Orders realized candidate boxes according to the requested policy. Args: - candidate_regions: Candidate remaining regions paired with the supported - shapes each region can realize for the current request. + candidate_regions: Candidate remaining regions paired with the best box each + region can realize for the current request. allocation_policy: ``COMPACT`` or ``PERFORMANCE``. Returns: - Candidate regions ordered according to the requested policy. + Candidate boxes ordered according to the requested policy. """ - # item = (region, supported_shapes); item[0].shape is the region's own size, - # item[1][0] is the best box shape it can realize. The two policies differ - # only in which of those two the primary sort key is: - # COMPACT -> rank by region size first (pack into the tightest region), - # PERFORMANCE -> rank by realizable box shape first (favor cubical boxes). + # item = (region, start, shape, coords). COMPACT still packs into the + # tightest remaining region first, but once a region is chosen it ranks the + # realized box by how simply it preserves the remainder. PERFORMANCE keeps the + # current box-shape-first behavior. if allocation_policy == _COMPACT_ALLOCATION_POLICY: return sorted( candidate_regions, key=lambda item: ( - _coord_box_score(item[0].start, item[0].shape), # region size first - _coord_box_score(item[0].start, item[1][0]), # then box shape + _coord_box_score(item[0].start, item[0].shape), + _allocation_box_score(item[0], item[1], item[2], allocation_policy), ), ) return sorted( candidate_regions, key=lambda item: ( - _coord_box_score(item[0].start, item[1][0]), # box shape first - _coord_box_score(item[0].start, item[0].shape), # then region size + _allocation_box_score(item[0], item[1], item[2], allocation_policy), + _coord_box_score(item[0].start, item[0].shape), ), ) @@ -1000,6 +1108,7 @@ def _find_best_candidate_box( coord_topology: CoordTopology, candidate_shapes: Sequence[tuple[int, ...]], *, + allocation_policy: str = _COMPACT_ALLOCATION_POLICY, region: CoordRegion | None = None, ) -> tuple[tuple[int, ...], tuple[int, ...], list[tuple[int, ...]]] | None: """Finds the best valid candidate box, optionally constrained to a region. @@ -1012,6 +1121,7 @@ def _find_best_candidate_box( Args: coord_topology: Normalized coord metadata for the candidate device pool. candidate_shapes: Exact box shapes to consider. + allocation_policy: Ranking policy for legal candidate boxes. region: Optional remaining coord region that candidate boxes must fit inside. When omitted, the scan considers the whole topology. @@ -1021,6 +1131,7 @@ def _find_best_candidate_box( """ best_candidate = None best_score = None + scoring_region = region or _full_coord_region(coord_topology) for shape in candidate_shapes: for start in sorted(coord_topology.coord_to_device): @@ -1031,7 +1142,9 @@ def _find_best_candidate_box( continue if not candidate_uses_whole_chips(coord_topology, candidate_coords): continue - score = _coord_box_score(start, shape) + score = _allocation_box_score( + scoring_region, start, shape, allocation_policy + ) if best_score is None or score < best_score: best_score = score best_candidate = (start, shape, candidate_coords) @@ -1043,6 +1156,8 @@ def find_best_candidate_coords( coord_topology: CoordTopology, required_devices: int, candidate_shapes: Sequence[tuple[int, ...]] | None = None, + *, + allocation_policy: str = _COMPACT_ALLOCATION_POLICY, ) -> list[tuple[int, ...]] | None: """Returns only the coord list for the best candidate box. @@ -1051,6 +1166,7 @@ def find_best_candidate_coords( required_devices: Number of devices needed for one mesh. candidate_shapes: Optional exact physical shapes to scan instead of enumerating every factorization of `required_devices`. + allocation_policy: Ranking policy for legal candidate boxes. Returns: The coord list for the best-ranked candidate box, or ``None`` when no @@ -1062,13 +1178,16 @@ def find_best_candidate_coords( the winning box start and shape, which the region-aware allocator still needs in order to split remaining regions. """ - shapes = candidate_shapes or _enumerate_box_shapes( - required_devices, - coord_topology.max_shape, - ) + shapes = candidate_shapes + if shapes is None: + shapes = _enumerate_box_shapes( + required_devices, + coord_topology.max_shape, + ) best_candidate = _find_best_candidate_box( coord_topology, shapes, + allocation_policy=allocation_policy, ) if best_candidate is None: return None @@ -1116,31 +1235,40 @@ def _allocate_devices_by_coords( return None, None allocation_policy = _normalize_allocation_policy(allocation_policy) - regions = tuple(coord_regions or (_full_coord_region(coord_topology),)) + regions = ( + (_full_coord_region(coord_topology),) + if coord_regions is None + else tuple(coord_regions) + ) candidate_regions = [] for region in regions: candidate_shapes = _supported_coord_box_shapes( devices, coord_topology, required_devices, + allocation_policy=allocation_policy, available_coord_shape=region.shape, ) if not candidate_shapes: continue - candidate_regions.append((region, candidate_shapes)) - - candidate_regions = _order_candidate_regions( - candidate_regions, allocation_policy - ) - for region, candidate_shapes in candidate_regions: best_region_candidate = _find_best_candidate_box( coord_topology, candidate_shapes, + allocation_policy=allocation_policy, region=region, ) if best_region_candidate is None: continue - candidate_start, candidate_shape, candidate_coords = best_region_candidate + start, shape, coords = best_region_candidate + candidate_regions.append((region, start, shape, coords)) + + candidate_regions = _order_region_candidates( + candidate_regions, allocation_policy + ) + if candidate_regions: + region, candidate_start, candidate_shape, candidate_coords = ( + candidate_regions[0] + ) selected_coords = set(candidate_coords) assigned_devices = [ device @@ -1160,17 +1288,35 @@ def _allocate_devices_by_coords( next_coord_regions = tuple(next_coord_regions_list) return assigned_devices, next_coord_regions + # This fallback runs only after every tracked remaining region failed to + # realize a valid box for this request. That can happen because the + # guillotine-style region bookkeeping is conservative: the remaining devices + # may still contain a contiguous supported box that spans the boundaries of + # the tracked regions even though no single tracked region can represent it. candidate_shapes = _supported_coord_box_shapes( devices, coord_topology, required_devices, + allocation_policy=allocation_policy, ) best_candidate_coords = find_best_candidate_coords( coord_topology, required_devices, candidate_shapes=candidate_shapes, + allocation_policy=allocation_policy, ) if best_candidate_coords is None: + logging.warning( + "Coord allocation could not construct a valid box: " + "required_devices=%d allocation_policy=%s topology_shape=%s " + "remaining_regions=%s supported_shapes=%s device_sample=%s.", + required_devices, + allocation_policy, + coord_topology.max_shape, + [{"start": region.start, "shape": region.shape} for region in regions], + candidate_shapes, + summarize_devices_for_logging(devices, debug=True, limit=8), + ) return None, tuple(regions) selected_coords = set(best_candidate_coords) @@ -1288,10 +1434,18 @@ def _allocate_devices_from_pool( allocation_policy, ) if assigned_devices is None: + remaining_region_summary = ( + None + if coord_regions is None + else [{"start": region.start, "shape": region.shape} for region in coord_regions] + ) raise ValueError( f"Mesh allocation requires {required_devices} devices for {mesh_name}, " "but coord-based allocation could not construct a valid box from the " - "remaining devices." + f"remaining devices (allocation_policy={allocation_policy}, " + f"remaining_device_count={len(remaining_devices)}, " + f"remaining_regions={remaining_region_summary}, " + f"device_sample={summarize_devices_for_logging(remaining_devices, debug=True, limit=8)})." ) remaining_devices = _remove_devices_by_identity( diff --git a/tunix/utils/topology.py b/tunix/utils/topology.py index 4d2c96f46..3a3016346 100644 --- a/tunix/utils/topology.py +++ b/tunix/utils/topology.py @@ -95,34 +95,107 @@ def _best_single_host_fish_shape( Returns: The best fitting within-host shape, or None when none holds the count. """ + supported_shapes = _supported_single_host_fish_shapes( + required_chips, available_chip_shape + ) + if not supported_shapes: + return None + return supported_shapes[0] + + +def _supported_single_host_fish_shapes( + required_chips: int, + available_chip_shape: Sequence[int] | None, +) -> list[tuple[int, int, int]]: + """Returns all exact single-host fish shapes ranked by compactness. + + Args: + required_chips: Number of chips requested. + available_chip_shape: Optional 3D per-axis bound; the per-host bound is + further capped to it. When its rank is not 3, no shape is returned. + + Returns: + All matching shapes sorted from more cubical to less cubical. + """ host_shape = _MULTI_HOST_BOUNDS if available_chip_shape is not None: if len(available_chip_shape) != 3: - return None + return [] host_shape = tuple( min(int(limit), host_limit) for limit, host_limit in zip(available_chip_shape, _MULTI_HOST_BOUNDS) ) - best_shape = None - best_key = None + supported_shapes = [] for x in range(1, host_shape[0] + 1): for y in range(1, host_shape[1] + 1): for z in range(1, host_shape[2] + 1): shape = (x, y, z) if math.prod(shape) != required_chips: continue - # Pack into earlier axes first: prefer larger x, then y, then z. - key = ( - max(shape), - tuple(sorted(shape, reverse=True)), - tuple(-dim for dim in shape), - shape, - ) - if best_key is None or key < best_key: - best_shape = shape - best_key = key - return best_shape + supported_shapes.append(shape) + + return sorted( + supported_shapes, + key=lambda shape: ( + max(shape), + tuple(sorted(shape, reverse=True)), + tuple(-dim for dim in shape), + shape, + ), + ) + + +def _supported_single_host_edge_shapes( + required_chips: int, + available_chip_shape: Sequence[int] | None, +) -> list[tuple[int, int, int]]: + """Returns all exact single-host edge shapes ranked by compactness. + + Edge families expose a 2D chip torus, but small requests may still occupy a + fraction of one host. This enumerates all 2D rectangular factorizations of + ``required_chips`` that fit within the per-host ``2x2`` bound, then + canonicalizes them to ``(x, y, z=1)``. + + Args: + required_chips: Number of chips requested. + available_chip_shape: Optional 2D/3D per-axis bound; when malformed it is + ignored to preserve the existing edge-family behavior. + + Returns: + All matching shapes sorted from more cubical to less cubical, breaking ties + by preferring earlier axes first. + """ + host_shape = _MULTI_HOST_BOUNDS + if available_chip_shape is not None: + canonical_available_shape = _canonicalize_chip_shape_to_3d( + available_chip_shape + ) + if canonical_available_shape is not None: + host_shape = tuple( + min(int(limit), host_limit) + for limit, host_limit in zip( + canonical_available_shape, _MULTI_HOST_BOUNDS + ) + ) + + supported_shapes = [] + for x in range(1, host_shape[0] + 1): + for y in range(1, host_shape[1] + 1): + shape = (x, y, 1) + if math.prod(shape) != required_chips: + continue + supported_shapes.append(shape) + + return sorted( + supported_shapes, + key=lambda shape: ( + max(shape), + tuple(sorted(shape, reverse=True)), + tuple(-dim for dim in shape), + shape, + ), + ) def _topology_shape_sort_key( @@ -356,13 +429,40 @@ def _best_fish_cube_shape( The most cubical valid cube-multiple shape, or None when `required_chips` is below one full cube or no arrangement fits the available bound. + Raises: + ValueError: If `required_chips` is at or above one full cube but not an + exact multiple of the cube volume. + """ + supported_shapes = _supported_fish_cube_shapes( + required_chips, available_chip_shape + ) + if not supported_shapes: + return None + return supported_shapes[0] + + +def _supported_fish_cube_shapes( + required_chips: int, + available_chip_shape: Sequence[int] | None, +) -> list[tuple[int, int, int]]: + """Returns all valid fish-family cube-multiple shapes, best first. + + Args: + required_chips: Total chip count requested. + available_chip_shape: Optional per-axis chip bound the shape must fit + within. When None, only the cube-count constraint applies. Expected to be + a 3-tuple when provided. + + Returns: + All valid shapes sorted from more cubical to less cubical. + Raises: ValueError: If `required_chips` is at or above one full cube but not an exact multiple of the cube volume. """ cube_volume = _FISH_CUBE_GRANULARITY**3 if required_chips < cube_volume: - return None + return [] cube_units, remainder = divmod(required_chips, cube_volume) if remainder != 0: @@ -382,8 +482,7 @@ def _best_fish_cube_shape( # once; `i * i <= cube_units` and `i * j <= cube_units` bound the loops to # that ordering (once i exceeds the cube root, no valid j >= i remains). The # actual shape multiplies each factor by the 4-chip cube edge. - best_shape = None - best_shape_key = None + supported_shapes = [] i = 1 while i <= max_i and i * i <= cube_units: j = i @@ -395,28 +494,24 @@ def _best_fish_cube_shape( _FISH_CUBE_GRANULARITY * j, _FISH_CUBE_GRANULARITY * int(k), ) - shape_key = _topology_shape_sort_key(shape) - if best_shape_key is None or shape_key < best_shape_key: - best_shape = shape - best_shape_key = shape_key + supported_shapes.append(shape) j += 1 i += 1 - return best_shape + return sorted(supported_shapes, key=_topology_shape_sort_key) -def best_topology_shapes_for_chip_count( +def supported_topology_shapes_for_chip_count( device_kind_or_family: str, required_chips: int, *, chip_rank: int = 3, available_chip_shape: Sequence[int] | None = None, ) -> list[tuple[int, ...]]: - """Returns the best legal topology shape(s) for a requested chip count. + """Returns all legal topology shapes for a requested chip count. - Shapes are ranked from more cubical to less cubical. For fish families this - helper returns only the best-ranked 3D shape. For edge families it returns the - single supported shape for the count, projected to the requested - ``chip_rank``. + Shapes are ranked from more cubical to less cubical. This is the full-shape + counterpart to `best_topology_shapes_for_chip_count()`, which returns only + the first entry. Args: device_kind_or_family: Either a raw device kind (``"TPU v5 lite"``) or an @@ -424,32 +519,15 @@ def best_topology_shapes_for_chip_count( required_chips: Number of physical chips the shape must contain. chip_rank: Rank of the returned shapes. Edge families support 2 or 3; fish families support only 3. - available_chip_shape: Optional per-axis chip bound the shape must fit within - (e.g. the remaining region of a partially consumed slice). + available_chip_shape: Optional per-axis chip bound the shape must fit + within. Returns: - A list with the best shape, or an empty list when no supported shape matches - the count, rank, or available bound (or the family is unknown). + A possibly-empty list of all legal shapes, sorted best-first. Raises: ValueError: If a fish-family request at or above ``4x4x4`` is not divisible by the cube granularity volume. - - Examples: - Edge family, 2D vs 3D projection:: - - best_topology_shapes_for_chip_count("TPU v6e", 8, chip_rank=2) # [(2, - 4)] - best_topology_shapes_for_chip_count("TPU v6e", 8, chip_rank=3) # [(2, - 4, 1)] - - Fish family picks the most cubical arrangement, and honors the bound:: - - best_topology_shapes_for_chip_count("TPU v7", 256) # [(4, - 8, 8)] - best_topology_shapes_for_chip_count( - "TPU v7", 576, available_chip_shape=(4, 12, 16)) # [(4, - 12, 12)] """ if required_chips <= 0: return [] @@ -463,23 +541,26 @@ def best_topology_shapes_for_chip_count( parsed_available_shape = tuple(available_chip_shape) if family in _EDGE_FAMILIES: - shape = _EDGE_SHAPE_BY_CHIP_COUNT.get(required_chips) - if shape is None: - return [] - if parsed_available_shape is not None: - canonical_edge_available_shape = _canonicalize_chip_shape_to_3d( - parsed_available_shape - ) - # A malformed available shape (not 2D/3D) canonicalizes to None and is - # treated as no constraint, matching the prior behavior. - if canonical_edge_available_shape is not None and not _shape_fits_within( - shape, canonical_edge_available_shape - ): + supported_shapes = _supported_single_host_edge_shapes( + required_chips, parsed_available_shape + ) + if not supported_shapes: + shape = _EDGE_SHAPE_BY_CHIP_COUNT.get(required_chips) + if shape is None: return [] + if parsed_available_shape is not None: + canonical_edge_available_shape = _canonicalize_chip_shape_to_3d( + parsed_available_shape + ) + if canonical_edge_available_shape is not None and not _shape_fits_within( + shape, canonical_edge_available_shape + ): + return [] + supported_shapes = [shape] if chip_rank == 2: - return [shape[:2]] + return [shape[:2] for shape in supported_shapes] if chip_rank == 3: - return [shape] + return supported_shapes return [] if chip_rank != 3: @@ -488,11 +569,6 @@ def best_topology_shapes_for_chip_count( if parsed_available_shape is not None and len(parsed_available_shape) != 3: return [] - # Exactly one chip-count range applies, so each branch already yields its own - # best shape -- no cross-branch ranking is needed. - # * below the first full cube, each sub-cube volume maps to one shape; - # * a sub-host fraction (1-2 chips) is searched within the per-host bound; - # * at or above the first full cube, shapes are 4i x 4j x 4k multiples. sub_cube_shape = _FISH_SUB_CUBE_SHAPE_BY_CHIP_COUNT.get(required_chips) if sub_cube_shape is not None and ( parsed_available_shape is None @@ -500,14 +576,72 @@ def best_topology_shapes_for_chip_count( ): return [sub_cube_shape] - single_host_shape = _best_single_host_fish_shape( + single_host_shapes = _supported_single_host_fish_shapes( required_chips, parsed_available_shape ) - if single_host_shape is not None: - return [single_host_shape] + if single_host_shapes: + return single_host_shapes - cube_shape = _best_fish_cube_shape(required_chips, parsed_available_shape) - if cube_shape is not None: - return [cube_shape] + cube_shapes = _supported_fish_cube_shapes(required_chips, parsed_available_shape) + if cube_shapes: + return cube_shapes return [] + + +def best_topology_shapes_for_chip_count( + device_kind_or_family: str, + required_chips: int, + *, + chip_rank: int = 3, + available_chip_shape: Sequence[int] | None = None, +) -> list[tuple[int, ...]]: + """Returns the best legal topology shape(s) for a requested chip count. + + Shapes are ranked from more cubical to less cubical. For fish families this + helper returns only the best-ranked 3D shape. For edge families it returns the + single supported shape for the count, projected to the requested + ``chip_rank``. + + Args: + device_kind_or_family: Either a raw device kind (``"TPU v5 lite"``) or an + already-normalized family key (``"v5e"``). + required_chips: Number of physical chips the shape must contain. + chip_rank: Rank of the returned shapes. Edge families support 2 or 3; fish + families support only 3. + available_chip_shape: Optional per-axis chip bound the shape must fit within + (e.g. the remaining region of a partially consumed slice). + + Returns: + A list with the best shape, or an empty list when no supported shape matches + the count, rank, or available bound (or the family is unknown). + + Raises: + ValueError: If a fish-family request at or above ``4x4x4`` is not divisible + by the cube granularity volume. + + Examples: + Edge family, 2D vs 3D projection:: + + best_topology_shapes_for_chip_count("TPU v6e", 8, chip_rank=2) # [(2, + 4)] + best_topology_shapes_for_chip_count("TPU v6e", 8, chip_rank=3) # [(2, + 4, 1)] + + Fish family picks the most cubical arrangement, and honors the bound:: + + best_topology_shapes_for_chip_count("TPU v7", 256) # [(4, + 8, 8)] + best_topology_shapes_for_chip_count( + "TPU v7", 576, available_chip_shape=(4, 12, 16)) # [(4, + 12, 12)] + """ + supported_shapes = supported_topology_shapes_for_chip_count( + device_kind_or_family, + required_chips, + chip_rank=chip_rank, + available_chip_shape=available_chip_shape, + ) + if not supported_shapes: + return [] + return [supported_shapes[0]]