# Copyright 2024 The Orbax 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
#
#     http://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.

"""Helper for running multi-process tests."""

import os
import pathlib
import re
import signal
import subprocess
import time

from absl import app
from absl import flags
from absl.testing import absltest
import jax
from jax import config
from orbax.checkpoint._src.multihost import multihost
import portpicker

NUM_PROCESSES = flags.DEFINE_integer(
    "num_processes", None, "Number of processes to use."
)

_GPUS_PER_PROCESS = flags.DEFINE_integer(
    "gpus_per_process",
    0,
    "Number of GPUs per worker process.",
)

TPU_CHIPS_PER_PROCESS = flags.DEFINE_integer(
    "tpu_chips_per_process",
    0,
    "Number of TPU chips per worker process.",
)

_WORKER_SHUTDOWN_TIMEOUT = flags.DEFINE_integer(
    "worker_shutdown_timeout",
    15,
    "JAX shutdown timeout duration in seconds for each subprocess worker. If "
    "your test is timing out, try increasing this value.",
)

_EXTRA_TEST_ARGS = flags.DEFINE_multi_string(
    "extra_test_args", [], "Extra flags to pass to worker process."
)

# For internal use.
MULTIPROCESS_TEST_WORKER_ID = flags.DEFINE_integer(
    "multiprocess_test_worker_id",
    -1,
    "TPU worker id. Set by main test process; should not be set by users.",
)

expect_failures_with_regex = None


def main():
  # We don't call googletest.main() because if we are the main process we
  # don't want to run the tests; instead we want to fork worker processes that
  # do. However, the googletest framework would otherwise report an error in the
  # main process if googletest.main() is not called, which the following line
  # suppresses.
  config.config_with_absl()
  app.run(_main)


class GracefulKiller:
  """Add a signal handler that sets a flag if SIGINT or SIGTERM are caught."""

  # From https://stackoverflow.com/a/31464349
  kill_now = False

  def __init__(self):
    signal.signal(signal.SIGINT, self.exit_gracefully)
    signal.signal(signal.SIGTERM, self.exit_gracefully)

  def exit_gracefully(self, sig_num, unused_stack_frame):
    print(f"Caught signal: {signal.Signals(sig_num).name} ({sig_num})")
    self.kill_now = True


