jax-ml/jax · error · RuntimeError

distributed.initialize should only be called once.

Error message

distributed.initialize should only be called once.

What it means

Only process 0 hosts the distributed coordination service, and it may host it exactly once per process. Calling jax.distributed.initialize again while self.service is already set raises this RuntimeError.

Source

Thrown at jax/_src/distributed.py:191

      proxy_vars = [key for key in os.environ.keys()
                    if '_proxy' in key.lower()]

    if len(proxy_vars) > 0:
      vars = " ".join(proxy_vars) + ". "
      warning = (
        f'JAX detected proxy variable(s) in the environment as distributed setup: {vars}'
        'On some systems, this may cause a hang of distributed.initialize and '
        'you may need to unset these ENV variable(s)'
      )
      logger.warning(warning)

    mtls_kwargs = _get_mtls_kwargs(
        mtls_cert_file, mtls_key_file, mtls_ca_file, mtls_peer_uri_prefix,
        verify_secure_credentials)

    if process_id == 0:
      if self.service is not None:
        raise RuntimeError('distributed.initialize should only be called once.')
      logger.info(
          'Starting JAX distributed service on %s', coordinator_bind_address
      )
      self.service = _jax.get_distributed_runtime_service(
          coordinator_bind_address,
          num_processes,
          heartbeat_timeout=heartbeat_timeout_seconds,
          shutdown_timeout=shutdown_timeout_seconds,
          recoverable=_ENABLE_RECOVERABILITY.value,
          **mtls_kwargs,
      )

    self.num_processes = num_processes

    if self.client is not None:
      raise RuntimeError('distributed.initialize should only be called once.')

    self.client = _jax.get_distributed_runtime_client(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Guard with the global state: if jax._src.distributed.global_state.service is None: initialize(...) — or simply call shutdown() before re-initializing
Call jax.distributed.shutdown() before a second initialize in rerunnable scripts/notebooks
Initialize once at process startup (entrypoint) and let libraries detect existing state

Example fix

# before
jax.distributed.initialize(...)  # run twice -> RuntimeError
# after
if jax.distributed.process_id == 0 and not jax._src.distributed.global_state.service:
    jax.distributed.initialize(...)
# or: jax.distributed.shutdown(); jax.distributed.initialize(...)

Defensive patterns

Strategy: try-catch

Validate before calling

import jax._src.distributed as dist
already = dist.global_state.service is not None
if not already:
    jax.distributed.initialize(...)

Try / catch

try:
    jax.distributed.initialize(...)
except RuntimeError as e:
    if 'only be called once' in str(e):
        pass  # already initialized; safe to continue
    else:
        raise

Prevention

Initialize once at process entrypoint
Call jax.distributed.shutdown() before re-init in notebooks
Guard init behind an idempotent helper

When it happens

Trigger: Calling jax.distributed.initialize() twice in the same Python process on rank 0 — e.g. a training script that re-initializes between runs, or library + user code both initializing.

Common situations: Notebook/REPL workflows rerunning the init cell; frameworks (e.g. Flax/JAX serving stacks) that initialize internally while user code also calls it; test suites reusing the process.

Related errors

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