Skip to content

Fix reached_preemption crash when JAX preemption service is disabled#3458

Open
marsunique wants to merge 1 commit into
google:mainfrom
marsunique:preemption-bug-fix
Open

Fix reached_preemption crash when JAX preemption service is disabled#3458
marsunique wants to merge 1 commit into
google:mainfrom
marsunique:preemption-bug-fix

Conversation

@marsunique

Copy link
Copy Markdown

Description

Fixes #3406

When JAX's preemption service is disabled via jax_enable_preemption_service=False,
multihost_utils.reached_preemption_sync_point raises a RuntimeError because the
preemption sync manager is None. This causes CheckpointManager.save() to crash.

This change catches the RuntimeError in reached_preemption() and returns False,
gracefully treating a disabled preemption service as "no preemption detected".

Type of change

  • Bug fix (non-breaking change that fixes an issue)
  • New feature (non-breaking change that adds functionality)
  • Breaking change (alters existing behavior or a public API)
  • Docs / tests / internal tooling (no user-facing behavior change)

Checklist

  • I have read the contribution guidelines.
  • Tests cover this change and pass locally.
  • Public APIs and non-trivial functions are documented.
  • If this change is user-facing, I have updated CHANGELOG.md under the [Unreleased] section (or the relevant checkpoint/ / export/ changelog).

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

CheckpointManager crashes when jax_enable_preemption_service=False

1 participant