import sys
import time
from collections import Counter

import pytest

import ray
from ray._common.test_utils import SignalActor, wait_for_condition
from ray._raylet import GcsClient
from ray.core.generated import autoscaler_pb2, common_pb2
from ray.util.scheduling_strategies import (
    NodeAffinitySchedulingStrategy,
    PlacementGroupSchedulingStrategy,
)
from ray.util.state import list_tasks


def test_idle_termination(ray_start_cluster):
    cluster = ray_start_cluster
    head_node = cluster.add_node(resources={"head": 1})
    ray.init(address=cluster.address)
    worker_node = cluster.add_node(resources={"worker": 1})
    cluster.wait_for_nodes()

    head_node_id = head_node.node_id
    worker_node_id = worker_node.node_id

    wait_for_condition(
        lambda: {node["NodeID"] for node in ray.nodes() if (node["Alive"])}
        == {head_node_id, worker_node_id}
    )

    @ray.remote(num_cpus=1, resources={"worker": 1})
    class Actor:
        def ping(self):
            pass

    actor = Actor.remote()
    ray.get(actor.ping.remote())

    gcs_client = GcsClient(address=ray.get_runtime_context().gcs_address)

    # The worker node is not idle so the drain request should be rejected.
    is_accepted, rejection_reason_message = gcs_client.drain_node(
        worker_node_id,
        autoscaler_pb2.DrainNodeReason.Value("DRAIN_NODE_REASON_IDLE_TERMINATION"),
        "idle for long enough",
        2**63 - 1,
    )
    assert not is_accepted
    assert (
        "The node to be idle terminated is no longer idle." in rejection_reason_message
    )

    ray.kill(actor)

    def drain_until_accept():
        # The worker node is idle now so the drain request should be accepted.
        is_accepted, _ = gcs_client.drain_node(
            worker_node_id,
            autoscaler_pb2.DrainNodeReason.Value("DRAIN_NODE_REASON_IDLE_TERMINATION"),
            "idle for long enough",
            2**63 - 1,
        )
        return is_accepted

    wait_for_condition(drain_until_accept)

    wait_for_condition(
        lambda: {node["NodeID"] for node in ray.nodes() if (node["Alive"])}
        == {head_node_id}
    )

    worker_node = [node for node in ray.nodes() if node["NodeID"] == worker_node_id][0]
    assert worker_node["DeathReason"] == common_pb2.NodeDeathInfo.Reason.Value(
        "AUTOSCALER_DRAIN_IDLE"
    )
    assert worker_node["DeathReasonMessage"] == "idle for long enough"

    # Draining a dead node is always accepted.
    is_accepted, _ = gcs_client.drain_node(
        worker_node_id,
        autoscaler_pb2.DrainNodeReason.Value("DRAIN_NODE_REASON_IDLE_TERMINATION"),
        "idle for long enough",
        2**63 - 1,
    )
    assert is_accepted


