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)
| 1111 | |
| 1112 | |
| 1113 | def get_scheduler(get=None, scheduler=None, collections=None, cls=None): |
| 1114 | """Get scheduler function |
| 1115 | |
| 1116 | There are various ways to specify the scheduler to use: |
| 1117 | |
| 1118 | 1. Passing in scheduler= parameters |
| 1119 | 2. Passing these into global configuration |
| 1120 | 3. Using a dask.distributed default Client |
| 1121 | 4. Using defaults of a dask collection |
| 1122 | |
| 1123 | This function centralizes the logic to determine the right scheduler to use |
| 1124 | from those many options |
| 1125 | """ |
| 1126 | if get: |
| 1127 | raise TypeError(get_err_msg) |
| 1128 | |
| 1129 | if scheduler is not None: |
| 1130 | if callable(scheduler): |
| 1131 | return scheduler |
| 1132 | elif "Client" in type(scheduler).__name__ and hasattr(scheduler, "get"): |
| 1133 | return _ensure_not_async(scheduler) |
| 1134 | elif isinstance(scheduler, str): |
| 1135 | scheduler = scheduler.lower() |
| 1136 | |
| 1137 | client_available = False |
| 1138 | if _distributed_available(): |
| 1139 | assert _DistributedClient is not None |
| 1140 | with suppress(ValueError): |
| 1141 | _DistributedClient.current(allow_global=True) |
| 1142 | client_available = True |
| 1143 | if scheduler in named_schedulers: |
| 1144 | return named_schedulers[scheduler] |
| 1145 | elif scheduler in ("dask.distributed", "distributed"): |
| 1146 | if not client_available: |
| 1147 | raise RuntimeError( |
| 1148 | f"Requested {scheduler} scheduler but no Client active." |
| 1149 | ) |
| 1150 | assert _get_distributed_client is not None |
| 1151 | client = _get_distributed_client() |
| 1152 | return _ensure_not_async(client) |
| 1153 | else: |
| 1154 | raise ValueError( |
| 1155 | "Expected one of [distributed, %s]" |
| 1156 | % ", ".join(sorted(named_schedulers)) |
| 1157 | ) |
| 1158 | elif isinstance(scheduler, Executor): |
| 1159 | # Get `num_workers` from `Executor`'s `_max_workers` attribute. |
| 1160 | # If undefined, fallback to `config` or worst case CPU_COUNT. |
| 1161 | num_workers = getattr(scheduler, "_max_workers", None) |
| 1162 | if num_workers is None: |
| 1163 | num_workers = config.get("num_workers", CPU_COUNT) |
| 1164 | assert isinstance(num_workers, Integral) and num_workers > 0 |
| 1165 | return partial(local.get_async, scheduler.submit, num_workers) |
| 1166 | else: |
| 1167 | raise ValueError("Unexpected scheduler: %s" % repr(scheduler)) |
| 1168 | # else: # try to connect to remote scheduler with this name |
| 1169 | # return get_client(scheduler).get |
| 1170 |