-
Notifications
You must be signed in to change notification settings - Fork 29
Expand file tree
/
Copy pathfuture.py
More file actions
224 lines (179 loc) · 8.07 KB
/
Copy pathfuture.py
File metadata and controls
224 lines (179 loc) · 8.07 KB
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
216
217
218
219
220
221
222
223
224
import json
from ..utils import _ExponentialBackoff
class DSSFuture(object):
"""
A future represents a long-running task on a DSS instance. It allows you to
track the state of the task, retrieve its result when it is ready or abort it.
Usage example:
.. code-block:: python
# In this example, 'folder' is a DSSManagedFolder
future = folder.compute_metrics()
# Waits until the metrics are computed
metrics = future.wait_for_result()
.. note:: This class does not need to be instantiated directly. A :class:`dataikuapi.dss.future.DSSFuture` is usually returned by the API calls that are initiating long running-tasks.
"""
def __init__(self, client, job_id, state=None, result_wrapper=lambda result: result):
self.client = client
self.job_id = job_id
self.state = state
self.state_is_peek = True
self.result_wrapper = result_wrapper
@staticmethod
def from_resp(client, resp, result_wrapper=lambda result: result):
"""
Creates a :class:`dataikuapi.dss.future.DSSFuture` from the response of an endpoint that initiated a long-running task.
:param client: An api client to connect to the DSS backend
:type client: :class:`dataikuapi.dssclient.DSSClient`
:param resp: The response of the API call that initiated a long-running task.
:type resp: dict
:param result_wrapper: (optional) a function to apply to the result of the future, before it is returned by `get_result` or `wait_for_result` methods.
:type result_wrapper: callable
:return: :class:`dataikuapi.dss.future.DSSFuture`
"""
return DSSFuture(client, resp.get('jobId', None), state=resp, result_wrapper=result_wrapper)
@classmethod
def get_result_wait_if_needed(cls, client, ret):
"""
:meta private:
"""
if 'jobId' in ret:
future = DSSFuture(client, ret["jobId"], ret)
future.wait_for_result()
return future.get_result()
else:
return ret['result']
def abort(self):
"""
Aborts the long-running task.
"""
self.client._perform_empty("DELETE", "/futures/%s" % self.job_id)
def get_state(self):
"""
Queries the state of the future, and fetches the result if it's ready.
:return: the state of the future, including the result if it is ready.
:rtype: dict
"""
self.state = self.client._perform_json(
"GET", "/futures/%s" % self.job_id, params={'peek': False})
self.state_is_peek = False
return self.state
def peek_state(self):
"""
Queries the state of the future, without fetching the result.
:return: the state of the future
:rtype: dict
"""
self.state = self.client._perform_json(
"GET", "/futures/%s" % self.job_id, params={'peek': True})
self.state_is_peek = True
return self.state
def get_result(self):
"""
Returns the future's result if it's ready.
.. note:: if a custom result_wrapper was provided, it will be applied on the result that will be returned.
:return: the result of the future
:rtype: object
:raises Exception: if the result is not ready (i.e. the task is still running, or it has failed)
"""
if self.state is None or not self.state.get('hasResult', False) or self.state_is_peek:
self.get_state()
if self.state.get('hasResult', False):
return self.result_wrapper(self.state.get('result', None))
else:
raise Exception("Result not ready")
def has_result(self):
"""
Checks whether the task is completed successfully and has a result.
:return: True if the task has successfully returned a result, False otherwise
:rtype: bool
"""
if self.state is None or not self.state.get('hasResult', False):
self.get_state()
return self.state.get('hasResult', False)
def wait_for_result(self):
"""
Waits for the completion of the long-running task, and returns its result.
.. note:: if a custom result_wrapper was provided, it will be applied on the result that will be returned.
:return: the result of the future
:rtype: object
:raises Exception: if the task failed
"""
if self.state is not None and self.state.get('hasResult', False):
# no future created in backend, result already in the state
return self.result_wrapper(self.state.get('result', None))
if self.state is None or not self.state.get('hasResult', False) or self.state_is_peek:
self.get_state()
eb = _ExponentialBackoff()
while not self.state.get('hasResult', False):
eb.sleep_next()
self.get_state()
if self.state.get('hasResult', False):
return self.result_wrapper(self.state.get('result', None))
else:
raise Exception("No result")
class DSSFutureWithHistory(object):
"""
A future with history item represents a long-running task on a DSS instance that can be tracked by several clients.
It allows you to track the state of the task, retrieve its result when it is ready or abort it.
Usage example:
.. code-block:: python
# In this example, 'agent_review' is a DSSAgentReview
future_with_history = agent_review.perform_run()
# Waits until the computation of the run is finished
run_result = future.wait_for_result()
.. note:: This class does not need to be instantiated directly. A :class:`dataikuapi.dss.future.DSSFutureWithHistory` is returned by some API calls that are
initiating long running-tasks and allow access to the result to multiple requested
"""
def __init__(self, client, name, result_wrapper=lambda result: result):
self.client = client
self.name = name
self.state = None
self.result_wrapper = result_wrapper
def abort(self):
"""
Aborts the long-running task.
"""
self.client._perform_empty("DELETE", "/futures/history/%s" % self.name)
def get_state(self):
"""
Queries the state of the future with history.
:return: the state of the future with history.
:rtype: dict
"""
raw_state = self.client._perform_text("GET", "/futures/history/%s" % self.name)
self.state = None if raw_state == '' else json.loads(raw_state)
return self.state
def has_result(self):
"""
Checks whether the task is completed successfully and has a result.
:return: True if the task has successfully returned a result, False otherwise
:rtype: bool
"""
if self.state is None or not (self.state.get("lastResponse") or {}).get("hasResult", False):
self.get_state()
if self.state is None: # Not found
return False
last_resp = self.state.get("lastResponse") or {}
return last_resp.get("hasResult", False)
def wait_for_result(self):
"""
Waits for the completion of the long-running task, and returns its result.
.. note:: if a custom result_wrapper was provided, it will be applied on the result that will be returned.
:return: the result of the future
:rtype: object
:raises Exception: if the task failed or was not found
"""
eb = _ExponentialBackoff()
refresh_needed = self.state is None
while True:
if refresh_needed:
self.get_state()
if self.state is None:
raise Exception("Future history item not found: %s" % self.name)
if self.state.get("storedError"):
raise Exception("Job failed: %s" % self.state.get("storedError"))
last_resp = self.state.get("lastResponse")
if last_resp and last_resp.get("hasResult"):
return self.result_wrapper(last_resp.get("result"))
eb.sleep_next()
refresh_needed = True