from __future__ import annotations

import asyncio
import glob
import logging
import os
import signal
import threading
from collections import Counter, UserDict
from time import sleep
import sys

import psutil
import pytest
from tlz import merge

import dask.config
from dask.utils import format_bytes, parse_bytes

import distributed.system
from distributed import Client, Event, KilledWorker, Nanny, Scheduler, Worker, wait
from distributed.compatibility import MACOS, WINDOWS
from distributed.core import Status
from distributed.metrics import monotonic
from distributed.spill import has_zict_210
from distributed.utils_test import (
    NO_AMM,
    captured_logger,
    gen_cluster,
    inc,
    wait_for_state,
)
from distributed.worker_memory import parse_memory_limit
from distributed.worker_state_machine import (
    ComputeTaskEvent,
    DigestMetric,
    ExecuteSuccessEvent,
    GatherDep,
    GatherDepSuccessEvent,
    TaskErredMsg,
)

requires_zict_210 = pytest.mark.skipif(
    not has_zict_210,
    reason="requires zict version >= 2.1.0",
)


def memory_monitor_running(dask_worker: Worker | Nanny) -> bool:
    return "memory_monitor" in dask_worker.periodic_callbacks


def test_parse_memory_limit_zero():
    logger = logging.getLogger(__name__)
    assert parse_memory_limit(0, 1, logger=logger) is None
    assert parse_memory_limit("0", 1, logger=logger) is None
    assert parse_memory_limit(None, 1, logger=logger) is None


def test_resource_limit(monkeypatch):
    logger = logging.getLogger(__name__)
    assert parse_memory_limit("250MiB", 1, 1, logger=logger) == 1024 * 1024 * 250

    new_limit = 1024 * 1024 * 200
    monkeypatch.setattr(distributed.system, "MEMORY_LIMIT", new_limit)
    assert parse_memory_limit("250MiB", 1, 1, logger=logger) == new_limit


@gen_cluster(nthreads=[("", 1)], worker_kwargs={"memory_limit": "2e3 MB"})
async def test_parse_memory_limit_worker(s, w):
    assert w.memory_manager.memory_limit == 2e9


@gen_cluster(nthreads=[("", 1)], worker_kwargs={"memory_limit": "0.5"})
async def test_parse_memory_limit_worker_relative(s, w):
    assert w.memory_manager.memory_limit > 0.5
    assert w.memory_manager.memory_limit == pytest.approx(
        distributed.system.MEMORY_LIMIT * 0.5
    )


@gen_cluster(
    client=True,
    nthreads=[("", 1)],
    Worker=Nanny,
    worker_kwargs={"memory_limit": "2e3 MB"},
)
async def test_parse_memory_limit_nanny(c, s, n):
    assert n.memory_manager.memory_limit == 2e9
    out = await c.run(lambda dask_worker: dask_worker.memory_manager.memory_limit)
    assert out[n.worker_address] == 2e9


@gen_cluster(
    nthreads=[("127.0.0.1", 1)],
    config={
        "distributed.worker.memory.spill": False,
        "distributed.worker.memory.target": False,
    },
)
async def test_dict_data_if_no_spill_to_disk(s, w):
    assert type(w.data) is dict


class WorkerData(dict):
    def __init__(self, **kwargs):
        super().__init__()
        self.kwargs = kwargs


class WorkerDataLocalDirectory(dict):
    def __init__(self, worker_local_directory, **kwargs):
        super().__init__()
        self.local_directory = worker_local_directory
        self.kwargs = kwargs


@gen_cluster(
    nthreads=[("", 1)], Worker=Worker, worker_kwargs={"data": WorkerDataLocalDirectory}
)
async def test_worker_data_callable_local_directory(s, w):
    assert type(w.memory_manager.data) is WorkerDataLocalDirectory
    assert w.memory_manager.data.local_directory == w.local_directory


@gen_cluster(
    nthreads=[("", 1)],
    Worker=Worker,
    worker_kwargs={"data": (WorkerDataLocalDirectory, {"a": "b"})},
)
async def test_worker_data_callable_local_directory_kwargs(s, w):
    assert type(w.memory_manager.data) is WorkerDataLocalDirectory
    assert w.memory_manager.data.local_directory == w.local_directory
    assert w.memory_manager.data.kwargs == {"a": "b"}


@gen_cluster(
    nthreads=[("", 1)], Worker=Worker, worker_kwargs={"data": (WorkerData, {"a": "b"})}
)
async def test_worker_data_callable_kwargs(s, w):
    assert type(w.memory_manager.data) is WorkerData
    assert w.memory_manager.data.kwargs == {"a": "b"}


class CustomError(Exception):
    pass


class FailToPickle:
    def __init__(self, *, reported_size=0):
        self.reported_size = int(reported_size)

    def __getstate__(self):
        raise CustomError()

    def __sizeof__(self):
        return self.reported_size


