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.

Solutions

  1. Guard with the global state: if jax._src.distributed.global_state.service is None: initialize(...) — or simply call shutdown() before re-initializing
  2. Call jax.distributed.shutdown() before a second initialize in rerunnable scripts/notebooks
  3. 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

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.

Appendix: 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)