merge luke new commit

Author: Sharon Yu

pathwaysutils/elastic/manager.py

@@ -43,7 +43,7 @@ class NewSliceAvailableError(RuntimeError):
 _F = TypeVar("_F", bound=Callable[..., Any])
 
 
-def _elastic_event_cleanup() -> None:
+def _elastic_event_cleanup():
   """Cleans up JAX profiles, caches, and live arrays."""
   try:
     _logger.info("Cleaning up any ongoing traces")
@@ -149,23 +149,11 @@ def _elasticity_retry_decorator(
       pre_callback: Callable[..., Any] | None = None,
       on_elastic_event_callback: Callable[..., Any] | None = None,
   ) -> Callable[[_F], _F]:
-    """Retries a function with elasticity fault tolerance.
+    """Retries a function with elasticity fault tolerance."""
 
-    Args:
-      max_retries: The maximum number of times to retry the function.
-      pre_callback: A callback to call before each attempt of the wrapped
-        function.
-      on_elastic_event_callback: A callback to call after an elastic failure
-        occurs.
-
-    Returns:
-      A function decorator.
-    """
-    if max_retries <= 0:
-      raise ValueError("max_retries must be positive.")
-    def decorator(func: _F) -> _F:
+    def decorator(func):
       @functools.wraps(func)
-      def wrapper(*args: Any, **kwargs: Any) -> Any:
+      def wrapper(*args, **kwargs):
         for retry_index in range(max_retries):
           try:
             _logger.info(
@@ -202,10 +190,6 @@ def wrapper(*args: Any, **kwargs: Any) -> Any:
         )
 
       return wrapper
-    else:
-          raise ElasticRuntimeError(
-              f"Elastic attempt {max_retries} out of {max_retries} failed."
-          )
 
     return decorator