async def assert_basic_futures(c: Client) -> None:
    futures = c.map(inc, range(10))
    results = await c.gather(futures)
    assert results == list(map(inc, range(10)))



@pytest.mark.skipif(
    CONDITION=(sys.maxsize < 2 ** 32),
    reason="fails on 32-bit, is it asking for large memory?")
@gen_cluster(client=True)
async def test_fail_to_pickle_execute_1(c, s, a, b):
    """Test failure to serialize triggered by computing a key which is individually
    larger than target. The data is lost and the task is marked as failed; the worker
    remains in usable condition.

    See also
    --------
    test_workerstate_fail_to_pickle_execute_1
    test_workerstate_fail_to_pickle_flight
    test_fail_to_pickle_execute_2
    test_fail_to_pickle_spill
    """
    x = c.submit(FailToPickle, reported_size=100e9, key="x")
    await wait(x)

    assert x.status == "error"

    with pytest.raises(TypeError, match="Could not serialize"):
        await x

    await assert_basic_futures(c)


class FailStoreDict(UserDict):
    def __setitem__(self, key, value):
        raise CustomError()


def test_workerstate_fail_to_pickle_execute_1(ws_with_running_task):
    """Same as test_fail_to_pickle_target_execute_1

    See also
    --------
    test_fail_to_pickle_execute_1
    test_workerstate_fail_to_pickle_flight
    test_fail_to_pickle_execute_2
    test_fail_to_pickle_spill
    """
    ws = ws_with_running_task
    assert not ws.data
    ws.data = FailStoreDict()

    instructions = ws.handle_stimulus(
        ExecuteSuccessEvent.dummy("x", None, stimulus_id="s1")
    )
    assert instructions == [
        DigestMetric(name="compute-duration", value=1.0, stimulus_id="s1"),
        TaskErredMsg.match(key="x", stimulus_id="s1"),
    ]
    assert ws.tasks["x"].state == "error"


@pytest.mark.xfail(reason="https://github.com/dask/distributed/issues/6705")
def test_workerstate_fail_to_pickle_flight(ws):
    """Same as test_workerstate_fail_to_pickle_execute_1, but the task was
    computed on another host and for whatever reason it did not fail to pickle when it
    was sent over the network.

    See also
    --------
    test_fail_to_pickle_execute_1
    test_workerstate_fail_to_pickle_execute_1
    test_fail_to_pickle_execute_2
    test_fail_to_pickle_spill

    See also test_worker_state_machine.py::test_gather_dep_failure, where the task
    instead fails to unpickle when leaving the network stack.
    """
    assert not ws.data
    ws.data = FailStoreDict()
    ws.total_resources = {"R": 1}
    ws.available_resources = {"R": 1}
    ws2 = "127.0.0.1:2"

    instructions = ws.handle_stimulus(
        ComputeTaskEvent.dummy(
            "y", who_has={"x": [ws2]}, resource_restrictions={"R": 1}, stimulus_id="s1"
        ),
        GatherDepSuccessEvent(
            worker=ws2, total_nbytes=1, data={"x": 123}, stimulus_id="s2"
        ),
    )
    assert instructions == [
        GatherDep(worker=ws2, to_gather={"x"}, total_nbytes=1, stimulus_id="s1"),
        TaskErredMsg.match(key="x", stimulus_id="s2"),
    ]
    assert ws.tasks["x"].state == "error"
    assert ws.tasks["y"].state == "waiting"  # Not constrained


@gen_cluster(
    client=True,
    nthreads=[("", 1)],
    worker_kwargs={"memory_limit": "1 kiB"},
    config={
        "distributed.worker.memory.target": 0.5,
        "distributed.worker.memory.spill": False,
        "distributed.worker.memory.pause": False,
    },
)
async def test_fail_to_pickle_execute_2(c, s, a):
    """Test failure to spill triggered by computing a key which is individually smaller
    than target, so it is not spilled immediately. The data is retained and the task is
    NOT marked as failed; the worker remains in usable condition.

    See also
    --------
    test_fail_to_pickle_execute_1
    test_workerstate_fail_to_pickle_execute_1
    test_workerstate_fail_to_pickle_flight
    test_fail_to_pickle_spill
    """
    x = c.submit(FailToPickle, reported_size=256, key="x")
    await wait(x)
    assert x.status == "finished"
    assert set(a.data.memory) == {"x"}

    y = c.submit(lambda: "y" * 256, key="y")
    await wait(y)
    if has_zict_210:
        assert set(a.data.memory) == {"x", "y"}
    else:
        assert set(a.data.memory) == {"y"}

    assert not a.data.disk

    await assert_basic_futures(c)


