# Copyright 2021 The JAX Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

from collections.abc import Sequence
import logging
import os
from typing import Any
import warnings

from jax._src import clusters
from jax._src import config
from jax._src import xla_bridge
from jax._src.lib import _jax

logger = logging.getLogger(__name__)


_CHECK_PROXY_ENVS = config.bool_flag(
    name="jax_check_proxy_envs",
    default=True,
    help="Checks proxy vars in user envs and emit warnings.",
)


_ENABLE_RECOVERABILITY = config.bool_state(
    name="jax_enable_recoverability",
    default=False,
    help=(
        "Allows a multi-controller JAX job to continue running, even after some"
        " tasks have failed."
    ),
)

_ENABLE_PREEMPTION_SERVICE = config.bool_state(
    name='jax_enable_preemption_service',
    default=True,
    help=(
        "Enables the preemption service. See"
        " multihost_utils.reached_preemption_sync_point for details."
    ),
)

class State:
  process_id: int = 0
  num_processes: int = 1
  service: _jax.DistributedRuntimeService | Any | None = None
  client: _jax.DistributedRuntimeClient | Any | None = None
  preemption_sync_manager: Any | None = None
  coordinator_address: str | None = None
  partition_index: int | None = None

  def initialize(self,
                 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,
                 coordinator_bind_address: str | None = None,
                 heartbeat_timeout_seconds: int = 100,
                 shutdown_timeout_seconds: int = 300,
                 partition_index: int | None = None):
    coordinator_address = (coordinator_address or
                           os.environ.get('JAX_COORDINATOR_ADDRESS'))
    if isinstance(local_device_ids, int):
      local_device_ids = [local_device_ids]

    if local_device_ids is None and (env_ids := os.environ.get('JAX_LOCAL_DEVICE_IDS')):
      local_device_ids = list(map(int, env_ids.split(",")))

    if (cluster_detection_method != 'deactivate' and
        None in (coordinator_address, num_processes, process_id, local_device_ids)):
      (coordinator_address, num_processes, process_id, local_device_ids) = (
          clusters.ClusterEnv.auto_detect_unset_distributed_params(
              coordinator_address,
              num_processes,
              process_id,
              local_device_ids,
              cluster_detection_method,
              initialization_timeout,
          )
      )

    if coordinator_address is None:
      raise ValueError('coordinator_address should be defined.')
    if num_processes is None:
      raise ValueError('Number of processes must be defined.')
    if process_id is None:
      raise ValueError('The process id of the current process must be defined.')
    if not isinstance(process_id, int):
      raise TypeError("process_id must be a nonnegative int. "
                      f"Got process_id={process_id} of type {type(process_id)}.")
    if not isinstance(num_processes, int):
      raise TypeError("num_processes must be a positive int. "
                      f"Got num_processes={num_processes} of type {type(num_processes)}.")
    if not (0 <= process_id < num_processes):
      raise ValueError("process_id and num_processes must be nonnegative, with process_id < num_processes. "
                       f"Got process_id={process_id}, num_processes={num_processes}.")

    self.coordinator_address = coordinator_address

    # The default value of [::]:port tells the coordinator to bind to all
    # available addresses on the same port as coordinator_address.
    default_coordinator_bind_address = '[::]:' + coordinator_address.rsplit(':', 1)[1]
    coordinator_bind_address = (coordinator_bind_address or
                                os.environ.get('JAX_COORDINATOR_BIND_ADDRESS',
                                               default_coordinator_bind_address))
    if coordinator_bind_address is None:
      raise ValueError('coordinator_bind_address should be defined.')

    if local_device_ids:
      visible_devices = ','.join(str(x) for x in local_device_ids)
      logger.info('JAX distributed initialized with visible devices: %s', visible_devices)
      config.update("jax_cuda_visible_devices", visible_devices)
      config.update("jax_rocm_visible_devices", visible_devices)

    self.process_id = process_id

    proxy_vars = []
    if _CHECK_PROXY_ENVS.value:
      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)

    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)

    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(
        coordinator_address, process_id, init_timeout=initialization_timeout,
        use_compression=True, heartbeat_timeout=heartbeat_timeout_seconds,
        recoverable=_ENABLE_RECOVERABILITY.value)  # type: ignore
    logger.info('Connecting to JAX distributed service on %s', coordinator_address)
    self.client.connect()

    self.initialize_preemption_sync_manager()

    if partition_index is None:
      jax_partition_index = os.environ.get('JAX_PARTITION_INDEX')
      jax_slice_index = os.environ.get('JAX_SLICE_INDEX')
      if jax_partition_index is not None:
        partition_index = int(jax_partition_index)  # type: ignore
      elif jax_slice_index is not None:
        # Deprecation added 2025-08-05. Should be removed after 3 months.
        warnings.warn(
            'JAX_SLICE_INDEX has been deprecated. Please use'
            ' JAX_PARTITION_INDEX instead.',
            DeprecationWarning,
        )
        partition_index = int(jax_slice_index)  # type: ignore
    self.partition_index = partition_index

  def shutdown(self):
    if self.preemption_sync_manager:
      # It's important to shut down the preemption sync manager before the
      # 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())
    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,
               partition_index: int | None = None):
  """Initializes the JAX distributed system.

  Calling :func:`~jax.distributed.initialize` prepares JAX for execution on
  multi-host GPU and Cloud TPU. :func:`~jax.distributed.initialize` must be
  called before performing any JAX computations.

  The JAX distributed system serves a number of roles:

    * It allows JAX processes to discover each other and share topology information,
    * It performs health checking, ensuring that all processes shut down if any process dies, and
    * It is used for distributed checkpointing.

  If you are using TPU, Slurm, or Open MPI, all arguments are optional: if omitted, they
  will be chosen automatically.

  The ``cluster_detection_method`` may be used to choose a specific method for detecting those
  distributed arguments. You may pass any of the automatic ``spec_detect_methods`` to this
  argument though it is not necessary in the TPU, Slurm, or Open MPI cases.  For other MPI
  installations, if you have a functional ``mpi4py`` installed, you may pass
  ``cluster_detection_method="mpi4py"`` to bootstrap the required arguments.

  Otherwise, you must provide the ``coordinator_address``,
  ``num_processes``, ``process_id``, and ``local_device_ids`` arguments
  to :func:`~jax.distributed.initialize`. When all four arguments are provided, cluster
  environment auto detection will be skipped.

  Please note: on some systems, particularly HPC clusters that only access external networks
  through proxy variables such as HTTP_PROXY, HTTPS_PROXY, etc., the call to
  :func:`~jax.distributed.initialize` may timeout.  You may need to unset these variables
  prior to application launch.

  Args:
    coordinator_address: the IP address of process `0` and a port on which that
      process should launch a coordinator service. The choice of
      port does not matter, so long as the port is available on the coordinator
      and all processes agree on the port.
      May be ``None`` only on supported environments, in which case it will be chosen automatically.
      Note that special addresses like ``localhost`` or ``127.0.0.1`` usually mean that the program
      will bind to a local interface and are not suitable when running in a multi-host environment.
    num_processes: Number of processes. May be ``None`` only on supported environments, in
      which case it will be chosen automatically.
    process_id: The ID number of the current process. The ``process_id`` values across
      the cluster must be a dense range ``0``, ``1``, ..., ``num_processes - 1``.
      May be ``None`` only on supported environments; if ``None`` it will be chosen automatically.
    local_device_ids: Restricts the visible devices of the current process to ``local_device_ids``.
      If ``None``, defaults to all local devices being visible to the process except when processes
      are launched via Slurm and Open MPI on GPUs. In that case, it will default to a single device per process.
    cluster_detection_method: An optional string to attempt to autodetect the configuration of the distributed
      run.  Note that "mpi4py" method requires you to have a working ``mpi4py`` install in your environment,
      and launch the applicatoin with an MPI-compatible job launcher such as ``mpiexec`` or ``mpirun``.
      Legacy auto-detect options "ompi" (OMPI) and "slurm" (Slurm) remain enabled. "deactivate" bypasses
      automatic cluster detection.
    initialization_timeout: Time period (in seconds) for which connection will
      be retried. If the initialization takes more than the timeout specified,
      the initialization will error. Defaults to 300 secs i.e. 5 mins.
    heartbeat_timeout_seconds: The time (in seconds) after which a process is
      considered dead if it hasn't successfully sent any heartbeats. Defaults
      to 100 seconds.
    shutdown_timeout_seconds: The time (in seconds) a terminating process will
      wait for all other processes to also terminate. Defaults to 300 seconds.
    coordinator_bind_address: the address and port to which the coordinator service
      on process `0` should bind. If this is not specified, the default is to bind to
      all available addresses on the same port as ``coordinator_address``. On systems
      that have multiple network interfaces per node it may be insufficient to only
      have the coordinator service listen on one address/interface.
    slice_index: DEPRECATED: Use ``partition_index`` instead.
    partition_index: The partition index assigned to this process' local devices. If any process sets ``partition_index``,
      then all processes must do so. If ``None`` the partition indices will be chosen automatically.

  Raises:
    RuntimeError: If :func:`~jax.distributed.initialize` is called more than once
      or if called after the backend is already initialized.

  Examples:

  Suppose there are two GPU processes, and process 0 is the designated coordinator
  with address ``10.0.0.1:1234``. To initialize the GPU cluster, run the
  following commands before anything else.

  On process 0:

  >>> jax.distributed.initialize(coordinator_address='10.0.0.1:1234', num_processes=2, process_id=0)  # doctest: +SKIP

  On process 1:

  >>> jax.distributed.initialize(coordinator_address='10.0.0.1:1234', num_processes=2, process_id=1)  # doctest: +SKIP
  """
  if xla_bridge.backends_are_initialized():
    raise RuntimeError("jax.distributed.initialize() must be called before "
                        "any JAX calls that might initialise the XLA backend. "
                        "This includes any computation, but also calls to jax.devices, jax.device_put, and others.")
  if partition_index is None:
    if slice_index is not None:
      # Deprecation added 2025-08-05. Should be removed after 3 months.
      warnings.warn(
          '`slice_index` has been deprecated. Please use `partition_index` instead.',
          DeprecationWarning,
      )
    partition_index = slice_index
  global_state.initialize(coordinator_address, num_processes, process_id,
                          local_device_ids, cluster_detection_method,
                          initialization_timeout, coordinator_bind_address,
                          heartbeat_timeout_seconds=heartbeat_timeout_seconds,
                          shutdown_timeout_seconds=shutdown_timeout_seconds,
                          partition_index=partition_index)


def is_initialized() -> bool:
  """Check if the JAX distributed system is initialized."""
  return global_state.client is not None

def shutdown():
  """Shuts down the distributed system.

  Does nothing if the distributed system is not running.
  """
  global_state.shutdown()