def _main(argv):
  if NUM_PROCESSES.value is None:
    raise ValueError("--num_processes must be specified.")

  if MULTIPROCESS_TEST_WORKER_ID.value >= 0:
    absltest.run_tests(argv, [], {})
    return

  num_processes = NUM_PROCESSES.value
  gpus_per_process = _GPUS_PER_PROCESS.value
  tpu_chips_per_process = TPU_CHIPS_PER_PROCESS.value
  num_tpu_chips = num_processes * tpu_chips_per_process
  if num_tpu_chips == 0:
    pass
  elif num_tpu_chips == 1:
    assert tpu_chips_per_process == 1
    deepsea_host_bounds = "1,1,1"
    deepsea_chips_per_host_bounds = "1,1,1"
  elif num_tpu_chips == 4:
    if tpu_chips_per_process == 1:
      deepsea_host_bounds = "2,2,1"
      deepsea_chips_per_host_bounds = "1,1,1"
    elif tpu_chips_per_process == 2:
      deepsea_host_bounds = "2,1,1"
      deepsea_chips_per_host_bounds = "1,2,1"
    elif tpu_chips_per_process == 4:
      deepsea_host_bounds = "1,1,1"
      deepsea_chips_per_host_bounds = "2,2,1"
    else:
      raise ValueError(
          "Invalid number of TPU chips per worker {}".format(
              tpu_chips_per_process
          )
      )
  elif num_tpu_chips == 8:
    if tpu_chips_per_process == 1:
      deepsea_host_bounds = "4,2,1"
      deepsea_chips_per_host_bounds = "1,1,1"
    else:
      # TODO(phawkins): implement other cases.
      raise ValueError(
          "Invalid number of TPU chips per worker {}".format(
              tpu_chips_per_process
          )
      )
  else:
    raise ValueError(f"Invalid number of TPU chips {num_tpu_chips}")

  slicebuilder_ports = [
      portpicker.pick_unused_port() for _ in range(num_processes)
  ]
  slicebuilder_addresses = ",".join(
      f"localhost:{port}" for port in slicebuilder_ports
  )
  jax_port = portpicker.pick_unused_port()

  subprocesses = []
  output_filenames = []
  output_files = []
  for i in range(num_processes):
    env = os.environ.copy()

    args = [
        "/proc/self/exe",
        f"--multiprocess_test_worker_id={i}",
        "--logtostderr",
        f"--num_processes={num_processes}",
        f"--jax_num_tasks={num_processes}",
        f"--jax_task_id={i}",
        f"--jax_controller_address=localhost:{jax_port}",
        "--jax_heartbeat_interval=1s",
        "--jax_max_missing_heartbeats=3",
        f"--jax_distributed_shutdown_timeout={_WORKER_SHUTDOWN_TIMEOUT.value}s",
        "--vmodule=client=10,service=10",
    ]
    if i == 0:
      args += [f"--jax_port={jax_port}"]

    if num_tpu_chips > 0:
      # We must set a CLOUD_TPU_TASK_ID, otherwise the TPU runtime expects to
      # be able to look up the Borg task number from the Borglet info.
      env["CLOUD_TPU_TASK_ID"] = str(i)
      chips = list(
          range(
              i * tpu_chips_per_process,
              (i + 1) * tpu_chips_per_process,
          )
      )
      excluded_chips = [d for d in range(num_tpu_chips) if d not in chips]
      args += [
          "--deepsea_hal_excluded_devs={}".format(
              ",".join(map(str, excluded_chips))
          ),
          f"--deepsea_host_bounds={deepsea_host_bounds}",
          f"--deepsea_chips_per_host_bounds={deepsea_chips_per_host_bounds}",
          f"--deepsea_slice_builder_worker_addresses={slicebuilder_addresses}",
          f"--deepsea_slice_builder_worker_port={slicebuilder_ports[i]}",
          # --jax_allow_unused_tpus suppresses a check in JAX that verifies the
          # number of JAX TPU devices found is equal to the hardware devices
          # that exist.
          "--jax_allow_unused_tpus",
      ]

    if gpus_per_process > 0:
      gpus = range(i * gpus_per_process, (i + 1) * gpus_per_process)
      env["CUDA_VISIBLE_DEVICES"] = ",".join(map(str, gpus))

    args += _EXTRA_TEST_ARGS.value

    undeclared_outputs = os.environ.get("TEST_UNDECLARED_OUTPUTS_DIR", "/tmp")
    stdout_name = f"{undeclared_outputs}/jax_{i}_stdout.log"
    stderr_name = f"{undeclared_outputs}/jax_{i}_stderr.log"
    stdout = open(stdout_name, "wb")
    stderr = open(stderr_name, "wb")
    print(f"Launching process {i}:")
    print(f"  stdout: {stdout_name}")
    print(f"  stderr: {stderr_name}")
    proc = subprocess.Popen(args, env=env, stdout=stdout, stderr=stderr)
    subprocesses.append(proc)
    output_filenames.append((stdout_name, stderr_name))
    output_files.append((stdout, stderr))

  print(" All launched, running ".center(80, "="), flush=True)

  # Wait for all the children to finish or for a SIGTERM from TAP. If we get
  # SIGTERM, we still want to collect their logs, so kill them and continue.
  killer = GracefulKiller()
  running_procs = {i: p for i, p in enumerate(subprocesses)}
  while not killer.kill_now and running_procs:
    time.sleep(0.1)
    for i, proc in list(running_procs.items()):
      if proc.poll() is not None:
        print(f"Process {i} finished.", flush=True)
        running_procs.pop(i)
  if killer.kill_now and running_procs:
    print("Caught termination, terminating remaining children.", flush=True)

    # Send a SIGTERM to each child process, to let it know it should terminate.
    for i, proc in running_procs.items():
      proc.terminate()
      print(f"Process {i} terminated.", flush=True)

    # On test timeout, Forge first sends a SIGTERM (a "soft" kill signal, that
    # the test can intercept, in order to do some cleanup, log flushing, etc).
    # After a grace period of 15 seconds, Forge sends a SIGKILL (a "hard" kill),
    # see http://yaqs/eng/q/4559876738121728#n4588728130600960. We give the
    # child process(es) a few seconds for their own cleanup, and keep the rest
    # (up to 15s) for copying the children logs into our own.
    time.sleep(5)

    # Send a SIGKILL (a "hard" kill) to each child process. This is CRITICAL:
    # without it, this process may end up waiting a long time on the proc.wait()
    # below, and never get to saving the children logs, making test timeouts
    # very hard to debug.
    for i, proc in running_procs.items():
      proc.kill()
      print(f"Process {i} killed.")
    print("Killed all child processes.", flush=True)

  retvals = []
  stdouts = []
  stderrs = []
  for proc, fds, (stdout, stderr) in zip(
      subprocesses, output_files, output_filenames
  ):
    retvals.append(proc.wait())
    for fd in fds:
      fd.close()
    stdouts.append(pathlib.Path(stdout).read_text(errors="replace"))
    stderrs.append(pathlib.Path(stderr).read_text(errors="replace"))

  print(" All finished ".center(80, "="), flush=True)

  print(" Summary ".center(80, "="))
  for i, (retval, stdout, stderr) in enumerate(zip(retvals, stdouts, stderrs)):
    m = re.search(r"Ran \d+ tests? in [\d.]+s\n\n.*", stderr, re.MULTILINE)
    result = m.group().replace("\n\n", "; ") if m else "Test crashed?"
    print(
        f"Process {i}, ret: {retval}, len(stdout): {len(stdout)}, "
        f"len(stderr): {len(stderr)}; {result}"
    )

  print(" Detailed logs ".center(80, "="))
  for i, (retval, stdout, stderr) in enumerate(zip(retvals, stdouts, stderrs)):
    print(f" Process {i}: return code: {retval} ".center(80, "="))
    if stdout:
      print(f" Process {i} stdout ".center(80, "-"))
      print(stdout)
    if stderr:
      print(f" Process {i} stderr ".center(80, "-"))
      print(stderr)

  print(" Done detailed logs ".center(80, "="), flush=True)
  for i, (retval, stderr) in enumerate(zip(retvals, stderrs)):
    if retval != 0:
      if expect_failures_with_regex is not None:
        assert re.search(
            expect_failures_with_regex, stderr
        ), f"process {i} failed, expected regex: {expect_failures_with_regex}"
      else:
        assert retval == 0, f"process {i} failed, return value: {retval}"


class MultiProcessTest(absltest.TestCase):
# TODO(b/378138653) Support TPUless MultiProcessTest.

  def setUp(self):
    """Start distributed service."""
    super().setUp()
    assert jax.process_count() == NUM_PROCESSES.value, (
        jax.process_count(),
        NUM_PROCESSES.value,
    )
    # Make sure all processes are at the same test case.
    client = multihost.get_jax_distributed_client()
    client.wait_at_barrier(self._testMethodName, 10000)