def test_idle_drain_rejected_while_holding_pinned_object(ray_start_cluster):
    # Tests that a node is not drained if it's holding a pinned object
    #
    # The resource-sync period is widened so the object-store-memory idle signal is NOT
    # refreshed between pinning the object and the drain request. The raylet must refresh
    # it at the drain decision regardless; without that, the cached signal is stale and
    # the drain is wrongly accepted once CPU/lease go idle.
    cluster = ray_start_cluster
    cluster.add_node(
        num_cpus=0,
        resources={"head": 1},
        _system_config={"raylet_report_resources_period_milliseconds": 120000},
    )
    ray.init(address=cluster.address)
    worker_node = cluster.add_node(num_cpus=1, resources={"worker": 1})
    cluster.wait_for_nodes()
    worker_node_id = worker_node.node_id

    gcs_client = GcsClient(address=ray.get_runtime_context().gcs_address)

    # >100 KiB so the return is stored in the worker's plasma (in_plasma), not inlined
    # into the owner's reply.
    @ray.remote(num_cpus=1, max_retries=0)
    def make_object():
        return b"x" * (10 * 1024 * 1024)

    # Runs on the worker (head has 0 CPU). Keep the ref but do NOT fetch it, so the only
    # copy stays pinned in the worker's plasma.
    ref = make_object.remote()
    ready, _ = ray.wait([ref], num_returns=1, timeout=60, fetch_local=False)
    assert ready

    def drain_idle():
        is_accepted, _ = gcs_client.drain_node(
            worker_node_id,
            autoscaler_pb2.DrainNodeReason.Value("DRAIN_NODE_REASON_IDLE_TERMINATION"),
            "idle drain while node holds a referenced object",
            2**63 - 1,
        )
        return is_accepted

    # Make idle drain requests during the window where worker lease is returned (CPU +
    # worker footprint go idle). The object held by the node should prevent it from draining.
    for _ in range(10):
        assert not drain_idle(), "idle drain accepted while the node held an object"
        time.sleep(0.5)

    # The worker survived and the object is still retrievable.
    assert worker_node_id in {node["NodeID"] for node in ray.nodes() if node["Alive"]}
    assert len(ray.get(ref)) == 10 * 1024 * 1024


def test_preemption(ray_start_cluster):
    cluster = ray_start_cluster
    head_node = cluster.add_node(resources={"head": 1})
    ray.init(address=cluster.address)
    worker_node = cluster.add_node(resources={"worker": 1})
    cluster.wait_for_nodes()

    head_node_id = head_node.node_id
    worker_node_id = worker_node.node_id

    @ray.remote(num_cpus=1, resources={"worker": 1})
    class Actor:
        def ping(self):
            pass

    actor = Actor.remote()
    ray.get(actor.ping.remote())

    gcs_client = GcsClient(address=ray.get_runtime_context().gcs_address)

    with pytest.raises(ray.exceptions.RaySystemError):
        # Test invalid draining deadline
        gcs_client.drain_node(
            worker_node_id,
            autoscaler_pb2.DrainNodeReason.Value("DRAIN_NODE_REASON_PREEMPTION"),
            "preemption",
            -1,
        )

    # The worker node is not idle but the drain request should be still accepted.
    is_accepted, _ = gcs_client.drain_node(
        worker_node_id,
        autoscaler_pb2.DrainNodeReason.Value("DRAIN_NODE_REASON_PREEMPTION"),
        "preemption",
        2**63 - 1,
    )
    assert is_accepted

    time.sleep(1)

    # Worker node should still be alive since it's not idle and cannot be drained.
    wait_for_condition(
        lambda: {node["NodeID"] for node in ray.nodes() if (node["Alive"])}
        == {head_node_id, worker_node_id}
    )

    ray.kill(actor)
    wait_for_condition(
        lambda: {node["NodeID"] for node in ray.nodes() if (node["Alive"])}
        == {head_node_id}
    )

    worker_node = [node for node in ray.nodes() if node["NodeID"] == worker_node_id][0]
    assert worker_node["DeathReason"] == common_pb2.NodeDeathInfo.Reason.Value(
        "AUTOSCALER_DRAIN_PREEMPTED"
    )
    assert worker_node["DeathReasonMessage"] == "preemption"