@requires_zict_210
@gen_cluster(
    client=True,
    nthreads=[("", 1)],
    worker_kwargs={"memory_limit": "1 kB"},
    config={
        "distributed.worker.memory.target": False,
        "distributed.worker.memory.spill": 0.7,
        "distributed.worker.memory.monitor-interval": "100ms",
    },
)
async def test_fail_to_pickle_spill(c, s, a):
    """Test failure to evict a key, triggered by the spill threshold.

    See also
    --------
    test_fail_to_pickle_execute_1
    test_workerstate_fail_to_pickle_execute_1
    test_workerstate_fail_to_pickle_flight
    test_fail_to_pickle_execute_2
    """
    a.monitor.get_process_memory = lambda: 701 if a.data.fast else 0

    with captured_logger(logging.getLogger("distributed.spill")) as logs:
        bad = c.submit(FailToPickle, key="bad")
        await wait(bad)

        # Must wait for memory monitor to kick in
        while True:
            logs_value = logs.getvalue()
            if logs_value:
                break
            await asyncio.sleep(0.01)

    assert "Failed to pickle" in logs_value
    assert "Traceback" in logs_value

    # key is in fast
    assert bad.status == "finished"
    assert bad.key in a.data.fast

    await assert_basic_futures(c)


@gen_cluster(
    client=True,
    nthreads=[("", 1)],
    worker_kwargs={"memory_limit": 1200 / 0.6},
    config={
        "distributed.worker.memory.target": 0.6,
        "distributed.worker.memory.spill": False,
        "distributed.worker.memory.pause": False,
    },
)
async def test_spill_target_threshold(c, s, a):
    """Test distributed.worker.memory.target threshold. Note that in this test we
    disabled spill and pause thresholds, which work on the process memory, and just left
    the target threshold, which works on managed memory so it is unperturbed by the
    several hundreds of MB of unmanaged memory that are typical of the test suite.
    """
    assert not memory_monitor_running(a)

    x = c.submit(lambda: "x" * 500, key="x")
    await wait(x)
    y = c.submit(lambda: "y" * 500, key="y")
    await wait(y)

    assert set(a.data) == {"x", "y"}
    assert set(a.data.memory) == {"x", "y"}

    z = c.submit(lambda: "z" * 500, key="z")
    await wait(z)
    assert set(a.data) == {"x", "y", "z"}
    assert set(a.data.memory) == {"y", "z"}
    assert set(a.data.disk) == {"x"}

    await x
    assert set(a.data.memory) == {"x", "z"}
    assert set(a.data.disk) == {"y"}


@requires_zict_210
@gen_cluster(
    client=True,
    nthreads=[("", 1)],
    worker_kwargs={"memory_limit": 1600},
    config={
        "distributed.worker.memory.target": 0.6,
        "distributed.worker.memory.spill": False,
        "distributed.worker.memory.pause": False,
        "distributed.worker.memory.max-spill": 600,
    },
)
async def test_spill_constrained(c, s, w):
    """Test distributed.worker.memory.max-spill parameter"""
    # spills starts at 1600*0.6=960 bytes of managed memory

    # size in memory ~200; size on disk ~400
    x = c.submit(lambda: "x" * 200, key="x")
    await wait(x)
    # size in memory ~500; size on disk ~700
    y = c.submit(lambda: "y" * 500, key="y")
    await wait(y)

    assert set(w.data) == {x.key, y.key}
    assert set(w.data.memory) == {x.key, y.key}

    z = c.submit(lambda: "z" * 500, key="z")
    await wait(z)

    assert set(w.data) == {x.key, y.key, z.key}

    # max_spill has not been reached
    assert set(w.data.memory) == {y.key, z.key}
    assert set(w.data.disk) == {x.key}

    # zb is individually larger than max_spill
    zb = c.submit(lambda: "z" * 1700, key="zb")
    await wait(zb)

    assert set(w.data.memory) == {y.key, z.key, zb.key}
    assert set(w.data.disk) == {x.key}

    del zb
    while "zb" in w.data:
        await asyncio.sleep(0.01)

    # zc is individually smaller than max_spill, but the evicted key together with
    # x it exceeds max_spill
    zc = c.submit(lambda: "z" * 500, key="zc")
    await wait(zc)
    assert set(w.data.memory) == {y.key, z.key, zc.key}
    assert set(w.data.disk) == {x.key}


@gen_cluster(
    nthreads=[("", 1)],
    client=True,
    worker_kwargs={"memory_limit": "1000 MB"},
    config={
        "distributed.worker.memory.target": False,
        "distributed.worker.memory.spill": 0.7,
        "distributed.worker.memory.pause": False,
        "distributed.worker.memory.monitor-interval": "10ms",
    },
)
async def test_spill_spill_threshold(c, s, a):
    """Test distributed.worker.memory.spill threshold.
    Test that the spill threshold uses the process memory and not the managed memory
    reported by sizeof(), which may be inaccurate.
    """
    assert memory_monitor_running(a)
    a.monitor.get_process_memory = lambda: 800_000_000 if a.data.fast else 0
    x = c.submit(inc, 0, key="x")
    while not a.data.disk:
        await asyncio.sleep(0.01)
    assert await x == 1


