1import datetime
2import os
3from importlib import import_module
4
5from omegaml.runtimes.proxies.jobproxy import OmegaJobProxy
6
7
8class DaskTask:
9 """
10 A dask remote function wrapper mimicking a Celery task
11 """
12
13 def __init__(self, fn, client, pure=True, **kwargs):
14 """
15 :param fn: (function) the function to be called
16 :param client: (dask client) the dask client to use
17 :param pure: (bool) whether this is a dask pure function (will
18 be cached or not). Defaults to True.
19 """
20 self.client = client
21 self.fn = fn
22 self.pure = pure
23 self.kwargs = kwargs
24
25 def delay(self, *args, **kwargs):
26 """
27 submit the function and execute on cluster.
28 """
29 kwargs['pure'] = kwargs.get('pure', self.pure)
30 if self.kwargs:
31 kwargs.update(self.kwargs)
32 return DaskAsyncResult(self.client.submit(self.fn, *args, **kwargs))
33
34
35class DaskAsyncResult:
36 """
37 A dask Future wrapper mimicking a Celery AsyncResult
38 """
39
40 def __init__(self, future):
41 self.future = future
42
43 def get(self):
44 import dask
45
46 if os.environ.get('DASK_DEBUG'):
47 with dask.set_options(get=dask.threaded.get):
48 return self.future.result()
49 return self.future.result()
50
51
52def daskhello(*args, **kwargs):
53 # test function for dask distributed
54 return f"hello from {os.getpid()} at {datetime.datetime.now()}"
55
56
[docs]
57class OmegaRuntimeDask:
58 """
59 omegaml compute cluster gateway to a dask distributed cluster
60
61 set environ DASK_DEBUG=1 to run dask tasks locally
62 """
63
64 def __init__(self, omega, dask_url=None):
65 self.dask_url = dask_url
66 self.omega = omega
67 self._client = None
68
69 @property
70 def client(self):
71 from distributed import Client, LocalCluster
72
73 if self._client is None:
74 if os.environ.get('DASK_DEBUG'):
75 # http://dask.pydata.org/en/latest/setup/single-distributed.html?highlight=single-threaded#localcluster
76 single_threaded = LocalCluster(processes=False)
77 self._client = Client(single_threaded)
78 else:
79 self._client = Client(self.dask_url)
80 return self._client
81
[docs]
82 def model(self, modelname):
83 """
84 return a model for remote execution
85 """
86 from omegaml.runtimes.proxies.modelproxy import OmegaModelProxy
87
88 return OmegaModelProxy(modelname, runtime=self)
89
90 def job(self, jobname):
91 """
92 return a job for remote exeuction
93 """
94 return OmegaJobProxy(jobname, runtime=self)
95
96 def task(self, name, **kwargs):
97 """
98 retrieve the task function from the task module
99
100 This retrieves the task function and wraps it into a
101 DaskTask. DaskTask mimicks a celery task and is
102 called on the cluster using .delay(), the same way we
103 call a celery task. .delay() will return a DaskAsyncResult,
104 supporting the celery .get() semantics. This way we can use
105 the same proxy objects, as all they do is call .delay() and
106 return an AsyncResult.
107 """
108 modname, funcname = name.rsplit('.', 1)
109 mod = import_module(modname)
110 func = getattr(mod, funcname)
111 # we pass pure=False to force dask to reevaluate the task
112 # http://distributed.readthedocs.io/en/latest/client.html?highlight=pure#pure-functions-by-default
113 return DaskTask(func, self.client, pure=False, **kwargs)
114
115 def settings(self):
116 """
117 return the runtimes's cluster settings
118 """
119 return self.task('omegaml.tasks.omega_settings').delay().get()
120
121 def ping(self):
122 return DaskTask(daskhello, self.client, pure=False)