-
-
Notifications
You must be signed in to change notification settings - Fork 711
/
config.py
215 lines (181 loc) · 7.64 KB
/
config.py
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
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",
}
# 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",
"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)