@pytest.mark.parametrize(
    "target,managed,expect_spilled",
    [
        # no target -> no hysteresis
        # Over-report managed memory to test that the automated LRU eviction based on
        # target is never triggered
        (False, int(10e9), 1),
        # Under-report managed memory, so that we reach the spill threshold for process
        # memory without first reaching the target threshold for managed memory
        # target == spill -> no hysteresis
        (0.7, 0, 1),
        # target < spill -> hysteresis from spill to target
        (0.4, 0, 7),
    ],
)
@gen_cluster(
    nthreads=[],
    client=True,
    config={
        "distributed.worker.memory.spill": 0.7,
        "distributed.worker.memory.pause": False,
        "distributed.worker.memory.monitor-interval": "10ms",
    },
)
async def test_spill_hysteresis(c, s, target, managed, expect_spilled):
    """
    1. Test that you can enable the spill threshold while leaving the target threshold
       to False
    2. Test the hysteresis system where, once you reach the spill threshold, the worker
       won't stop spilling until the target threshold is reached
    """

    class C:
        def __sizeof__(self):
            return managed

    with dask.config.set({"distributed.worker.memory.target": target}):
        async with Worker(s.address, memory_limit="1000 MB") as a:
            a.monitor.get_process_memory = lambda: 50_000_000 * len(a.data.fast)

            # Add 500MB (reported) process memory. Spilling must not happen.
            futures = [c.submit(C, pure=False) for _ in range(10)]
            await wait(futures)
            await asyncio.sleep(0.1)
            assert not a.data.disk

            # Add another 250MB unmanaged memory. This must trigger the spilling.
            futures += [c.submit(C, pure=False) for _ in range(5)]
            await wait(futures)

            # Wait until spilling starts. Then, wait until it stops.
            prev_n = 0
            while not a.data.disk or len(a.data.disk) > prev_n:
                prev_n = len(a.data.disk)
                await asyncio.sleep(0)

            assert len(a.data.disk) == expect_spilled


@gen_cluster(
    nthreads=[("", 1)],
    client=True,
    config={
        "distributed.worker.memory.target": False,
        "distributed.worker.memory.spill": False,
        "distributed.worker.memory.pause": False,
    },
)
async def test_pause_executor_manual(c, s, a):
    assert not memory_monitor_running(a)

    # Task that is running when the worker pauses
    ev_x = Event()

    def f(ev):
        ev.wait()
        return 1

    # Task that is running on the worker when the worker pauses
    x = c.submit(f, ev_x, key="x")
    while a.state.executing_count != 1:
        await asyncio.sleep(0.01)

    # Task that is queued on the worker when the worker pauses
    y = c.submit(inc, 1, key="y")
    while "y" not in a.state.tasks:
        await asyncio.sleep(0.01)

    a.status = Status.paused
    # Wait for sync to scheduler
    while s.workers[a.address].status != Status.paused:
        await asyncio.sleep(0.01)

    # Task that is queued on the scheduler when the worker pauses.
    # It is not sent to the worker.
    z = c.submit(inc, 2, key="z")
    while "z" not in s.tasks or s.tasks["z"].state != "no-worker":
        await asyncio.sleep(0.01)
    assert s.unrunnable == {s.tasks["z"]}

    # Test that a task that already started when the worker paused can complete
    # and its output can be retrieved. Also test that the now free slot won't be
    # used by other tasks.
    await ev_x.set()
    assert await x == 1
    await asyncio.sleep(0.05)

    assert a.state.executing_count == 0
    assert len(a.state.ready) == 1
    assert a.state.tasks["y"].state == "ready"
    assert "z" not in a.state.tasks

    # Unpause. Tasks that were queued on the worker are executed.
    # Tasks that were stuck on the scheduler are sent to the worker and executed.
    a.status = Status.running
    assert await y == 2
    assert await z == 3


