jax-ml/jax · error · RuntimeError

Preemption sync manager should only be initialized once.

Error message

Preemption sync manager should only be initialized once.

What it means

Raised when initialize_preemption_sync_manager() is called after a preemption sync manager already exists for this process. The preemption sync manager (used to coordinate graceful shutdown across distributed workers) is a singleton bound to the distributed client.

Source

Thrown at jax/_src/distributed.py:258

      # client because the preemption sync manager depends on the client.
      self.preemption_sync_manager.shutdown()
      self.preemption_sync_manager = None
    if self.client:
      self.client.shutdown()
      self.client = None
    if self.service:
      self.service.shutdown()
      self.service = None

  def initialize_preemption_sync_manager(self):
    if not _ENABLE_PREEMPTION_SERVICE.value:
      logger.info(
          'The JAX preemption service is disabled. You can enable it using the'
          ' jax_enable_preemption_service configuration option.'
      )
      return
    if self.preemption_sync_manager is not None:
      raise RuntimeError(
          'Preemption sync manager should only be initialized once.')
    self.preemption_sync_manager = (
        _jax.create_preemption_sync_manager())
    assert self.client is not None
    self.preemption_sync_manager.initialize(self.client)

global_state = State()

def initialize(coordinator_address: str | None = None,
               num_processes: int | None = None,
               process_id: int | None = None,
               local_device_ids: int | Sequence[int] | None = None,
               cluster_detection_method: str | None = None,
               initialization_timeout: int = 300,
               heartbeat_timeout_seconds: int = 100,
               shutdown_timeout_seconds: int = 300,
               coordinator_bind_address: str | None = None,
               slice_index: int | None = None,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Initialize distributed/preemption only once per process; remove duplicate calls
2. Restart the process/kernel to reset the singleton state
3. Check `global_state.preemption_sync_manager is None` before initializing

Example fix

// before
initialize()  # preemption enabled
initialize()  # RuntimeError

// after
from jax._src.distributed import global_state
if global_state.preemption_sync_manager is None:
    initialize()

from jax._src.distributed import global_state
if global_state.preemption_sync_manager is None:
    initialize_preemption_sync_manager()

try:
    initialize_preemption_sync_manager()
except RuntimeError as e:
    if 'initialized once' not in str(e): raise

When it happens

Trigger: Calling jax.distributed.initialize() with preemption service enabled twice, or calling the internal initialize_preemption_sync_manager() a second time in one process.

Common situations: Re-running distributed setup in notebooks or test harnesses; combining a launcher that enables jax_enable_preemption_service with user code that also initializes it.

Related errors

• Mapped away dimension of inputs passed to vmap should be sha
• The 'sharding' attribute is not available on {self._error_re
• The is_fully_addressable property was called on {self._error
• mTLS for the JAX distributed service requires jaxlib 0.11.2
• coordinator_address should be defined.

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/3d747f105826a7e4. Report an issue: GitHub.