File: config.py

package info (click to toggle)
dask.distributed 2024.12.1%2Bds-1
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid, trixie
  • size: 12,588 kB
  • sloc: python: 96,954; javascript: 1,549; sh: 390; makefile: 220
file content (217 lines) | stat: -rw-r--r-- 7,923 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
from __future__ import annotations

import asyncio
import logging.config
import os
import sys
from collections.abc import Callable
from typing import Any

import yaml

import dask
from dask.utils import import_required

config = dask.config.config


fn = os.path.join(os.path.dirname(__file__), "distributed.yaml")

with open(fn) as f:
    defaults = yaml.safe_load(f)

dask.config.update_defaults(defaults)

deprecations = {
    "allowed-failures": "distributed.scheduler.allowed-failures",
    "bandwidth": "distributed.scheduler.bandwidth",
    "default-data-size": "distributed.scheduler.default-data-size",
    "work-stealing": "distributed.scheduler.work-stealing",
    "worker-ttl": "distributed.scheduler.worker-ttl",
    "multiprocessing-method": "distributed.worker.multiprocessing-method",
    "use-file-locking": "distributed.worker.use-file-locking",
    "profile-interval": "distributed.worker.profile.interval",
    "profile-cycle-interval": "distributed.worker.profile.cycle",
    "worker-memory-target": "distributed.worker.memory.target",
    "worker-memory-spill": "distributed.worker.memory.spill",
    "worker-memory-pause": "distributed.worker.memory.pause",
    "worker-memory-terminate": "distributed.worker.memory.terminate",
    "heartbeat-interval": "distributed.client.heartbeat",
    "compression": "distributed.comm.compression",
    "connect-timeout": "distributed.comm.timeouts.connect",
    "tcp-timeout": "distributed.comm.timeouts.tcp",
    "default-scheme": "distributed.comm.default-scheme",
    "socket-backlog": "distributed.comm.socket-backlog",
    "diagnostics-link": "distributed.dashboard.link",
    "bokeh-export-tool": "distributed.dashboard.export-tool",
    "tick-time": "distributed.admin.tick.interval",
    "tick-maximum-delay": "distributed.admin.tick.limit",
    "log-length": "distributed.admin.log-length",
    "log-format": "distributed.admin.log-format",
    "pdb-on-err": "distributed.admin.pdb-on-err",
    "ucx": "distributed.comm.ucx",
    "rmm": "distributed.rmm",
    # low-level-log-length aliases
    "transition-log-length": "distributed.admin.low-level-log-length",
    "distributed.scheduler.transition-log-length": "distributed.admin.low-level-log-length",
    "distributed.scheduler.events-log-length": "distributed.admin.low-level-log-length",
    "recent-messages-log-length": "distributed.admin.low-level-log-length",
    "distributed.comm.recent-messages-log-length": "distributed.admin.low-level-log-length",
    "distributed.p2p.disk": "distributed.p2p.storage.disk",
}

# Affects yaml and env variables configs, as well as calls to dask.config.set()
# before importing distributed
dask.config.rename(deprecations)
# Affects dask.config.set() from now on
dask.config.deprecations.update(deprecations)


#########################
# Logging specific code #
#########################
#
# Here we enact the policies in the logging part of the configuration

logger = logging.getLogger(__name__)


if sys.version_info >= (3, 11):
    _logging_get_level_names_mapping = logging.getLevelNamesMapping
else:

    def _logging_get_level_names_mapping() -> dict[str, int]:
        return logging._nameToLevel.copy()


