Skip to content

Make elastic_retry the single entry point for every elastic mode. #337

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Make elastic_retry the single entry point for every elastic mode. #337

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 15 additions & 7 deletions pathwaysutils/elastic/manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,14 @@ class ScaleUpSignalError(Exception):


def _elastic_event_cleanup() -> None:
"""Cleans up JAX profiles, caches, and live arrays."""
"""Stops any ongoing JAX profiler trace.

Live arrays are not deleted and the JAX caches are not cleared: JAX frees a
buffer once its last reference is dropped, which happens when the failed
attempt unwinds, and compiled programs for slices that went away are never
called again. Anything that must not keep the previous attempt's arrays
alive, such as a snapshot still being copied, is the caller's responsibility.
"""
try:
_logger.info("Cleaning up any ongoing traces")
jax.profiler.stop_trace()
Expand All @@ -62,10 +69,6 @@ def _elastic_event_cleanup() -> None:
_logger.exception("Error cleaning up ongoing traces")
raise

jax.clear_caches()
for array in jax.live_arrays():
array.delete()


class ElasticRetryLimit:
"""A retry callback that limits the number of attempts."""
Expand Down Expand Up @@ -299,8 +302,7 @@ def elastic_retry(

This decorator wraps a function to automatically retry execution in case of
`jax.errors.JaxRuntimeError` caused by slice down events. It waits for
`minimum_slice_count` active slices before each attempt and cleans up JAX
caches on failure.
`minimum_slice_count` active slices before each attempt.

If `minimum_slice_count` is not met, the function will wait until at least
`minimum_slice_count` slices are active before execution. If
Expand All @@ -313,6 +315,12 @@ def elastic_retry(
boundaries) and raise a `ScaleUpSignalError` to gracefully interrupt the
current execution and trigger a retry with the expanded hardware.

On an elastic event the arrays of the failed attempt are not deleted
explicitly; JAX frees them once the attempt has unwound and nothing
references them anymore. The wrapped function and the callbacks must not
keep references to them, or the memory stays allocated during the next
attempt.

Often, the function will dispatch JAX operations and wait for them to
complete while creating a log message. If using Python logging, it is
recommended to set `logging.raiseExceptions=True` to ensure that the
Expand Down
Loading