Get scheduler function There are various ways to specify the scheduler to use: 1. Passing in scheduler= parameters 2. Passing these into global configuration 3. Using a dask.distributed default Client 4. Using defaults of a dask collection This function centralizes the
(get=None, scheduler=None, collections=None, cls=None)
| 1223 | |
| 1224 | |
| 1225 | def get_scheduler(get=None, scheduler=None, collections=None, cls=None): |
| 1226 | """Get scheduler function |
| 1227 | |
| 1228 | There are various ways to specify the scheduler to use: |
| 1229 | |
| 1230 | 1. Passing in scheduler= parameters |
| 1231 | 2. Passing these into global configuration |
| 1232 | 3. Using a dask.distributed default Client |
| 1233 | 4. Using defaults of a dask collection |
| 1234 | |
| 1235 | This function centralizes the logic to determine the right scheduler to use |
| 1236 | from those many options |
| 1237 | """ |
| 1238 | if get: |
| 1239 | raise TypeError(get_err_msg) |
| 1240 | |
| 1241 | if scheduler is not None: |
| 1242 | if callable(scheduler): |
| 1243 | return scheduler |
| 1244 | elif "Client" in type(scheduler).__name__ and hasattr(scheduler, "get"): |
| 1245 | return _ensure_not_async(scheduler) |
| 1246 | elif isinstance(scheduler, str): |
| 1247 | scheduler = scheduler.lower() |
| 1248 | |
| 1249 | client_available = False |
| 1250 | if _distributed_available(): |
| 1251 | assert _DistributedClient is not None |
| 1252 | with suppress(ValueError): |
| 1253 | _DistributedClient.current(allow_global=True) |
| 1254 | client_available = True |
| 1255 | if scheduler in named_schedulers: |
| 1256 | return named_schedulers[scheduler] |
| 1257 | elif scheduler in ("dask.distributed", "distributed"): |
| 1258 | if not client_available: |
| 1259 | raise RuntimeError( |
| 1260 | f"Requested {scheduler} scheduler but no Client active." |
| 1261 | ) |
| 1262 | assert _get_distributed_client is not None |
| 1263 | client = _get_distributed_client() |
| 1264 | return _ensure_not_async(client) |
| 1265 | else: |
| 1266 | raise ValueError( |
| 1267 | "Expected one of [distributed, {}]".format( |
| 1268 | ", ".join(sorted(named_schedulers)) |
| 1269 | ) |
| 1270 | ) |
| 1271 | elif isinstance(scheduler, Executor): |
| 1272 | # Get `num_workers` from `Executor`'s `_max_workers` attribute. |
| 1273 | # If undefined, fallback to `config` or worst case CPU_COUNT. |
| 1274 | num_workers = getattr(scheduler, "_max_workers", None) |
| 1275 | if num_workers is None: |
| 1276 | num_workers = config.get("num_workers", CPU_COUNT) |
| 1277 | assert isinstance(num_workers, Integral) and num_workers > 0 |
| 1278 | return partial(local.get_async, scheduler.submit, num_workers) |
| 1279 | else: |
| 1280 | raise ValueError(f"Unexpected scheduler: {scheduler!r}") |
| 1281 | # else: # try to connect to remote scheduler with this name |
| 1282 | # return get_client(scheduler).get |