From bd4589ae53e79ddeb004796b98b3ec32e13e1a6e Mon Sep 17 00:00:00 2001 From: Luke Baumann Date: Fri, 23 Oct 2026 22:31:29 -0700 Subject: [PATCH] Remove unused _cleanup_on_retry helper in elastic manager. PiperOrigin-RevId: 1006105618 --- pathwaysutils/elastic/manager.py | 14 -------------- 1 file changed, 14 deletions(-) diff --git a/pathwaysutils/elastic/manager.py b/pathwaysutils/elastic/manager.py index 7b4dd4a..01bc100 100644 --- a/pathwaysutils/elastic/manager.py +++ b/pathwaysutils/elastic/manager.py @@ -214,20 +214,6 @@ def scale_by_active_slices(self, x: int | float) -> int | float: else: raise ValueError(f"Unsupported type: {type(x)=}") - def _cleanup_on_retry(self): - """Cleans up JAX caches and traces on retry.""" - try: - _logger.debug("Cleaning up any ongoing traces") - jax.profiler.stop_trace() - except (RuntimeError, ValueError): - _logger.debug("No ongoing traces to clean up") - except Exception: # pylint: disable=broad-exception-caught - _logger.exception("Error cleaning up ongoing traces") - - jax.clear_caches() - for array in jax.live_arrays(): - array.delete() - def _check_inactive_slices(self) -> None: """Checks inactive slices and updates available_inactive_slices.""" if not self.inactive_slice_indices: