from __future__ import annotations

import asyncio
import contextlib
import errno
import functools
import logging
import multiprocessing
import os
import shutil
import threading
import uuid
import warnings
import weakref
from collections.abc import Callable, Collection
from inspect import isawaitable
from queue import Empty
from typing import ClassVar, Literal, cast

from toolz import merge
from tornado.ioloop import IOLoop

import dask
from dask.system import CPU_COUNT
from dask.utils import parse_timedelta

from distributed import preloading
from distributed._async_taskgroup import AsyncTaskGroupClosedError
from distributed.comm import get_address_host
from distributed.comm.addressing import address_from_user_args
from distributed.compatibility import asyncio_run
from distributed.config import get_loop_factory
from distributed.core import (
    CommClosedError,
    ErrorMessage,
    OKMessage,
    RPCClosed,
    Status,
    coerce_to_address,
    error_message,
)
from distributed.diagnostics.plugin import NannyPlugin, _get_plugin_name
from distributed.metrics import time
from distributed.node import ServerNode
from distributed.process import AsyncProcess
from distributed.proctitle import enable_proctitle_on_children
from distributed.protocol import pickle
from distributed.protocol.serialize import _is_dumpable
from distributed.security import Security
from distributed.utils import (
    get_ip,
    get_mp_context,
    json_load_robust,
    log_errors,
    parse_ports,
    silence_logging_cmgr,
    wait_for,
)
from distributed.worker import Worker, run
from distributed.worker_memory import (
    DeprecatedMemoryManagerAttribute,
    DeprecatedMemoryMonitor,
    NannyMemoryManager,
)

logger = logging.getLogger(__name__)


