From fbd93d75f1bacc9cdacd8d9f9592a2701509e1ff Mon Sep 17 00:00:00 2001 From: Luke Baumann Date: Fri, 23 Oct 2026 22:31:28 -0700 Subject: [PATCH] Validate wait_for_slices parameters and clean up monitor thread state in elastic manager. PiperOrigin-RevId: 1006105615 --- pathwaysutils/elastic/elastic.py | 4 ++++ pathwaysutils/elastic/manager.py | 5 ++++- 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/pathwaysutils/elastic/elastic.py b/pathwaysutils/elastic/elastic.py index 1ef4bf8..509c52a 100644 --- a/pathwaysutils/elastic/elastic.py +++ b/pathwaysutils/elastic/elastic.py @@ -198,6 +198,10 @@ def wait_for_slices( TimeoutError: If the timeout is reached before the slices become active. """ + if poll_interval < 0: + raise ValueError(f"poll_interval must be non-negative, got {poll_interval}") + if timeout is not None and timeout < 0: + raise ValueError(f"timeout must be non-negative, got {timeout}") if slice_to_devices is None: _logger.debug("slice_to_devices is None. Getting from jax.devices().") slice_to_devices = get_slice_to_devices(jax.devices()) diff --git a/pathwaysutils/elastic/manager.py b/pathwaysutils/elastic/manager.py index 7b4dd4a..ca8b5af 100644 --- a/pathwaysutils/elastic/manager.py +++ b/pathwaysutils/elastic/manager.py @@ -158,7 +158,10 @@ def close(self) -> None: pass else: raise - self._monitor_thread = None + finally: + self._monitor_thread = None + self._stop_event = None + else: self._stop_event = None @functools.cached_property