@pytest.mark.parametrize(
    "graceful",
    [True, False],
)
def test_preemption_after_draining_deadline(monkeypatch, ray_start_cluster, graceful):
    monkeypatch.setenv("RAY_health_check_failure_threshold", "3")
    monkeypatch.setenv("RAY_health_check_timeout_ms", "100")
    monkeypatch.setenv("RAY_health_check_period_ms", "1000")
    monkeypatch.setenv("RAY_health_check_initial_delay_ms", "0")

    cluster = ray_start_cluster
    head_node = cluster.add_node(resources={"head": 1})
    ray.init(address=cluster.address)
    worker_node = cluster.add_node(resources={"worker": 1})
    cluster.wait_for_nodes()

    head_node_id = head_node.node_id
    worker_node_id = worker_node.node_id

    wait_for_condition(
        lambda: {node["NodeID"] for node in ray.nodes() if (node["Alive"])}
        == {head_node_id, worker_node_id}
    )

    @ray.remote(num_cpus=1, resources={"worker": 1})
    class Actor:
        def ping(self):
            pass

    actor = Actor.remote()
    ray.get(actor.ping.remote())

    gcs_client = GcsClient(address=ray.get_runtime_context().gcs_address)

    # The worker node is not idle but the drain request should be still accepted.
    is_accepted, _ = gcs_client.drain_node(
        worker_node_id,
        autoscaler_pb2.DrainNodeReason.Value("DRAIN_NODE_REASON_PREEMPTION"),
        "preemption",
        1,
    )
    assert is_accepted

    # Simulate autoscaler terminates the worker node after the draining deadline.
    cluster.remove_node(worker_node, graceful)

    wait_for_condition(
        lambda: {node["NodeID"] for node in ray.nodes() if (node["Alive"])}
        == {head_node_id},
    )

    worker_node = [node for node in ray.nodes() if node["NodeID"] == worker_node_id][0]
    assert worker_node["DeathReason"] == common_pb2.NodeDeathInfo.Reason.Value(
        "AUTOSCALER_DRAIN_PREEMPTED"
    )
    assert worker_node["DeathReasonMessage"] == "preemption"


def test_node_death_before_draining_deadline(monkeypatch, ray_start_cluster):
    monkeypatch.setenv("RAY_health_check_failure_threshold", "3")
    monkeypatch.setenv("RAY_health_check_timeout_ms", "100")
    monkeypatch.setenv("RAY_health_check_period_ms", "1000")
    monkeypatch.setenv("RAY_health_check_initial_delay_ms", "0")

    cluster = ray_start_cluster
    head_node = cluster.add_node(resources={"head": 1})
    ray.init(address=cluster.address)
    worker_node = cluster.add_node(resources={"worker": 1})
    cluster.wait_for_nodes()

    head_node_id = head_node.node_id
    worker_node_id = worker_node.node_id

    wait_for_condition(
        lambda: {node["NodeID"] for node in ray.nodes() if (node["Alive"])}
        == {head_node_id, worker_node_id}
    )

    @ray.remote(num_cpus=1, resources={"worker": 1})
    class Actor:
        def ping(self):
            pass

    actor = Actor.remote()
    ray.get(actor.ping.remote())

    gcs_client = GcsClient(address=ray.get_runtime_context().gcs_address)

    # The worker node is not idle but the drain request should be still accepted.
    is_accepted, _ = gcs_client.drain_node(
        worker_node_id,
        autoscaler_pb2.DrainNodeReason.Value("DRAIN_NODE_REASON_PREEMPTION"),
        "preemption",
        2**63 - 1,
    )
    assert is_accepted

    # Simulate the worker node crashes before the draining deadline.
    cluster.remove_node(worker_node, False)

    wait_for_condition(
        lambda: {node["NodeID"] for node in ray.nodes() if (node["Alive"])}
        == {head_node_id},
    )

    # Since worker node failure is detected to be before the draining deadline,
    # this is considered as an unexpected termination.
    worker_node = [node for node in ray.nodes() if node["NodeID"] == worker_node_id][0]
    assert worker_node["DeathReason"] == common_pb2.NodeDeathInfo.Reason.Value(
        "UNEXPECTED_TERMINATION"
    )
    assert (
        worker_node["DeathReasonMessage"]
        == "health check failed due to missing too many heartbeats"
    )