@gen_cluster(
    nthreads=[("", 1)],
    client=True,
    worker_kwargs={"memory_limit": "10 GB"},
    config={
        "distributed.worker.memory.target": False,
        "distributed.worker.memory.spill": False,
        "distributed.worker.memory.pause": 0.8,
        "distributed.worker.memory.monitor-interval": "10ms",
    },
)
async def test_pause_executor_with_memory_monitor(c, s, a):
    assert memory_monitor_running(a)
    mocked_rss = 0
    a.monitor.get_process_memory = lambda: mocked_rss

    # Task that is running when the worker pauses
    ev_x = Event()

    def f(ev):
        ev.wait()
        return 1

    # Task that is running on the worker when the worker pauses
    x = c.submit(f, ev_x, key="x")
    while a.state.executing_count != 1:
        await asyncio.sleep(0.01)

    with captured_logger(logging.getLogger("distributed.worker.memory")) as logger:
        # Task that is queued on the worker when the worker pauses
        y = c.submit(inc, 1, key="y")
        while "y" not in a.state.tasks:
            await asyncio.sleep(0.01)

        # Hog the worker with 900GB unmanaged memory
        mocked_rss = 900 * 1000**3
        while s.workers[a.address].status != Status.paused:
            await asyncio.sleep(0.01)

        assert "Pausing worker" in logger.getvalue()

        # Task that is queued on the scheduler when the worker pauses.
        # It is not sent to the worker.
        z = c.submit(inc, 2, key="z")
        while "z" not in s.tasks or s.tasks["z"].state != "no-worker":
            await asyncio.sleep(0.01)
        assert s.unrunnable == {s.tasks["z"]}

        # Test that a task that already started when the worker paused can complete
        # and its output can be retrieved. Also test that the now free slot won't be
        # used by other tasks.
        await ev_x.set()
        assert await x == 1
        await asyncio.sleep(0.05)

        assert a.state.executing_count == 0
        assert len(a.state.ready) == 1
        assert a.state.tasks["y"].state == "ready"
        assert "z" not in a.state.tasks

        # Release the memory. Tasks that were queued on the worker are executed.
        # Tasks that were stuck on the scheduler are sent to the worker and executed.
        mocked_rss = 0
        assert await y == 2
        assert await z == 3

        assert a.status == Status.running
        assert "Resuming worker" in logger.getvalue()


@gen_cluster(
    client=True,
    nthreads=[("", 1), ("", 1)],
    config=merge(
        NO_AMM,
        {
            "distributed.worker.memory.target": False,
            "distributed.worker.memory.spill": False,
            "distributed.worker.memory.pause": False,
        },
    ),
)
async def test_pause_prevents_deps_fetch(c, s, a, b):
    """A worker is paused while there are dependencies ready to fetch, but all other
    workers are in flight
    """
    a_addr = a.address

    class X:
        def __sizeof__(self):
            return 2**40  # Disable clustering in select_keys_for_gather

        def __reduce__(self):
            return X.pause_on_unpickle, ()

        @staticmethod
        def pause_on_unpickle():
            # Note: outside of task execution, distributed.get_worker()
            # returns a random worker running in the process
            for w in Worker._instances:
                if w.address == a_addr:
                    w.status = Status.paused
                    return X()
            assert False

    x = c.submit(X, key="x", workers=[b.address])
    y = c.submit(inc, 1, key="y", workers=[b.address])
    await wait([x, y])
    w = c.submit(lambda _: None, x, key="w", priority=1, workers=[a.address])
    z = c.submit(inc, y, key="z", priority=0, workers=[a.address])

    # - w and z reach worker a within the same message
    # - w and z respectively make x and y go into fetch state.
    #   w has a higher priority than z, therefore w's dependency x has a higher priority
    #   than z's dependency y.
    #   a.state.data_needed[b.address] = ["x", "y"]
    # - ensure_communicating decides to fetch x but not to fetch y together with it, as
    #   it thinks x is 1TB in size
    # - x fetch->flight; a is added to in_flight_workers
    # - y is skipped by ensure_communicating since all workers that hold a replica are
    #   in flight
    # - x reaches a and sends a into paused state
    # - x flight->memory; a is removed from in_flight_workers
    # - ensure_communicating is triggered again
    # - ensure_communicating refuses to fetch y because the worker is paused

    await wait_for_state("y", "fetch", a)
    await asyncio.sleep(0.1)
    assert a.state.tasks["y"].state == "fetch"
    assert "y" not in a.data
    assert [ts.key for ts in a.state.data_needed[b.address]] == ["y"]

    # Unpausing kicks off ensure_communicating again
    a.status = Status.running
    assert await z == 3
    assert a.state.tasks["y"].state == "memory"
    assert "y" in a.data


@gen_cluster(
    client=True,
    nthreads=[("", 1)],
    worker_kwargs={"memory_limit": 0},
    config={"distributed.worker.memory.monitor-interval": "10ms"},
)
async def test_avoid_memory_monitor_if_zero_limit_worker(c, s, a):
    assert type(a.data) is dict
    assert not memory_monitor_running(a)

    future = c.submit(inc, 1)
    assert await future == 2
    await asyncio.sleep(0.05)
    assert await c.submit(inc, 2) == 3  # worker doesn't pause


@gen_cluster(
    client=True,
    nthreads=[("", 1)],
    Worker=Nanny,
    worker_kwargs={"memory_limit": 0},
    config={"distributed.worker.memory.monitor-interval": "10ms"},
)
async def test_avoid_memory_monitor_if_zero_limit_nanny(c, s, nanny):
    typ = await c.run(lambda dask_worker: type(dask_worker.data))
    assert typ == {nanny.worker_address: dict}
    assert not memory_monitor_running(nanny)
    assert not (await c.run(memory_monitor_running))[nanny.worker_address]

    future = c.submit(inc, 1)
    assert await future == 2
    await asyncio.sleep(0.02)
    assert await c.submit(inc, 2) == 3  # worker doesn't pause


