From 58ba851ba22edffe3817d2014aaf8e71b5f0f2e1 Mon Sep 17 00:00:00 2001 From: Luke Baumann Date: Sat, 25 Jul 2026 22:00:11 -0700 Subject: [PATCH] Fix active slice pruning and dynamic inactive slice tracking in the elastic manager. PiperOrigin-RevId: 954057947 --- pathwaysutils/elastic/manager.py | 34 +++++++++++++++++--------------- 1 file changed, 18 insertions(+), 16 deletions(-) diff --git a/pathwaysutils/elastic/manager.py b/pathwaysutils/elastic/manager.py index e7b51f1..62995aa 100644 --- a/pathwaysutils/elastic/manager.py +++ b/pathwaysutils/elastic/manager.py @@ -116,11 +116,12 @@ def __init__(self, devices: Sequence[jax.Device] | None = None) -> None: self.all_slice_indices = frozenset(self.slice_to_devices.keys()) + self._active_slice_indices = frozenset() + self.available_inactive_slices = frozenset() + self.active_slice_indices = elastic.get_active_slice_indices( slice_to_devices=self.slice_to_devices ) - self.inactive_slice_indices = self.all_slice_indices - self.active_slice_indices - self.available_inactive_slices = frozenset() self._stop_event = None self._monitor_thread = None @@ -180,22 +181,26 @@ def default_device(self) -> jax.Device: except StopIteration as error: raise ValueError("No active slices") from error + @property + def active_slice_indices(self) -> frozenset[int]: + """The indices of active slices.""" + return self._active_slice_indices + + @active_slice_indices.setter + def active_slice_indices(self, value: set[int] | frozenset[int]) -> None: + self._active_slice_indices = frozenset(value) + self.inactive_slice_indices = ( + self.all_slice_indices - self._active_slice_indices + ) + self.available_inactive_slices = frozenset( + self.available_inactive_slices - self._active_slice_indices + ) + @property def active_slice_count(self) -> int: """The number of active slices.""" return len(self.active_slice_indices) - @property - def new_slice_event(self) -> threading.Event: - """Deprecated compatibility property for un-updated MaxText code. - - TODO: b/527183831 - Remove this property once MaxText CL 2 is submitted. - """ - event = threading.Event() - if self.available_inactive_slices: - event.set() - return event - def scale_by_active_slices(self, x: int | float) -> int | float: """Scale x by the number of active slices.""" if isinstance(x, int): @@ -372,9 +377,6 @@ def attempt_execution(attempt: int) -> Any: poll_interval=poll_interval, timeout=timeout, ) - self.inactive_slice_indices = ( - self.all_slice_indices - self.active_slice_indices - ) # Reset available_inactive_slices at attempt start since # active_slice_indices has just been updated by wait_for_slices. self.available_inactive_slices = frozenset()