def _initialize_logging_old_style(config: dict[Any, Any]) -> None:
    """
    Initialize logging using the "old-style" configuration scheme, e.g.:
        {
        'logging': {
            'distributed': 'info',
            'tornado': 'critical',
            'tornado.application': 'error',
            }
        }
    """
    loggers: dict[str, str | int] = {  # default values
        "distributed": "info",
        "distributed.client": "warning",
        "distributed.gc": "warning",
        "bokeh": "error",
        "tornado": "critical",
        "tornado.application": "error",
    }
    base_config = _find_logging_config(config)
    loggers.update(base_config.get("logging", {}))

    handler = logging.StreamHandler(sys.stderr)
    handler.setFormatter(
        logging.Formatter(
            dask.config.get("distributed.admin.log-format", config=config)
        )
    )
    logging_names = _logging_get_level_names_mapping()
    for name, raw_level in sorted(loggers.items()):
        level = (
            logging_names[raw_level.upper()]
            if isinstance(raw_level, str)
            else raw_level
        )
        logger = logging.getLogger(name)
        logger.setLevel(level)

        # Ensure that we're not registering the logger twice in this hierarchy.
        anc = logging.getLogger(None)
        already_registered = False

        for ancestor in name.split("."):
            if anc.handlers:
                already_registered = True
                break
            anc.getChild(ancestor)

        if not already_registered:
            logger.addHandler(handler)
            logger.propagate = False


def _initialize_logging_new_style(config: dict[Any, Any]) -> None:
    """
    Initialize logging using logging's "Configuration dictionary schema".
    (ref.: https://docs.python.org/3/library/logging.config.html#configuration-dictionary-schema)
    """
    base_config = _find_logging_config(config)
    logging.config.dictConfig(base_config.get("logging"))  # type: ignore[arg-type]


def _initialize_logging_file_config(config: dict[Any, Any]) -> None:
    """
    Initialize logging using logging's "Configuration file format".
    (ref.: https://docs.python.org/3/howto/logging.html#configuring-logging)
    """
    base_config = _find_logging_config(config)
    logging.config.fileConfig(
        base_config.get("logging-file-config"),  # type: ignore[arg-type]
        disable_existing_loggers=False,
    )


def _find_logging_config(config: dict[Any, Any]) -> dict[Any, Any]:
    """
    Look for the dictionary containing logging-specific configurations,
    starting in the 'distributed' dictionary and then trying the top-level
    """
    logging_keys = {"logging-file-config", "logging"}
    if logging_keys & config.get("distributed", {}).keys():
        return config["distributed"]
    else:
        return config


def initialize_logging(config: dict[Any, Any]) -> None:
    base_config = _find_logging_config(config)
    if "logging-file-config" in base_config:
        if "logging" in base_config:
            raise RuntimeError(
                "Config options 'logging-file-config' and 'logging' are mutually exclusive."
            )
        _initialize_logging_file_config(config)
    else:
        log_config = base_config.get("logging", {})
        if "version" in log_config:
            # logging module mandates version to be an int
            log_config["version"] = int(log_config["version"])
            _initialize_logging_new_style(config)
        else:
            _initialize_logging_old_style(config)


def get_loop_factory() -> Callable[[], asyncio.AbstractEventLoop] | None:
    event_loop = dask.config.get("distributed.admin.event-loop")
    if event_loop == "uvloop":
        uvloop = import_required(
            "uvloop",
            "The distributed.admin.event-loop configuration value "
            "is set to 'uvloop' but the uvloop module is not installed"
            "\n\n"
            "Please either change the config value or install one of the following\n"
            "    conda install uvloop\n"
            "    pip install uvloop",
        )
        return uvloop.new_event_loop
    if event_loop in {"asyncio", "tornado"}:
        if sys.platform == "win32":
            # ProactorEventLoop is not compatible with tornado 6
            # fallback to the pre-3.8 default of Selector
            # https://github.com/tornadoweb/tornado/issues/2608
            return asyncio.SelectorEventLoop
        return None
    raise ValueError(
        "Expected distributed.admin.event-loop to be in ('asyncio', 'tornado', 'uvloop'), got %s"
        % dask.config.get("distributed.admin.event-loop")
    )


initialize_logging(dask.config.config)