@gen_cluster(nthreads=[])
async def test_override_data_worker(s):
    # Use a UserDict to sidestep potential special case handling for dict
    async with Worker(s.address, data=UserDict) as w:
        assert type(w.data) is UserDict

    data = UserDict()
    async with Worker(s.address, data=data) as w:
        assert w.data is data


@gen_cluster(
    client=True,
    nthreads=[("", 1)],
    Worker=Nanny,
    worker_kwargs={"data": UserDict},
)
async def test_override_data_nanny(c, s, n):
    r = await c.run(lambda dask_worker: type(dask_worker.data))
    assert r[n.worker_address] is UserDict


@gen_cluster(
    client=True,
    nthreads=[("", 1)],
    worker_kwargs={"memory_limit": "10 GB", "data": UserDict},
    config={"distributed.worker.memory.monitor-interval": "10ms"},
)
async def test_override_data_vs_memory_monitor(c, s, a):
    a.monitor.get_process_memory = lambda: 8_100_000_000 if a.data else 0
    assert memory_monitor_running(a)

    # Push a key that would normally trip both the target and the spill thresholds
    class C:
        def __sizeof__(self):
            return 8_100_000_000

    # Capture output of log_errors()
    with captured_logger(logging.getLogger("distributed.utils")) as logger:
        x = c.submit(C)
        await wait(x)

        # The pause subsystem of the memory monitor has been tripped.
        # The spill subsystem hasn't.
        while a.status != Status.paused:
            await asyncio.sleep(0.01)
        await asyncio.sleep(0.05)

    # This would happen if memory_monitor() tried to blindly call SpillBuffer.evict()
    assert "Traceback" not in logger.getvalue()

    assert type(a.data) is UserDict
    assert a.data.keys() == {x.key}


class ManualEvictDict(UserDict):
    """A MutableMapping which implements distributed.spill.ManualEvictProto"""

    def __init__(self):
        super().__init__()
        self.evicted = set()

    @property
    def fast(self):
        # Any Sized of bool will do
        return self.keys() - self.evicted

    def evict(self):
        # Evict a random key
        k = next(iter(self.fast))
        self.evicted.add(k)
        return 1


@gen_cluster(
    client=True,
    nthreads=[("", 1)],
    worker_kwargs={"memory_limit": "1 GB", "data": ManualEvictDict},
    config={
        "distributed.worker.memory.pause": False,
        "distributed.worker.memory.monitor-interval": "10ms",
    },
)
async def test_manual_evict_proto(c, s, a):
    """data is a third-party dict-like which respects the ManualEvictProto duck-type
    API. spill threshold is respected.
    """
    a.monitor.get_process_memory = lambda: 701_000_000 if a.data else 0
    assert memory_monitor_running(a)
    assert isinstance(a.data, ManualEvictDict)

    futures = await c.scatter({"x": None, "y": None, "z": None})
    while a.data.evicted != {"x", "y", "z"}:
        await asyncio.sleep(0.01)


async def leak_until_restart(c: Client, s: Scheduler) -> None:
    s.allowed_failures = 0

    def leak():
        L = []
        while True:
            L.append(b"0" * 5_000_000)
            sleep(0.01)

    (addr,) = s.workers
    pid = (await c.run(os.getpid))[addr]

    future = c.submit(leak, key="leak")

    # Wait until the worker is restarted
    while len(s.workers) != 1 or set(s.workers) == {addr}:
        await asyncio.sleep(0.01)

    # Test that the process has been properly waited for and not just left there
    with pytest.raises(psutil.NoSuchProcess):
        psutil.Process(pid)

    with pytest.raises(KilledWorker):
        await future
    assert s.tasks["leak"].suspicious == 1
    assert not any(
        (await c.run(lambda dask_worker: "leak" in dask_worker.state.tasks)).values()
    )
    future.release()
    while "leak" in s.tasks:
        await asyncio.sleep(0.01)


@pytest.mark.slow
@gen_cluster(
    nthreads=[("", 1)],
    client=True,
    Worker=Nanny,
    worker_kwargs={"memory_limit": "400 MiB"},
    config={"distributed.worker.memory.monitor-interval": "10ms"},
)
async def test_nanny_terminate(c, s, a):
    await leak_until_restart(c, s)