def test_scheduling_placement_groups_during_draining(ray_start_cluster):
    """Test that the draining node is unschedulable for new pgs."""
    cluster = ray_start_cluster
    node1 = cluster.add_node(num_cpus=1, resources={"node1": 1})
    ray.init(address=cluster.address)
    node2 = cluster.add_node(num_cpus=1, resources={"node2": 1})
    cluster.add_node(num_cpus=2, resources={"node3": 1})
    cluster.wait_for_nodes()

    node1_id = node1.node_id
    node2_id = node2.node_id
    node3_id = node2.node_id

    gcs_client = GcsClient(address=ray.get_runtime_context().gcs_address)

    # The node is idle so the draining request should be accepted.
    is_accepted, _ = gcs_client.drain_node(
        node3_id,
        autoscaler_pb2.DrainNodeReason.Value("DRAIN_NODE_REASON_PREEMPTION"),
        "preemption",
        2**63 - 1,
    )
    assert is_accepted

    @ray.remote
    def get_node_id():
        return ray.get_runtime_context().get_node_id()

    # Even though node3 is the best for pack but it's draining
    # so the pg should be on node1 and node2
    pg = ray.util.placement_group(bundles=[{"CPU": 1}, {"CPU": 1}], strategy="PACK")
    {
        ray.get(
            get_node_id.options(
                scheduling_strategy=PlacementGroupSchedulingStrategy(
                    placement_group=pg,
                    placement_group_bundle_index=0,
                )
            ).remote()
        ),
        ray.get(
            get_node_id.options(
                scheduling_strategy=PlacementGroupSchedulingStrategy(
                    placement_group=pg,
                    placement_group_bundle_index=1,
                )
            ).remote()
        ),
    } == {node1_id, node2_id}


def test_scheduling_tasks_and_actors_during_draining(ray_start_cluster):
    """Test that the draining node is unschedulable for new tasks and actors."""
    cluster = ray_start_cluster
    head_node = cluster.add_node(num_cpus=1, resources={"head": 1})
    ray.init(address=cluster.address)
    worker_node = cluster.add_node(num_cpus=1, resources={"worker": 1})
    cluster.wait_for_nodes()

    head_node_id = head_node.node_id
    worker_node_id = worker_node.node_id

    @ray.remote
    class Actor:
        def ping(self):
            pass

    actor = Actor.options(num_cpus=0, resources={"worker": 1}).remote()
    ray.get(actor.ping.remote())

    gcs_client = GcsClient(address=ray.get_runtime_context().gcs_address)

    # The worker node is not idle but the drain request should be still accepted.
    is_accepted, _ = gcs_client.drain_node(
        worker_node_id,
        autoscaler_pb2.DrainNodeReason.Value("DRAIN_NODE_REASON_PREEMPTION"),
        "preemption",
        2**63 - 1,
    )
    assert is_accepted

    @ray.remote
    def get_node_id():
        return ray.get_runtime_context().get_node_id()

    assert (
        ray.get(get_node_id.options(scheduling_strategy="SPREAD").remote())
        == head_node_id
    )
    assert (
        ray.get(get_node_id.options(scheduling_strategy="SPREAD").remote())
        == head_node_id
    )

    assert (
        ray.get(
            get_node_id.options(
                scheduling_strategy=NodeAffinitySchedulingStrategy(
                    worker_node_id, soft=True
                )
            ).remote()
        )
        == head_node_id
    )

    with pytest.raises(ray.exceptions.TaskUnschedulableError):
        ray.get(
            get_node_id.options(
                label_selector={ray._raylet.RAY_NODE_ID_KEY: worker_node_id}
            ).remote()
        )

    head_actor = Actor.options(num_cpus=1, resources={"head": 1}).remote()
    ray.get(head_actor.ping.remote())

    obj = get_node_id.remote()

    # Cannot run on the draining worker node even though it has resources.
    with pytest.raises(ray.exceptions.GetTimeoutError):
        ray.get(obj, timeout=2)

    ray.kill(head_actor)
    ray.get(obj, timeout=2) == head_node_id