class Nanny(ServerNode):
    """A process to manage worker processes

    The nanny spins up Worker processes, watches them, and kills or restarts
    them as necessary. It is necessary if you want to use the
    ``Client.restart`` method, or to restart the worker automatically if
    it gets to the terminate fraction of its memory limit.

    The parameters for the Nanny are mostly the same as those for the Worker
    with exceptions listed below.

    Parameters
    ----------
    env: dict, optional
        Environment variables set at time of Nanny initialization will be
        ensured to be set in the Worker process as well. This argument allows to
        overwrite or otherwise set environment variables for the Worker. It is
        also possible to set environment variables using the option
        ``distributed.nanny.environ``. Precedence as follows

            1. Nanny arguments
            2. Existing environment variables
            3. Dask configuration

        .. note::
           Some environment variables, like ``OMP_NUM_THREADS``, must be set before
           importing numpy to have effect. Others, like ``MALLOC_TRIM_THRESHOLD_`` (see
           :ref:`memtrim`), must be set before starting the Linux process. Such
           variables would be ineffective if set here or in
           ``distributed.nanny.environ``; they must be set in
           ``distributed.nanny.pre-spawn-environ`` so that they are set before spawning
           the subprocess, even if this means poisoning the process running the Nanny.

           For the same reason, be warned that changing
           ``distributed.worker.multiprocessing-method`` from ``spawn`` to ``fork`` or
           ``forkserver`` may inhibit some environment variables; if you do, you should
           set the variables yourself in the shell before you start ``dask-worker``.

    See Also
    --------
    Worker
    """

    _instances: ClassVar[weakref.WeakSet[Nanny]] = weakref.WeakSet()
    process: WorkerProcess | None
    memory_manager: NannyMemoryManager

    env: dict[str, str]
    pre_spawn_env: dict[str, str]

    # Inputs to parse_ports()
    _given_worker_port: int | str | Collection[int] | None
    _start_port: int | str | Collection[int] | None

    def __init__(  # type: ignore[no-untyped-def]
        self,
        scheduler_ip=None,
        scheduler_port=None,
        scheduler_file=None,
        worker_port: int | str | Collection[int] | None = 0,
        nthreads=None,
        loop=None,
        local_directory=None,
        services=None,
        name=None,
        memory_limit="auto",
        reconnect=True,
        validate=False,
        quiet=False,
        resources=None,
        silence_logs=None,
        death_timeout=None,
        preload=None,
        preload_argv=None,
        preload_nanny=None,
        preload_nanny_argv=None,
        security=None,
        contact_address=None,
        listen_address=None,
        worker_class=None,
        env=None,
        interface=None,
        host=None,
        port: int | str | Collection[int] | None = None,
        protocol=None,
        config=None,
        **worker_kwargs,
    ):
        if loop is not None:
            warnings.warn(
                "the `loop` kwarg to `Nanny` is ignored, and will be removed in a future release. "
                "The Nanny always binds to the current loop.",
                DeprecationWarning,
                stacklevel=2,
            )

        self.__exit_stack = stack = contextlib.ExitStack()
        self.process = None
        self._setup_logging(logger)
        self.loop = self.io_loop = IOLoop.current()

        if isinstance(security, dict):
            security = Security(**security)
        self.security = security or Security()
        assert isinstance(self.security, Security)
        self.connection_args = self.security.get_connection_args("worker")

        self.preload = preload
        if self.preload is None:
            self.preload = dask.config.get("distributed.worker.preload")
        self.preload_argv = preload_argv
        if self.preload_argv is None:
            self.preload_argv = dask.config.get("distributed.worker.preload-argv")

        if preload_nanny is None:
            preload_nanny = dask.config.get("distributed.nanny.preload")
        if preload_nanny_argv is None:
            preload_nanny_argv = dask.config.get("distributed.nanny.preload-argv")

        handlers = {
            "instantiate": self.instantiate,
            "kill": self.kill,
            "restart": self.restart,
            "get_logs": self.get_logs,
            # cannot call it 'close' on the rpc side for naming conflict
            "terminate": self.close,
            "close_gracefully": self.close_gracefully,
            "run": self.run,
            "plugin_add": self.plugin_add,
            "plugin_remove": self.plugin_remove,
        }
        super().__init__(
            handlers=handlers,
            connection_args=self.connection_args,
            local_directory=local_directory,
            needs_workdir=False,
        )

        self.preloads = preloading.process_preloads(
            self, preload_nanny, preload_nanny_argv, file_dir=self.local_directory
        )

        self.death_timeout = parse_timedelta(death_timeout)
        if scheduler_file:
            cfg = json_load_robust(scheduler_file, timeout=self.death_timeout)
            self.scheduler_addr = cfg["address"]
        elif scheduler_ip is None and dask.config.get("scheduler-address"):
            self.scheduler_addr = dask.config.get("scheduler-address")
        elif scheduler_port is None:
            self.scheduler_addr = coerce_to_address(scheduler_ip)
        else:
            self.scheduler_addr = coerce_to_address((scheduler_ip, scheduler_port))

        if protocol is None:
            protocol_address = self.scheduler_addr.split("://")
            if len(protocol_address) == 2:
                protocol = protocol_address[0]

        self._given_worker_port = worker_port
        self.nthreads: int = nthreads or CPU_COUNT
        self.reconnect = reconnect
        self.validate = validate
        self.resources = resources

        self.Worker = Worker if worker_class is None else worker_class

        self.pre_spawn_env = _get_env_variables("distributed.nanny.pre-spawn-environ")
        # To get consistent hashing on subprocesses, we need to set a consistent seed for
        # the Python hash algorithm; xref https://github.com/dask/distributed/pull/8400
        if self.pre_spawn_env.get("PYTHONHASHSEED") in (None, "0"):
            # This number is arbitrary; it was chosen to commemorate
            # https://github.com/dask/dask/issues/6640.
            self.pre_spawn_env.update({"PYTHONHASHSEED": "6640"})

        self.env = merge(
            self.pre_spawn_env,
            _get_env_variables("distributed.nanny.environ"),
            {k: str(v) for k, v in env.items()} if env else {},
        )

        self.config = merge(dask.config.config, config or {})
        worker_kwargs.update(
            {
                "port": worker_port,
                "interface": interface,
                "protocol": protocol,
                "host": host,
            }
        )
        self.worker_kwargs = worker_kwargs

        self.contact_address = contact_address

        self.services = services
        self.name = name
        self.quiet = quiet

        if silence_logs:
            stack.enter_context(silence_logging_cmgr(level=silence_logs))
        self.silence_logs = silence_logs

        self.plugins: dict[str, NannyPlugin] = {}
        self.scheduler = self.rpc(self.scheduler_addr)
        self.memory_manager = NannyMemoryManager(self, memory_limit=memory_limit)

        if (
            not host
            and not interface
            and not self.scheduler_addr.startswith("inproc://")
        ):
            host = get_ip(get_address_host(self.scheduler.address))

        self._start_port = port
        self._start_host = host
        self._interface = interface
        self._protocol = protocol

        self._listen_address = listen_address
        Nanny._instances.add(self)

    # Deprecated attributes; use Nanny.memory_manager.<name> instead
    memory_limit = DeprecatedMemoryManagerAttribute()
    memory_terminate_fraction = DeprecatedMemoryManagerAttribute()
    memory_monitor = DeprecatedMemoryMonitor()

    def __repr__(self):
        return "<Nanny: %s, threads: %d>" % (self.worker_address, self.nthreads)

    async def _unregister(self, timeout=10):
        if self.process is None:
            return
        worker_address = self.process.worker_address
        if worker_address is None:
            return

        try:
            await wait_for(
                self.scheduler.unregister(
                    address=self.worker_address, stimulus_id=f"nanny-close-{time()}"
                ),
                timeout,
            )
        except (asyncio.TimeoutError, CommClosedError, OSError, RPCClosed):
            pass

    @property
    def worker_address(self):
        return None if self.process is None else self.process.worker_address

    @property
    def worker_dir(self):
        return None if self.process is None else self.process.worker_dir

    async def start_unsafe(self):
        """Start nanny, start local process, start watching"""

        await super().start_unsafe()

        ports = parse_ports(self._start_port)
        for port in ports:
            start_address = address_from_user_args(
                host=self._start_host,
                port=port,
                interface=self._interface,
                protocol=self._protocol,
                security=self.security,
            )
            try:
                await self.listen(
                    start_address, **self.security.get_listen_args("worker")
                )
            except OSError as e:
                if len(ports) > 1 and e.errno == errno.EADDRINUSE:
                    continue
                else:
                    raise
            else:
                self._start_address = start_address
                break
        else:
            raise ValueError(
                f"Could not start Nanny on host {self._start_host} "
                f"with port {self._start_port}"
            )

        self.ip = get_address_host(self.address)

        await self.preloads.start()
        saddr = self.scheduler.addr
        comm = await self.rpc.connect(saddr)
        comm.name = "Nanny->Scheduler (registration)"

        try:
            await comm.write({"op": "register_nanny", "address": self.address})
            msg = await comm.read()
            try:
                for name, plugin in msg["nanny-plugins"].items():
                    await self.plugin_add(plugin=plugin, name=name)

                logger.info("        Start Nanny at: %r", self.address)
                response = await self.instantiate()

                if response != Status.running:
                    raise RuntimeError("Nanny failed to start worker process")
            except Exception:
                try:
                    await comm.write({"status": "error"})

                # If self.instantiate() failed, the comm will already be closed.
                except CommClosedError:
                    pass
                await self.close(reason="nanny-start-failed")
                raise
            else:
                await comm.write({"status": "ok"})
        finally:
            await comm.close()

        assert self.worker_address

        self.start_periodic_callbacks()

        return self

    async def kill(self, timeout: float = 5, reason: str = "nanny-kill") -> None:
        """Kill the local worker process

        Blocks until both the process is down and the scheduler is properly
        informed
        """
        if self.process is not None:
            await self.process.kill(reason=reason, timeout=timeout)

    async def instantiate(self) -> Status:
        """Start a local worker process

        Blocks until the process is up and the scheduler is properly informed
        """
        if self.process is None:
            worker_kwargs = dict(
                scheduler_ip=self.scheduler_addr,
                nthreads=self.nthreads,
                local_directory=self._original_local_dir,
                services=self.services,
                nanny=self.address,
                name=self.name,
                memory_limit=self.memory_manager.memory_limit,
                resources=self.resources,
                validate=self.validate,
                silence_logs=self.silence_logs,
                death_timeout=self.death_timeout,
                preload=self.preload,
                preload_argv=self.preload_argv,
                security=self.security,
                contact_address=self.contact_address,
            )
            worker_kwargs.update(self.worker_kwargs)
            self.process = WorkerProcess(
                worker_kwargs=worker_kwargs,
                silence_logs=self.silence_logs,
                on_exit=self._on_worker_exit_sync,
                worker=self.Worker,
                env=self.env,
                pre_spawn_env=self.pre_spawn_env,
                config=self.config,
            )

        if self.death_timeout:
            try:
                result = await wait_for(self.process.start(), self.death_timeout)
            except asyncio.TimeoutError:
                logger.error(
                    "Timed out connecting Nanny '%s' to scheduler '%s'",
                    self,
                    self.scheduler_addr,
                )
                await self.close(
                    timeout=self.death_timeout, reason="nanny-instantiate-timeout"
                )
                raise

        else:
            try:
                result = await self.process.start()
            except Exception:
                logger.error("Failed to start process", exc_info=True)
                await self.close(reason="nanny-instantiate-failed")
                raise
        return result

    @log_errors
    async def plugin_add(
        self, plugin: NannyPlugin | bytes, name: str | None = None
    ) -> ErrorMessage | OKMessage:
        if isinstance(plugin, bytes):
            plugin = pickle.loads(plugin)
        if not isinstance(plugin, NannyPlugin):
            warnings.warn(
                "Registering duck-typed plugins has been deprecated. "
                "Please make sure your plugin inherits from `NannyPlugin`.",
                DeprecationWarning,
                stacklevel=2,
            )
        plugin = cast(NannyPlugin, plugin)

        if name is None:
            name = _get_plugin_name(plugin)
        assert name

        self.plugins[name] = plugin

        logger.info("Starting Nanny plugin %s", name)
        if hasattr(plugin, "setup"):
            try:
                result = plugin.setup(nanny=self)
                if isawaitable(result):
                    result = await result
            except Exception as e:
                logger.exception("Nanny plugin %s failed to setup", name)
                return error_message(e)
        if getattr(plugin, "restart", False):
            await self.restart(reason=f"nanny-plugin-{name}-restart")

        return {"status": "OK"}

    @log_errors
    async def plugin_remove(self, name: str) -> ErrorMessage | OKMessage:
        logger.info(f"Removing Nanny plugin {name}")
        try:
            plugin = self.plugins.pop(name)
            if hasattr(plugin, "teardown"):
                result = plugin.teardown(nanny=self)
                if isawaitable(result):
                    result = await result
        except Exception as e:
            logger.exception("Nanny plugin %s failed to teardown", name)
            msg = error_message(e)
            return msg

        return {"status": "OK"}

    async def restart(
        self, timeout: float = 30, reason: str = "nanny-restart"
    ) -> Literal["OK", "timed out"] | ErrorMessage:
        async def _():
            if self.process is not None:
                await self.kill(reason=reason)
                await self.instantiate()

        try:
            await wait_for(asyncio.shield(_()), timeout)
        except asyncio.TimeoutError:
            logger.error(
                f"Restart timed out after {timeout}s; returning before finished"
            )
            return "timed out"
        except Exception as e:
            return error_message(e)
        else:
            return "OK"

    def is_alive(self):
        return self.process is not None and self.process.is_alive()

    def run(self, comm, *args, **kwargs):
        return run(self, comm, *args, **kwargs)

    def _on_worker_exit_sync(self, exitcode):
        try:
            self._ongoing_background_tasks.call_soon(self._on_worker_exit, exitcode)
        except (
            AsyncTaskGroupClosedError
        ):  # Async task group has already been closed, so the nanny is already clos(ed|ing).
            pass

    @log_errors
    async def _on_worker_exit(self, exitcode):
        if self.status not in (
            Status.init,
            Status.closing,
            Status.closed,
            Status.closing_gracefully,
            Status.failed,
        ):
            try:
                await self._unregister()
            except OSError:
                logger.exception("Failed to unregister")
                if not self.reconnect:
                    await self.close(reason="nanny-unregister-failed")
                    return

        try:
            if self.status not in (
                Status.closing,
                Status.closed,
                Status.closing_gracefully,
                Status.failed,
            ):
                logger.warning("Restarting worker")
                await self.instantiate()
            elif self.status == Status.closing_gracefully:
                await self.close(reason="nanny-close-gracefully")

        except Exception:
            logger.error(
                "Failed to restart worker after its process exited", exc_info=True
            )

    @property
    def pid(self):
        return self.process and self.process.pid

    def _close(self, *args, **kwargs):
        warnings.warn("Worker._close has moved to Worker.close", stacklevel=2)
        return self.close(*args, **kwargs)

    def close_gracefully(self, reason: str = "nanny-close-gracefully") -> None:
        """
        A signal that we shouldn't try to restart workers if they go away

        This is used as part of the cluster shutdown process.
        """
        self.status = Status.closing_gracefully
        logger.info(
            "Closing Nanny gracefully at %r. Reason: %s", self.address_safe, reason
        )

    async def close(  # type:ignore[override]
        self, timeout: float = 5, reason: str = "nanny-close"
    ) -> Literal["OK"]:
        """
        Close the worker process, stop all comms.
        """
        if self.status == Status.closing:
            await self.finished()
            assert self.status == Status.closed

        if self.status == Status.closed:
            return "OK"

        self.status = Status.closing
        logger.info("Closing Nanny at %r. Reason: %s", self.address_safe, reason)

        await self.preloads.teardown()

        await asyncio.gather(*(self.plugin_remove(name) for name in self.plugins))

        self.stop()
        if self.process is not None:
            await self.kill(timeout=timeout, reason=reason)

        self.process = None
        await self.rpc.close()
        self.status = Status.closed
        await super().close()
        self.__exit_stack.__exit__(None, None, None)
        logger.info("Nanny at %r closed.", self.address_safe)
        return "OK"

    async def _log_event(self, topic, msg):
        await self.scheduler.log_event(
            topic=topic,
            msg=msg,
        )

    def log_event(self, topic, msg):
        """Log an event under a given topic

        Parameters
        ----------
        topic : str, list[str]
            Name of the topic under which to log an event. To log the same
            event under multiple topics, pass a list of topic names.
        msg
            Event message to log. Note this must be msgpack serializable.

        See also
        --------
        Client.log_event
        """
        if not _is_dumpable(msg):
            raise TypeError(
                f"Message must be msgpack serializable. Got {type(msg)=} instead."
            )
        self._ongoing_background_tasks.call_soon(self._log_event, topic, msg)