@pytest.mark.slow
@pytest.mark.parametrize(
    "ignore_sigterm",
    [
        False,
        pytest.param(True, marks=pytest.mark.skipif(WINDOWS, reason="Needs SIGKILL")),
    ],
)
@gen_cluster(
    nthreads=[("", 1)],
    client=True,
    Worker=Nanny,
    worker_kwargs={"memory_limit": "400 MiB"},
    config={"distributed.worker.memory.monitor-interval": "10ms"},
)
async def test_disk_cleanup_on_terminate(c, s, a, ignore_sigterm):
    """Test that the spilled data on disk is cleaned up when the nanny kills the worker.

    Unlike in a regular worker shutdown, where the worker deletes its own spill
    directory, the cleanup in case of termination from the monitor is performed by the
    nanny.

    The worker may be slow to accept SIGTERM, for whatever reason.
    At the next iteration of the memory manager, if the process is still alive, the
    nanny sends SIGKILL.
    """

    def do_ignore_sigterm():
        # ignore the return value of signal.signal:  it may not be serializable
        signal.signal(signal.SIGTERM, signal.SIG_IGN)

    if ignore_sigterm:
        await c.run(do_ignore_sigterm)

    fut = c.submit(inc, 1, key="myspill")
    await wait(fut)
    await c.run(lambda dask_worker: dask_worker.data.evict())
    glob_out = await c.run(
        lambda dask_worker: glob.glob(dask_worker.local_directory + "/**/myspill")
    )
    spill_fname = next(iter(glob_out.values()))[0]
    assert os.path.exists(spill_fname)

    await leak_until_restart(c, s)
    assert not os.path.exists(spill_fname)


@gen_cluster(
    nthreads=[("", 1)],
    client=True,
    worker_kwargs={"memory_limit": "2 GiB"},
    # ^ must be smaller than system memory limit, otherwise that will take precedence
    config={
        "distributed.worker.memory.target": False,
        "distributed.worker.memory.spill": 0.5,
        "distributed.worker.memory.pause": 0.8,
        "distributed.worker.memory.monitor-interval": "10ms",
    },
)
async def test_pause_while_spilling(c, s, a):
    N_PAUSE = 3
    N_TOTAL = 5

    if a.memory_manager.memory_limit < parse_bytes("2 GiB"):
        pytest.fail(
            f"Set 2 GiB memory limit, got {format_bytes(a.memory_manager.memory_limit)}."
        )

    def get_process_memory():
        if len(a.data) < N_PAUSE:
            # Don't trigger spilling until after all tasks have completed
            return 0
        elif a.data.fast and not a.data.slow:
            # Trigger spilling
            return parse_bytes("1.6 GiB")
        else:
            # Trigger pause, but only after we started spilling
            return parse_bytes("1.9 GiB")

    a.monitor.get_process_memory = get_process_memory

    class SlowSpill:
        def __init__(self, _):
            # Can't pickle a Semaphore, so instead of a default value, we create it
            # here. Don't worry about race conditions; the worker is single-threaded.
            if not hasattr(type(self), "sem"):
                type(self).sem = threading.Semaphore(N_PAUSE)
            # Block if there are N_PAUSE tasks in a.data.fast
            self.sem.acquire()

        def __reduce__(self):
            paused = distributed.get_worker().status == Status.paused
            if not paused:
                sleep(0.1)
            self.sem.release()
            return bool, (paused,)

    futs = c.map(SlowSpill, range(N_TOTAL))
    while len(a.data.slow) < (N_PAUSE + 1 if a.state.ready else N_PAUSE):
        await asyncio.sleep(0.01)

    assert a.status == Status.paused
    # Worker should have become paused after the first `SlowSpill` was evicted, because
    # the spill to disk took longer than the memory monitor interval.
    assert len(a.data.fast) == 0
    # With queuing enabled, after the 3rd `SlowSpill` has been created, there's a race
    # between the scheduler sending the worker a new task, and the memory monitor
    # running and pausing the worker. If the worker gets paused before the 4th task
    # lands, only 3 will be in memory. If after, the 4th will block on the semaphore
    # until one of the others is spilled.
    assert len(a.data.slow) in (N_PAUSE, N_PAUSE + 1)
    n_spilled_while_not_paused = sum(paused is False for paused in a.data.slow.values())
    assert 0 <= n_spilled_while_not_paused <= 1


@pytest.mark.slow
@pytest.mark.skipif(
    condition=MACOS, reason="https://github.com/dask/distributed/issues/6233"
)
@gen_cluster(
    nthreads=[("", 1)],
    client=True,
    worker_kwargs={"memory_limit": "10 GiB"},
    config={
        "distributed.worker.memory.target": False,
        "distributed.worker.memory.spill": 0.6,
        "distributed.worker.memory.pause": False,
        "distributed.worker.memory.monitor-interval": "10ms",
    },
)
async def test_release_evloop_while_spilling(c, s, a):
    N = 100

    def get_process_memory():
        if len(a.data) < N:
            # Don't trigger spilling until after all tasks have completed
            return 0
        return 10 * 2**30

    a.monitor.get_process_memory = get_process_memory

    class SlowSpill:
        def __reduce__(self):
            sleep(0.01)
            return SlowSpill, ()

    futs = [c.submit(SlowSpill, pure=False) for _ in range(N)]
    while len(a.data) < N:
        await asyncio.sleep(0)

    ts = [monotonic()]
    while a.data.fast:
        await asyncio.sleep(0)
        ts.append(monotonic())

    # 100 tasks taking 0.01s to pickle each = 2s to spill everything
    # (this is because everything is pickled twice:
    # https://github.com/dask/distributed/issues/1371).
    # We should regain control of the event loop every 0.5s.
    c = Counter(round(t1 - t0, 1) for t0, t1 in zip(ts, ts[1:]))
    # Depending on the implementation of WorkerMemoryMonitor._maybe_spill:
    # if it calls sleep(0) every 0.5s:
    #   {0.0: 315, 0.5: 4}
    # if it calls sleep(0) after spilling each key:
    #   {0.0: 233}
    # if it never yields:
    #   {0.0: 359, 2.0: 1}
    # Make sure we remain in the first use case.
    assert 1 < sum(v for k, v in c.items() if 0.5 <= k <= 1.9), dict(c)
    assert not any(v for k, v in c.items() if k >= 2.0), dict(c)


