-
Notifications
You must be signed in to change notification settings - Fork 1
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
1 changed file
with
46 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,46 @@ | ||
""" | ||
用于启动一个单一主线程的dask worker,防止一些存在线程安全问题的库出现问题 | ||
""" | ||
|
||
import click | ||
from distributed import Worker | ||
from distributed.cli.dask_worker import main | ||
from concurrent.futures import Future, Executor | ||
|
||
|
||
class DummyWorker(Worker): | ||
def __init__(self, *args, **kwargs): | ||
print('Dummy啦') | ||
super().__init__(*args, executor=DummyLinearExecutor(), **kwargs) | ||
|
||
|
||
class DummyLinearExecutor(Executor): | ||
def __init__(self): | ||
pass | ||
|
||
def submit(self, fn, /, *args, **kwargs): | ||
fut = Future() | ||
result = fn(*args, **kwargs) | ||
fut.set_result(result) | ||
|
||
return fut | ||
|
||
|
||
if __name__ == "__main__": | ||
use_dummy_executor = True # 是否使用上文实现的单线程执行器,否则使用dask的offload执行器 | ||
if use_dummy_executor: | ||
click.option( | ||
"--worker-class", | ||
type=str, | ||
default='startup.DummyWorker', | ||
show_default=True, | ||
help="Default executor.", | ||
)(main)() | ||
else: | ||
click.option( | ||
"--executor", | ||
type=str, | ||
default="offload", | ||
show_default=True, | ||
help="Default executor.", | ||
)(main)() |