class WorkerProcess:
    running: asyncio.Event
    stopped: asyncio.Event

    process: AsyncProcess | None
    env: dict[str, str]
    pre_spawn_env: dict[str, str]

    # The interval how often to check the msg queue for init
    _init_msg_interval = 0.05

    def __init__(
        self,
        worker_kwargs,
        silence_logs,
        on_exit,
        worker,
        env,
        pre_spawn_env,
        config,
    ):
        self.status = Status.init
        self.silence_logs = silence_logs
        self.worker_kwargs = worker_kwargs
        self.on_exit = on_exit
        self.process = None
        self.Worker = worker
        self.env = env
        self.pre_spawn_env = pre_spawn_env
        self.config = config.copy()

        # Ensure default clients don't propagate to subprocesses
        try:
            from distributed.client import default_client

            default_client()
            self.config.pop("scheduler", None)
            self.config.pop("shuffle", None)
        except ValueError:
            pass

        # Initialized when worker is ready
        self.worker_dir = None
        self.worker_address = None

    async def start(self) -> Status:
        """
        Ensure the worker process is started.
        """
        enable_proctitle_on_children()
        if self.status == Status.running:
            return self.status
        if self.status == Status.starting:
            await self.running.wait()
            return self.status
        mp_ctx = get_mp_context()
        self.init_result_q = mp_ctx.Queue()
        self.child_stop_q = mp_ctx.Queue()
        # put an empty message to start the background self._child_stop_q._thread
        # otherwise calling .put({"op": "stop"}) later in an atexit handler
        # results in a RuntimeError starting the thread
        # get the message before passing child_stop_q to AsyncProcess to avoid
        # dealing with it in watch_stop_q
        self.child_stop_q.put(None)
        self.child_stop_q.get()
        uid = uuid.uuid4().hex

        self.process = AsyncProcess(
            target=functools.partial(
                self._run,
                silence_logs=self.silence_logs,
                init_result_q=self.init_result_q,
                child_stop_q=self.child_stop_q,
                uid=uid,
                worker_factory=functools.partial(self.Worker, **self.worker_kwargs),
                env=self.env,
                config=self.config,
            ),
            name="Dask Worker process (from Nanny)",
            kwargs=dict(),
        )
        self.process.daemon = dask.config.get("distributed.worker.daemon", default=True)
        self.process.set_exit_callback(self._on_exit)
        self.running = asyncio.Event()
        self.stopped = asyncio.Event()
        self.status = Status.starting

        # Set selected environment variables before spawning the subprocess.
        # See note in Nanny docstring.
        os.environ.update(self.pre_spawn_env)

        try:
            try:
                await self.process.start()
            except OSError:
                # This can only happen if the actual process creation failed, e.g.
                # multiprocessing.Process.start failed. This is not tested!
                logger.exception("Nanny failed to start process", exc_info=True)
                # NOTE: doesn't wait for process to terminate, just for terminate signal to be sent
                await self.process.terminate()
                self.status = Status.failed
            try:
                msg = await self._wait_until_connected(uid)
            except Exception:
                # NOTE: doesn't wait for process to terminate, just for terminate signal to be sent
                await self.process.terminate()
                self.status = Status.failed
                raise
        finally:
            self.running.set()
        if not msg:
            return self.status
        self.worker_address = msg["address"]
        self.worker_dir = msg["dir"]
        assert self.worker_address
        self.status = Status.running

        return self.status

    def _on_exit(self, proc):
        if proc is not self.process:
            # Ignore exit of old process instance
            return
        self.mark_stopped()

    def _death_message(self, pid, exitcode):
        assert exitcode is not None
        if exitcode == 255:
            return "Worker process %d was killed by unknown signal" % (pid,)
        elif exitcode >= 0:
            return "Worker process %d exited with status %d" % (pid, exitcode)
        else:
            return "Worker process %d was killed by signal %d" % (pid, -exitcode)

    def is_alive(self):
        return self.process is not None and self.process.is_alive()

    @property
    def pid(self):
        return self.process.pid if self.process and self.process.is_alive() else None

    def mark_stopped(self):
        if self.status != Status.stopped:
            assert self.process is not None
            r = self.process.exitcode
            assert r is not None
            if r != 0:
                msg = self._death_message(self.process.pid, r)
                logger.info(msg)
            self.status = Status.stopped
            self.running.clear()
            self.stopped.set()
            # Release resources
            self.process.close()
            self.init_result_q = None
            self.child_stop_q = None
            self.process = None
            # Best effort to clean up worker directory
            if self.worker_dir and os.path.exists(self.worker_dir):
                shutil.rmtree(self.worker_dir, ignore_errors=True)
            self.worker_dir = None
            # User hook
            if self.on_exit is not None:
                self.on_exit(r)

    async def kill(
        self,
        timeout: float = 5,
        executor_wait: bool = True,
        reason: str = "workerprocess-kill",
    ) -> None:
        """
        Ensure the worker process is stopped, waiting at most
        ``timeout * 0.8`` seconds before killing it abruptly.

        When `kill` returns, the worker process has been joined.

        If the worker process does not terminate within ``timeout`` seconds,
        even after being killed, `asyncio.TimeoutError` is raised.
        """
        deadline = time() + timeout

        # If the process is not properly up it will not watch the closing queue
        # and we may end up leaking this process
        # Therefore wait for it to be properly started before killing it
        if self.status == Status.starting:
            await self.running.wait()

        assert self.status in (
            Status.stopping,
            Status.stopped,
            Status.running,
            Status.failed,  # process failed to start, but hasn't been joined yet
            Status.closing_gracefully,
        ), self.status
        if self.status == Status.stopped:
            return
        if self.status == Status.stopping:
            await self.stopped.wait()
            return
        self.status = Status.stopping
        logger.info("Nanny asking worker to close. Reason: %s", reason)

        process = self.process
        wait_timeout = timeout * 0.8
        self.child_stop_q.put(
            {
                "op": "stop",
                "timeout": wait_timeout,
                "executor_wait": executor_wait,
                "reason": reason,
            }
        )
        self.child_stop_q.close()
        assert process is not None
        try:
            try:
                await process.join(wait_timeout)
                return
            except asyncio.TimeoutError:
                pass

            logger.warning(
                f"Worker process still alive after {wait_timeout:.1f} seconds, killing"
            )
            await process.kill()
            await process.join(max(0, deadline - time()))
        except ValueError as e:
            if "invalid operation on closed AsyncProcess" in str(e):
                return
            raise

    async def _wait_until_connected(self, uid):
        while True:
            if self.status != Status.starting:
                return
            # This is a multiprocessing queue and we'd block the event loop if
            # we simply called get
            try:
                msg = self.init_result_q.get_nowait()
            except Empty:
                await asyncio.sleep(self._init_msg_interval)
                continue

            if msg["uid"] != uid:  # ensure that we didn't cross queues
                continue

            if "exception" in msg:
                raise msg["exception"]
            else:
                return msg

    @classmethod
    def _run(
        cls,
        silence_logs: bool,
        init_result_q: multiprocessing.Queue,
        child_stop_q: multiprocessing.Queue,
        uid: str,
        env: dict,
        config: dict,
        worker_factory: Callable[[], Worker],
    ) -> None:  # pragma: no cover
        async def do_stop(
            *,
            worker: Worker,
            timeout: float = 5,
            executor_wait: bool = True,
            reason: str = "workerprocess-stop",
        ) -> None:
            await worker.close(
                nanny=False,
                executor_wait=executor_wait,
                timeout=timeout,
                reason=reason,
            )

        def watch_stop_q(loop: IOLoop, worker: Worker) -> None:
            """
            Wait for an incoming stop message and then stop the
            worker cleanly.
            """
            try:
                msg = child_stop_q.get()
            except (TypeError, OSError, EOFError):
                logger.error("Worker process died unexpectedly")
                msg = {"op": "stop"}
            finally:
                child_stop_q.close()
                assert msg["op"] == "stop", msg
                del msg["op"]
                loop.add_callback(do_stop, worker=worker, **msg)

        async def run() -> None:
            """
            Try to start worker and inform parent of outcome.
            """
            failure_type: str | None = "initialize"
            try:
                worker = worker_factory()
                failure_type = "start"
                thread = threading.Thread(
                    target=functools.partial(
                        watch_stop_q,
                        worker=worker,
                        loop=IOLoop.current(),
                    ),
                    name="Nanny stop queue watch",
                    daemon=True,
                )
                thread.start()
                stack.callback(thread.join, timeout=2)
                async with worker:
                    failure_type = None

                    try:
                        assert worker.address
                    except ValueError:
                        pass
                    else:
                        init_result_q.put(
                            {
                                "address": worker.address,
                                "dir": worker.local_directory,
                                "uid": uid,
                            }
                        )
                        init_result_q.close()
                        await worker.finished()
                        logger.info("Worker closed")
            except Exception as e:
                if failure_type is None:
                    raise

                logger.exception(f"Failed to {failure_type} worker")
                init_result_q.put({"uid": uid, "exception": e})
                init_result_q.close()
            finally:
                init_result_q.join_thread()

        with contextlib.ExitStack() as stack:

            @stack.callback
            def close_stop_q() -> None:
                try:
                    child_stop_q.put({"op": "stop"})  # usually redundant
                except ValueError:
                    pass

                try:
                    child_stop_q.close()  # usually redundant
                except ValueError:
                    pass
                child_stop_q.join_thread()

            os.environ.update(env)
            dask.config.refresh()
            dask.config.set(config)

            from dask.multiprocessing import default_initializer

            default_initializer()

            if silence_logs:
                logger.setLevel(silence_logs)

            asyncio_run(run(), loop_factory=get_loop_factory())


def _get_env_variables(config_key: str) -> dict[str, str]:
    cfg = dask.config.get(config_key)
    if not isinstance(cfg, dict):
        raise TypeError(  # pragma: nocover
            f"{config_key} configuration must be of type dict. Instead got {type(cfg)}"
        )
    # Override dask config with explicitly defined env variables from the OS
    # Allow unsetting a variable in a config override by setting its value to None.
    cfg = {k: os.environ.get(k, str(v)) for k, v in cfg.items() if v is not None}
    return cfg