# Copyright 2022 The HuggingFace Team. All rights reserved.
#
# 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.
"""
Test file to ensure that in general certain situational setups for notebooks work.
"""

import os
import time

from pytest import mark, raises
from torch.distributed.elastic.multiprocessing.errors import ChildFailedError

from accelerate import PartialState, notebook_launcher
from accelerate.test_utils import require_bnb
from accelerate.utils import is_bnb_available, is_rocm_available, is_xpu_available


def basic_function():
    # Just prints the PartialState
    print(f"PartialState:\n{PartialState()}")


def tough_nut_function(queue):
    if queue.empty():
        return
    trial = queue.get()
    if trial > 0:
        queue.put(trial - 1)
        raise RuntimeError("The nut hasn't cracked yet! Try again.")

    print(f"PartialState:\n{PartialState()}")


def bipolar_sleep_function(sleep_sec: int):
    state = PartialState()
    if state.process_index % 2 == 0:
        raise RuntimeError("I'm an even process. I don't like to sleep.")
    else:
        time.sleep(sleep_sec)


def assert_clean_cuda_fork_function():
    # Regression guard: notebook_launcher must not load the CUDA driver in the
    # parent before forking children. If it does, _is_in_bad_fork is set in each
    # forked child and any later CUDA call (e.g. Accelerator() -> set_device)
    # fails with "Cannot re-initialize CUDA in forked subprocess".
    import torch

    if torch.cuda.is_available():
        assert not torch._C._cuda_isInBadFork(), (
            "Child process started in bad-fork state — notebook_launcher tainted "
            "the parent's CUDA driver before forking."
        )


NUM_PROCESSES = int(os.environ.get("ACCELERATE_NUM_PROCESSES", 1))
require_multi_process = mark.skipif(NUM_PROCESSES < 2, reason="Need at least 2 processes")


def test_can_initialize():
    notebook_launcher(basic_function, (), num_processes=NUM_PROCESSES)


@require_multi_process
def test_clean_cuda_fork():
    # Regression test for the CUDA-driver-loaded-in-parent bug that broke
    # notebook_launcher on accelerate >=1.9 and torch >=2.10 (see issue #3925).
    notebook_launcher(assert_clean_cuda_fork_function, (), num_processes=NUM_PROCESSES)


@require_multi_process
def test_static_rdzv_backend():
    notebook_launcher(basic_function, (), num_processes=NUM_PROCESSES, rdzv_backend="static")


@require_multi_process
def test_c10d_rdzv_backend():
    notebook_launcher(basic_function, (), num_processes=NUM_PROCESSES, rdzv_backend="c10d")


@require_multi_process
def test_fault_tolerant(max_restarts: int = 3):
    # Use torch.multiprocessing to get the right context for the current device
    import torch.multiprocessing as mp

    # Get appropriate context - 'spawn' for XPU/ROCm (matches notebook_launcher's start_method),
    # 'fork' for others
    if is_xpu_available() or is_rocm_available():
        ctx = mp.get_context("spawn")
    else:
        ctx = mp.get_context("fork")
    queue = ctx.Queue()
    queue.put(max_restarts)
    notebook_launcher(tough_nut_function, (queue,), num_processes=NUM_PROCESSES, max_restarts=max_restarts)


@require_multi_process
def test_monitoring(monitor_interval: float = 0.01, sleep_sec: int = 100):
    start_time = time.time()
    with raises(ChildFailedError, match="I'm an even process. I don't like to sleep."):
        notebook_launcher(
            bipolar_sleep_function,
            (sleep_sec,),
            num_processes=NUM_PROCESSES,
            monitor_interval=monitor_interval,
        )
    assert time.time() - start_time < sleep_sec, "Monitoring did not stop the process in time."


@require_bnb
def test_problematic_imports():
    with raises(RuntimeError, match="Please keep these imports"):
        import bitsandbytes as bnb  # noqa: F401

        notebook_launcher(basic_function, (), num_processes=NUM_PROCESSES)


def main():
    print("Test basic notebook can be ran")
    test_can_initialize()
    print("Test clean CUDA fork (regression test for #3925)")
    test_clean_cuda_fork()
    print("Test static rendezvous backend")
    test_static_rdzv_backend()
    print("Test c10d rendezvous backend")
    test_c10d_rdzv_backend()
    print("Test fault tolerant")
    test_fault_tolerant()
    print("Test monitoring")
    test_monitoring()
    if is_bnb_available():
        print("Test problematic imports (bnb)")
        test_problematic_imports()
    if NUM_PROCESSES > 1:
        PartialState().destroy_process_group()


if __name__ == "__main__":
    main()