@pytest.mark.parametrize(
    "graceful",
    [False, True],
)
def test_draining_reason(ray_start_cluster, graceful):
    cluster = ray_start_cluster
    cluster.add_node(num_cpus=1, resources={"node1": 1})
    ray.init(
        address=cluster.address,
    )
    node2 = cluster.add_node(num_cpus=1, resources={"node2": 1})

    @ray.remote
    class Actor:
        def ping(self):
            pass

    gcs_client = GcsClient(address=ray.get_runtime_context().gcs_address)

    node2_id = node2.node_id

    # Schedule actor
    actor = Actor.options(num_cpus=0, resources={"node2": 1}).remote()
    ray.get(actor.ping.remote())

    drain_reason_message = "testing node preemption."
    # Preemption is always accepted.
    is_accepted, _ = gcs_client.drain_node(
        node2_id,
        autoscaler_pb2.DrainNodeReason.Value("DRAIN_NODE_REASON_PREEMPTION"),
        drain_reason_message,
        1,
    )
    assert is_accepted

    # Simulate autoscaler terminates the worker node after the draining deadline.
    cluster.remove_node(node2, graceful)

    def check_actor_died_error():
        try:
            ray.get(actor.ping.remote())
            return False
        except ray.exceptions.ActorDiedError as e:
            assert e.preempted
            if graceful:
                assert "The actor died because its node has died." in str(e)
                assert "the actor's node was preempted: " + drain_reason_message in str(
                    e
                )
        return True

    wait_for_condition(check_actor_died_error)


def test_drain_node_actor_restart(ray_start_cluster):
    cluster = ray_start_cluster
    cluster.add_node(num_cpus=1, resources={"head": 1})
    ray.init(address=cluster.address)

    gcs_client = GcsClient(address=ray.get_runtime_context().gcs_address)

    @ray.remote(max_restarts=1)
    class Actor:
        def get_node_id(self):
            return ray.get_runtime_context().get_node_id()

    # Prepare the first worker node for the actor.
    cur_worker = cluster.add_node(num_cpus=1, resources={"worker": 1})
    cluster.wait_for_nodes()

    actor = Actor.options(num_cpus=0, resources={"worker": 1}).remote()

    def actor_started():
        node_id = ray.get(actor.get_node_id.remote())
        return node_id == cur_worker.node_id

    wait_for_condition(actor_started, timeout=5)

    # Kill the current worker node.
    cluster.remove_node(cur_worker, True)

    # Prepare a new worker node for the actor to be restarted on later.
    cur_worker = cluster.add_node(num_cpus=1, resources={"worker": 1})
    cluster.wait_for_nodes()

    # Make sure the actor is restarted on the new worker node.
    # This should be counted into the max_restarts of the actor.
    wait_for_condition(actor_started, timeout=5)

    # Preemption the current worker node.
    is_accepted, _ = gcs_client.drain_node(
        cur_worker.node_id,
        autoscaler_pb2.DrainNodeReason.Value("DRAIN_NODE_REASON_PREEMPTION"),
        "preemption",
        1,
    )
    assert is_accepted
    cluster.remove_node(cur_worker, True)

    # Prepare a new worker node for the actor to be restarted on later.
    cur_worker = cluster.add_node(num_cpus=1, resources={"worker": 1})
    cluster.wait_for_nodes()

    # Make sure the actor is restarted on the new worker node.
    # This should not be counted into the max_restarts of the actor because the actor was preempted.
    wait_for_condition(actor_started, timeout=5)

    # Kill the current worker node.
    cluster.remove_node(cur_worker, True)

    # Prepare a new worker node, however, the actor should not be restarted on this node, since
    # the max_restarts is reached.
    cur_worker = cluster.add_node(num_cpus=1, resources={"worker": 1})
    cluster.wait_for_nodes()

    # The actor should not be restarted, thus an exception should be raised.
    with pytest.raises(RuntimeError):
        wait_for_condition(actor_started, timeout=5)