@gen_cluster(
    client=True,
    worker_kwargs={"memory_limit": "100 MiB"},
    # ^ must be smaller than system memory limit, otherwise that will take precedence
    config={
        "distributed.worker.memory.target": 0.5,
        "distributed.worker.memory.spill": 1.0,
        "distributed.worker.memory.pause": False,
        "distributed.worker.memory.monitor-interval": "10ms",
    },
)
async def test_digests(c, s, a, b):
    trigger_spill = False

    def get_process_memory():
        return 2**30 if trigger_spill else 0

    a.monitor.get_process_memory = get_process_memory

    x1 = c.submit(inc, 1, key="x1", workers=[a.address])  # store to fast
    x2 = c.submit(inc, x1, key="x2", workers=[a.address])  # read from fast for execute
    y1 = c.submit(inc, x2, key="y1", workers=[b.address])  # read from fast for get-data
    # digest happens only if write/read takes more than 5ms
    await wait([x2, y1])
    assert "disk-load-duration" not in a.digests_total
    assert "get-data-load-duration" not in a.digests_total
    assert "disk-write-target-duration" not in a.digests_total
    assert "disk-write-spill-duration" not in a.digests_total

    # Pass target threshold (50 MiB)
    # We need substantial data to be sure that spilling it will take more than 5ms.
    x3 = c.submit(lambda: "x" * 40_000_000, key="x3", workers=[a.address])
    x4 = c.submit(lambda: "x" * 40_000_000, key="x4", workers=[a.address])
    await wait([x3, x4])
    x5 = c.submit(lambda: "x" * 40_000_000, key="x5", workers=[a.address])
    x6 = c.submit(lambda: "x" * 40_000_000, key="x6", workers=[a.address])
    await wait([x5, x6])
    assert "x3" in a.data.slow
    assert "x4" in a.data.slow
    assert a.digests_total["disk-write-target-duration"] > 0
    assert "disk-write-spill-duration" not in a.digests_total
    x7 = c.submit(lambda x: None, x3, key="x7", workers=[a.address])
    await wait(x7)
    assert a.digests_total["disk-load-duration"] > 0

    y2 = c.submit(lambda x: None, x4, key="y2", workers=[b.address])
    await wait(y2)
    assert a.digests_total["get-data-load-duration"] > 0

    trigger_spill = True
    while a.data.fast:
        await asyncio.sleep(0.01)
    assert a.digests_total["disk-write-spill-duration"] > 0


@pytest.mark.parametrize(
    "cls,name,value",
    [
        (Worker, "memory_limit", 123e9),
        (Worker, "memory_target_fraction", 0.789),
        (Worker, "memory_spill_fraction", 0.789),
        (Worker, "memory_pause_fraction", 0.789),
        (Nanny, "memory_limit", 123e9),
        (Nanny, "memory_terminate_fraction", 0.789),
    ],
)
@gen_cluster(nthreads=[])
async def test_deprecated_attributes(s, cls, name, value):
    async with cls(s.address) as a:
        with pytest.warns(FutureWarning, match=name):
            setattr(a, name, value)
        with pytest.warns(FutureWarning, match=name):
            assert getattr(a, name) == value
        assert getattr(a.memory_manager, name) == value


@gen_cluster(nthreads=[("", 1)])
async def test_deprecated_memory_monitor_method_worker(s, a):
    with pytest.warns(FutureWarning, match="memory_monitor"):
        await a.memory_monitor()


@gen_cluster(nthreads=[("", 1)], Worker=Nanny)
async def test_deprecated_memory_monitor_method_nanny(s, a):
    with pytest.warns(FutureWarning, match="memory_monitor"):
        a.memory_monitor()


@pytest.mark.parametrize(
    "name",
    ["memory_target_fraction", "memory_spill_fraction", "memory_pause_fraction"],
)
@gen_cluster(nthreads=[])
async def test_deprecated_params(s, name):
    with pytest.warns(FutureWarning, match=name):
        async with Worker(s.address, **{name: 0.789}) as a:
            assert getattr(a.memory_manager, name) == 0.789