def test_drain_node_task_retry(ray_start_cluster):
    cluster = ray_start_cluster
    cluster.add_node(num_cpus=1, resources={"head": 100})
    ray.init(address=cluster.address)

    cur_worker = cluster.add_node(num_cpus=1, resources={"worker": 1})
    cluster.wait_for_nodes()
    node_ids = Counter()

    gcs_client = GcsClient(address=ray.get_runtime_context().gcs_address)

    @ray.remote(resources={"head": 1})
    class NodeTracker:
        def __init__(self):
            self._node_ids = Counter()

        def add_node(self, node_id):
            self._node_ids.update([node_id])

        def nodes(self):
            return self._node_ids

    @ray.remote(max_retries=1, resources={"worker": 1})
    def func(signal, nodes):
        node_id = ray.get_runtime_context().get_node_id()
        ray.get(nodes.add_node.remote(node_id))
        ray.get(signal.wait.remote())
        return node_id

    signal = SignalActor.options(resources={"head": 1}).remote()
    node_tracker = NodeTracker.remote()
    r1 = func.remote(signal, node_tracker)

    # Verify the first node is added to the counter by the func.remote task.
    node_ids.update([cur_worker.node_id])
    wait_for_condition(lambda: ray.get(node_tracker.nodes.remote()) == node_ids)

    # Remove the current worker node and add a new one to trigger a retry.
    cluster.remove_node(cur_worker, True)
    cur_worker = cluster.add_node(num_cpus=1, resources={"worker": 1})

    # Verify the second node is added to the counter by the task after a retry.
    node_ids.update([cur_worker.node_id])
    wait_for_condition(lambda: ray.get(node_tracker.nodes.remote()) == node_ids)

    # Preempt the second node and add a new one to trigger a retry.
    is_accepted, _ = gcs_client.drain_node(
        cur_worker.node_id,
        autoscaler_pb2.DrainNodeReason.Value("DRAIN_NODE_REASON_PREEMPTION"),
        "preemption",
        1,
    )
    assert is_accepted
    cluster.remove_node(cur_worker, True)
    cur_worker = cluster.add_node(num_cpus=1, resources={"worker": 1})

    # Verify the third node is added to the counter after a preemption retry.
    node_ids.update([cur_worker.node_id])
    wait_for_condition(lambda: ray.get(node_tracker.nodes.remote()) == node_ids)

    # Remove the third node and add a new one, but the task should not retry.
    cluster.remove_node(cur_worker, True)
    cur_worker = cluster.add_node(num_cpus=1, resources={"worker": 1})

    # max_retries is reached, the task should fail.
    with pytest.raises(ray.exceptions.NodeDiedError):
        ray.get(r1)


def test_leases_rescheduling_during_draining(ray_start_cluster):
    """Test that when a node is being drained, leases inside local lease manager
    will be cancelled and re-added to the cluster lease manager for rescheduling
    instead of being marked as permanently infeasible.

    This is regression test for https://github.com/ray-project/ray/pull/57834/
    """
    cluster = ray_start_cluster
    cluster.add_node(num_cpus=0)
    ray.init(address=cluster.address)

    worker1 = cluster.add_node(num_cpus=1)
    cluster.wait_for_nodes()

    gcs_client = GcsClient(address=ray.get_runtime_context().gcs_address)

    @ray.remote(num_cpus=1)
    class Actor:
        def ping(self):
            pass

    actor = Actor.remote()
    ray.get(actor.ping.remote())

    @ray.remote(num_cpus=1)
    def get_node_id():
        return ray.get_runtime_context().get_node_id()

    obj_ref = get_node_id.options(name="f1").remote()

    def verify_f1_pending_node_assignment():
        tasks = list_tasks(filters=[("name", "=", "f1")])
        assert len(tasks) == 1
        assert tasks[0]["state"] == "PENDING_NODE_ASSIGNMENT"
        return True

    # f1 should be in the local lease manager of worker1,
    # waiting for resource to be available.
    wait_for_condition(verify_f1_pending_node_assignment)

    is_accepted, _ = gcs_client.drain_node(
        worker1.node_id,
        autoscaler_pb2.DrainNodeReason.Value("DRAIN_NODE_REASON_PREEMPTION"),
        "preemption",
        2**63 - 1,
    )
    assert is_accepted

    # The task should be rescheduled on another node.
    worker2 = cluster.add_node(num_cpus=1)
    assert ray.get(obj_ref) == worker2.node_id


if __name__ == "__main__":

    sys.exit(pytest.main(["-sv", __file__]))