pintel 0.3.0.dev0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- pintel/__init__.py +7 -0
- pintel/hub.py +317 -0
- pintel/worker.py +231 -0
- pintel-0.3.0.dev0.dist-info/METADATA +230 -0
- pintel-0.3.0.dev0.dist-info/RECORD +8 -0
- pintel-0.3.0.dev0.dist-info/WHEEL +5 -0
- pintel-0.3.0.dev0.dist-info/licenses/LICENSE +21 -0
- pintel-0.3.0.dev0.dist-info/top_level.txt +1 -0
pintel/__init__.py
ADDED
pintel/hub.py
ADDED
|
@@ -0,0 +1,317 @@
|
|
|
1
|
+
import asyncio
|
|
2
|
+
import time
|
|
3
|
+
import uuid
|
|
4
|
+
from collections import namedtuple
|
|
5
|
+
from typing import Any, Tuple
|
|
6
|
+
|
|
7
|
+
import msgpack
|
|
8
|
+
import zmq
|
|
9
|
+
import zmq.asyncio
|
|
10
|
+
from loguru import logger
|
|
11
|
+
|
|
12
|
+
from pintel.worker import CODE_FAILURE, DEFAULT_HEARTBEAT_INTERVAL, DEFAULT_PORT, PONG
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
RESPONSE_BUFFER_SIZE: int = 4096
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
class UserCodeException(Exception):
|
|
19
|
+
"""Custom Exception for user-defined catched errors.
|
|
20
|
+
|
|
21
|
+
Args:
|
|
22
|
+
t (str): Traceback returned by the worker.
|
|
23
|
+
"""
|
|
24
|
+
|
|
25
|
+
def __init__(self, t: str):
|
|
26
|
+
super().__init__(f"Exception raised in user-defined code. Traceback :\n{t}")
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
class WorkerManager:
|
|
30
|
+
"""A helper class that takes care of managing workers.
|
|
31
|
+
Workers' address can be registered as available, and this class will make
|
|
32
|
+
sure to return address of workers that are available and alive.
|
|
33
|
+
|
|
34
|
+
A worker is considered as dead if we didn't receive any heartbeat within a
|
|
35
|
+
given interval.
|
|
36
|
+
|
|
37
|
+
Args:
|
|
38
|
+
heartbeat_interval (float): Interval of time (in seconds) after which we
|
|
39
|
+
consider a worker to be dead.
|
|
40
|
+
"""
|
|
41
|
+
|
|
42
|
+
def __init__(self, heartbeat_interval: float):
|
|
43
|
+
super().__init__()
|
|
44
|
+
|
|
45
|
+
self.heartbeat_t = heartbeat_interval
|
|
46
|
+
|
|
47
|
+
self.w_ts = {}
|
|
48
|
+
self.w_access = asyncio.Condition()
|
|
49
|
+
|
|
50
|
+
async def reckon(self, address: str):
|
|
51
|
+
"""Register the given address as available.
|
|
52
|
+
|
|
53
|
+
Args:
|
|
54
|
+
address (str): Address of the worker to register as available.
|
|
55
|
+
"""
|
|
56
|
+
async with self.w_access:
|
|
57
|
+
self.w_ts[address] = time.time()
|
|
58
|
+
self.w_access.notify()
|
|
59
|
+
|
|
60
|
+
async def get_next_worker(self) -> str:
|
|
61
|
+
"""Retrieve the next available and alive worker's address.
|
|
62
|
+
|
|
63
|
+
Returns:
|
|
64
|
+
str: Address of the available and alive worker.
|
|
65
|
+
"""
|
|
66
|
+
async with self.w_access:
|
|
67
|
+
# Iterate workers until we find one that was alive recently
|
|
68
|
+
w_alive = False
|
|
69
|
+
while not w_alive:
|
|
70
|
+
# If no workers are available, wait...
|
|
71
|
+
if not self.w_ts:
|
|
72
|
+
await self.w_access.wait()
|
|
73
|
+
|
|
74
|
+
address, ts = self.w_ts.popitem()
|
|
75
|
+
w_alive = time.time() - ts < self.heartbeat_t
|
|
76
|
+
|
|
77
|
+
return address
|
|
78
|
+
|
|
79
|
+
|
|
80
|
+
Response = namedtuple("Response", ["code", "content"])
|
|
81
|
+
|
|
82
|
+
|
|
83
|
+
class RequestManager:
|
|
84
|
+
"""A helper class that takes care of storing responses and waiting for the
|
|
85
|
+
right response.
|
|
86
|
+
|
|
87
|
+
Args:
|
|
88
|
+
resp_buffer_size (int): Maximum size of the response buffer.
|
|
89
|
+
"""
|
|
90
|
+
|
|
91
|
+
def __init__(self, resp_buffer_size: int):
|
|
92
|
+
super().__init__()
|
|
93
|
+
|
|
94
|
+
self.resp_buffer_size = resp_buffer_size
|
|
95
|
+
|
|
96
|
+
self.responses = {}
|
|
97
|
+
self.req_states = {}
|
|
98
|
+
|
|
99
|
+
def pin(self, req_id: str):
|
|
100
|
+
"""Pin a request ID. This is a necessary step when sending a request,
|
|
101
|
+
so that the request can be awaited until a response is received.
|
|
102
|
+
This method should be called for each `req_id` before calling `wait_for`.
|
|
103
|
+
|
|
104
|
+
Args:
|
|
105
|
+
req_id (str): Request unique identifier.
|
|
106
|
+
"""
|
|
107
|
+
self.req_states[req_id] = asyncio.Event()
|
|
108
|
+
|
|
109
|
+
# Ensure we don't store too many requests
|
|
110
|
+
if len(self.req_states) > self.resp_buffer_size:
|
|
111
|
+
# If it's the case, forget the oldest one
|
|
112
|
+
k = list(self.req_states.keys())[0]
|
|
113
|
+
logger.warning(f"Response buffer overflow (>{self.resp_buffer_size}). Forgetting oldest request : {k}")
|
|
114
|
+
self.req_states.pop(k)
|
|
115
|
+
self.responses.pop(k, None)
|
|
116
|
+
|
|
117
|
+
async def wait_for(self, req_id: str) -> Tuple[int, Any]:
|
|
118
|
+
"""Async method that waits until we received the response corresponding
|
|
119
|
+
to the given request ID.
|
|
120
|
+
|
|
121
|
+
The method `pin` should be called before waiting with this method.
|
|
122
|
+
|
|
123
|
+
Args:
|
|
124
|
+
req_id (str): Request unique identifier.
|
|
125
|
+
|
|
126
|
+
Raises:
|
|
127
|
+
KeyError: Exception raised if the request wasn't registered previously.
|
|
128
|
+
|
|
129
|
+
Returns:
|
|
130
|
+
Tuple[int, Any]: Code and content of the received response.
|
|
131
|
+
"""
|
|
132
|
+
if req_id not in self.req_states:
|
|
133
|
+
raise KeyError(f"Request #{req_id} was not pinned, or was removed because of buffer overflow")
|
|
134
|
+
|
|
135
|
+
# Wait for the receiving loop to receive the response
|
|
136
|
+
await self.req_states[req_id].wait()
|
|
137
|
+
|
|
138
|
+
# Once we get it, access the result
|
|
139
|
+
r = self.responses.pop(req_id)
|
|
140
|
+
|
|
141
|
+
# Don't forget to remove the event
|
|
142
|
+
self.req_states.pop(req_id)
|
|
143
|
+
|
|
144
|
+
return r.code, r.content
|
|
145
|
+
|
|
146
|
+
def store(self, req_id: str, code: int, response: Any):
|
|
147
|
+
"""Store a response, to be consumed later.
|
|
148
|
+
|
|
149
|
+
Args:
|
|
150
|
+
req_id (str): Request unique identifier.
|
|
151
|
+
code (int): Code of the response.
|
|
152
|
+
response (Any): Content of the response.
|
|
153
|
+
"""
|
|
154
|
+
# Store the response if the req_id is recognized
|
|
155
|
+
if req_id in self.req_states:
|
|
156
|
+
self.responses[req_id] = Response(code, response)
|
|
157
|
+
|
|
158
|
+
# Notify that we received the response
|
|
159
|
+
self.req_states[req_id].set()
|
|
160
|
+
else:
|
|
161
|
+
logger.warning(
|
|
162
|
+
f"Request #{req_id} was previously removed from response buffer. "
|
|
163
|
+
f"Ignoring the response from this request..."
|
|
164
|
+
)
|
|
165
|
+
|
|
166
|
+
|
|
167
|
+
class Hub:
|
|
168
|
+
"""Class acting as a hub for all the requests to send. All requests sent
|
|
169
|
+
through this class will be automatically dispatched to connected workers,
|
|
170
|
+
will wait for the responses and return it.
|
|
171
|
+
|
|
172
|
+
Args:
|
|
173
|
+
port (int): Port number to use for the sockets. Defaults to
|
|
174
|
+
DEFAULT_PORT.
|
|
175
|
+
heartbeat_interval (float): Heartbeat interval used by the
|
|
176
|
+
workers, in seconds. Defaults to DEFAULT_HEARTBEAT_INTERVAL.
|
|
177
|
+
resp_buffer_size (int): Maximum response buffer size. Defaults
|
|
178
|
+
to RESPONSE_BUFFER_SIZE.
|
|
179
|
+
"""
|
|
180
|
+
|
|
181
|
+
def __init__(
|
|
182
|
+
self,
|
|
183
|
+
port: int = DEFAULT_PORT,
|
|
184
|
+
heartbeat_interval: float = DEFAULT_HEARTBEAT_INTERVAL,
|
|
185
|
+
resp_buffer_size: int = RESPONSE_BUFFER_SIZE,
|
|
186
|
+
):
|
|
187
|
+
super().__init__()
|
|
188
|
+
|
|
189
|
+
self.port = port
|
|
190
|
+
self.heartbeat_t = heartbeat_interval
|
|
191
|
+
self.socket = None
|
|
192
|
+
self.w_manager = None
|
|
193
|
+
self.req_manager = RequestManager(resp_buffer_size=resp_buffer_size)
|
|
194
|
+
|
|
195
|
+
async def receive_loop(self):
|
|
196
|
+
"""Infinite loop for receiving responses from the workers."""
|
|
197
|
+
while True:
|
|
198
|
+
# Receive stuff
|
|
199
|
+
logger.debug("Receiving...")
|
|
200
|
+
address, *frames = await self.socket.recv_multipart()
|
|
201
|
+
logger.debug(f"Received something from worker #{address}")
|
|
202
|
+
|
|
203
|
+
# Since we received a response from this worker, it means it's ready for more !
|
|
204
|
+
await self.w_manager.reckon(address)
|
|
205
|
+
|
|
206
|
+
if len(frames) == 1:
|
|
207
|
+
# Answer the Ping
|
|
208
|
+
logger.debug("Answering the ping")
|
|
209
|
+
await self.socket.send_multipart([address, b"", PONG])
|
|
210
|
+
continue
|
|
211
|
+
|
|
212
|
+
_, resp = frames
|
|
213
|
+
req_id, code, res = msgpack.unpackb(resp)
|
|
214
|
+
logger.debug(f"Received response from request #{req_id}")
|
|
215
|
+
|
|
216
|
+
self.req_manager.store(req_id, code, res)
|
|
217
|
+
|
|
218
|
+
def _start_if_not_started(self):
|
|
219
|
+
"""Helper method to ensure everything is properly started (socket is
|
|
220
|
+
initialized, receiving loop is started, etc...).
|
|
221
|
+
"""
|
|
222
|
+
if self.socket is None:
|
|
223
|
+
# Create what we need here in this process/context
|
|
224
|
+
context = zmq.asyncio.Context()
|
|
225
|
+
self.socket = context.socket(zmq.ROUTER)
|
|
226
|
+
self.socket.bind(f"tcp://*:{self.port}")
|
|
227
|
+
self.w_manager = WorkerManager(heartbeat_interval=self.heartbeat_t)
|
|
228
|
+
|
|
229
|
+
# Fire and forget : infinite loop, taking care of receiving stuff from the socket
|
|
230
|
+
logger.info("Starting receiving loop...")
|
|
231
|
+
asyncio.create_task(self.receive_loop())
|
|
232
|
+
|
|
233
|
+
async def _request(self, *args: Any, **kwargs: Any) -> Any:
|
|
234
|
+
"""Raw method to send a request. This method will do the following :
|
|
235
|
+
* Start the receiving loop if it was not started
|
|
236
|
+
* Get the address of the next worker ready (blocking)
|
|
237
|
+
* Send the request to the worker
|
|
238
|
+
* Wait for the response (blocking)
|
|
239
|
+
* Return the result
|
|
240
|
+
|
|
241
|
+
Args:
|
|
242
|
+
*args: Positional arguments for the request.
|
|
243
|
+
**kwargs: Keywords arguments for the request.
|
|
244
|
+
|
|
245
|
+
Raises:
|
|
246
|
+
UserCodeException: Exception raised if an exception was raised inside
|
|
247
|
+
user-defined code on the worker side.
|
|
248
|
+
|
|
249
|
+
Returns:
|
|
250
|
+
Any: The response for the request.
|
|
251
|
+
"""
|
|
252
|
+
# Before anything, if the receiving loop was not started, start it
|
|
253
|
+
self._start_if_not_started()
|
|
254
|
+
|
|
255
|
+
# Assign a unique ID to the request
|
|
256
|
+
req_id = uuid.uuid4().hex
|
|
257
|
+
|
|
258
|
+
# Let the manager know that we are waiting for this request
|
|
259
|
+
logger.debug(f"Pinning request #{req_id}")
|
|
260
|
+
self.req_manager.pin(req_id)
|
|
261
|
+
|
|
262
|
+
# Send the request
|
|
263
|
+
address = await self.w_manager.get_next_worker()
|
|
264
|
+
logger.debug(f"Sending request #{req_id} to worker #{address}")
|
|
265
|
+
await self.socket.send_multipart([address, b"", msgpack.packb([req_id, args, kwargs])])
|
|
266
|
+
|
|
267
|
+
# Wait until we receive the response
|
|
268
|
+
code, res = await self.req_manager.wait_for(req_id)
|
|
269
|
+
logger.debug(f"Received result for request #{req_id}")
|
|
270
|
+
|
|
271
|
+
# Depending on what is the response, deal with it properly
|
|
272
|
+
if code == CODE_FAILURE:
|
|
273
|
+
raise UserCodeException(res)
|
|
274
|
+
else:
|
|
275
|
+
return res
|
|
276
|
+
|
|
277
|
+
async def request(self, *args: Any, pintel_timeout: float = None, pintel_retries: int = 0, **kwargs: Any) -> Any:
|
|
278
|
+
"""Main method, used to send a request to workers and get the result.
|
|
279
|
+
|
|
280
|
+
This method is a wrapper around `_request`, providing additional
|
|
281
|
+
functionalities :
|
|
282
|
+
* Timeout
|
|
283
|
+
* Automatic retries
|
|
284
|
+
|
|
285
|
+
Args:
|
|
286
|
+
*args: Positional arguments for the request.
|
|
287
|
+
pintel_timeout (float): Timeout for the request, in seconds.
|
|
288
|
+
If `None` is given, block until the request is complete.
|
|
289
|
+
Defaults to None.
|
|
290
|
+
pintel_retries (int): Number of retries. This argument is
|
|
291
|
+
used only if `pintel_timeout` is not `None`. If `-1` is given,
|
|
292
|
+
indefinitely retry. Defaults to 0.
|
|
293
|
+
**kwargs: Keywords arguments for the request.
|
|
294
|
+
|
|
295
|
+
Raises:
|
|
296
|
+
asyncio.TimeoutError: Error raised if no response is received within
|
|
297
|
+
the given timeout.
|
|
298
|
+
|
|
299
|
+
Returns:
|
|
300
|
+
Any: The response for the request.
|
|
301
|
+
"""
|
|
302
|
+
try:
|
|
303
|
+
return await asyncio.wait_for(self._request(*args, **kwargs), timeout=pintel_timeout)
|
|
304
|
+
except asyncio.TimeoutError:
|
|
305
|
+
if pintel_retries == 0:
|
|
306
|
+
logger.error(f"Request failed : no response received within {pintel_timeout}s")
|
|
307
|
+
raise
|
|
308
|
+
else:
|
|
309
|
+
retries_left = max(pintel_retries - 1, -1)
|
|
310
|
+
logger.warning(
|
|
311
|
+
f"Request failed : no response received within {pintel_timeout}s. Retrying {retries_left} times"
|
|
312
|
+
)
|
|
313
|
+
return await self.request(*args, pintel_timeout=pintel_timeout, pintel_retries=retries_left, **kwargs)
|
|
314
|
+
|
|
315
|
+
def __del__(self):
|
|
316
|
+
if self.socket is not None:
|
|
317
|
+
self.socket.close()
|
pintel/worker.py
ADDED
|
@@ -0,0 +1,231 @@
|
|
|
1
|
+
import signal
|
|
2
|
+
import traceback
|
|
3
|
+
import uuid
|
|
4
|
+
from multiprocessing import Process
|
|
5
|
+
from typing import Any, Callable
|
|
6
|
+
|
|
7
|
+
import msgpack
|
|
8
|
+
import zmq
|
|
9
|
+
from loguru import logger
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
DEFAULT_PORT: int = 5019
|
|
13
|
+
DEFAULT_HEARTBEAT_INTERVAL: float = 1
|
|
14
|
+
DEFAULT_RESET_AFTER_N_MISS: int = 2
|
|
15
|
+
MS: int = 1000
|
|
16
|
+
CODE_SUCCESS: int = 0
|
|
17
|
+
CODE_FAILURE: int = 1
|
|
18
|
+
PING: bytes = b""
|
|
19
|
+
PONG: bytes = b""
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
class Worker(Process):
|
|
23
|
+
"""Define a worker process. This worker process indefinitely waits for
|
|
24
|
+
requests on a socket. Upon receiving a request, it processes it with the
|
|
25
|
+
worker class provided, and return the response.
|
|
26
|
+
|
|
27
|
+
After creating the Worker object, you have 2 different ways to run it :
|
|
28
|
+
* `worker.run()` : It will run in the current process directly (blocking,
|
|
29
|
+
infinite loop).
|
|
30
|
+
* `worker.start()` : It will start a different process and start the code
|
|
31
|
+
there (non-blocking).
|
|
32
|
+
|
|
33
|
+
Args:
|
|
34
|
+
worker_cls (Callable): Worker class containing the code that will be used
|
|
35
|
+
to process requests.
|
|
36
|
+
pintel_host (str): Host of the Hub. Defaults to "localhost".
|
|
37
|
+
pintel_port (int): Port of the Hub. Defaults to DEFAULT_PORT.
|
|
38
|
+
pintel_heartbeat_interval (float): Heartbeat interval between the
|
|
39
|
+
worker and the Hub. Defaults to DEFAULT_HEARTBEAT_INTERVAL.
|
|
40
|
+
pintel_reset_after_n_miss (int): Number of missed heartbeats
|
|
41
|
+
allowed before hard-resetting the socket and retrying. Defaults to
|
|
42
|
+
DEFAULT_RESET_AFTER_N_MISS.
|
|
43
|
+
"""
|
|
44
|
+
|
|
45
|
+
def __init__(
|
|
46
|
+
self,
|
|
47
|
+
worker_cls: Callable,
|
|
48
|
+
*args: Any,
|
|
49
|
+
pintel_host: str = "localhost",
|
|
50
|
+
pintel_port: int = DEFAULT_PORT,
|
|
51
|
+
pintel_heartbeat_interval: float = DEFAULT_HEARTBEAT_INTERVAL,
|
|
52
|
+
pintel_reset_after_n_miss: int = DEFAULT_RESET_AFTER_N_MISS,
|
|
53
|
+
**kwargs: Any,
|
|
54
|
+
):
|
|
55
|
+
super().__init__()
|
|
56
|
+
|
|
57
|
+
self.worker_cls = worker_cls
|
|
58
|
+
self.worker_args = args
|
|
59
|
+
self.worker_kwargs = kwargs
|
|
60
|
+
|
|
61
|
+
self.identity = uuid.uuid4().hex
|
|
62
|
+
self.host = pintel_host
|
|
63
|
+
self.port = pintel_port
|
|
64
|
+
self.heartbeat_t = pintel_heartbeat_interval
|
|
65
|
+
self.reset_n_miss = pintel_reset_after_n_miss
|
|
66
|
+
|
|
67
|
+
self.waiting_pong = 0
|
|
68
|
+
|
|
69
|
+
def create_socket(self, context: zmq.Context) -> zmq.Socket:
|
|
70
|
+
"""Helper method to create a socket, setting its identity and connecting
|
|
71
|
+
to the Hub.
|
|
72
|
+
|
|
73
|
+
Args:
|
|
74
|
+
context (zmq.Context): ZMQ context to use.
|
|
75
|
+
|
|
76
|
+
Returns:
|
|
77
|
+
zmq.Socket: Initialized and connected socket, ready to use.
|
|
78
|
+
"""
|
|
79
|
+
# Create the socket, set its identity
|
|
80
|
+
socket = context.socket(zmq.DEALER)
|
|
81
|
+
socket.setsockopt_string(zmq.IDENTITY, self.identity)
|
|
82
|
+
|
|
83
|
+
# Connect to the Hub
|
|
84
|
+
socket.connect(f"tcp://{self.host}:{self.port}")
|
|
85
|
+
|
|
86
|
+
return socket
|
|
87
|
+
|
|
88
|
+
def create_term_socket(self, context: zmq.Context) -> zmq.Socket:
|
|
89
|
+
"""Helper method to create a termination socket.
|
|
90
|
+
|
|
91
|
+
Basically it creates 2 sockets, bind/connect them together. One socket
|
|
92
|
+
will be used to send a termination signal, and the other is returned and
|
|
93
|
+
used to receive the termination signal.
|
|
94
|
+
|
|
95
|
+
Args:
|
|
96
|
+
context (zmq.Context): ZMQ context to use.
|
|
97
|
+
|
|
98
|
+
Returns:
|
|
99
|
+
zmq.Socket: Initialized and connected socket, ready to use.
|
|
100
|
+
"""
|
|
101
|
+
# Create the socket than will send the termination ping
|
|
102
|
+
term_snd_socket = context.socket(zmq.REQ)
|
|
103
|
+
port = term_snd_socket.bind_to_random_port("tcp://127.0.0.1")
|
|
104
|
+
|
|
105
|
+
# Then define the behavior on how to send the termination ping
|
|
106
|
+
def send_term(*args, **kwargs):
|
|
107
|
+
logger.debug("Sending termination ping...")
|
|
108
|
+
# Send something on the termination socket, it doesn't matter what
|
|
109
|
+
term_snd_socket.send(PING)
|
|
110
|
+
|
|
111
|
+
# We send the termination ping upon receiving these signals
|
|
112
|
+
signal.signal(signal.SIGTERM, send_term)
|
|
113
|
+
signal.signal(signal.SIGINT, send_term)
|
|
114
|
+
|
|
115
|
+
# And finally create the socket that we will use to receive the termination signal
|
|
116
|
+
term_rcv_socket = context.socket(zmq.REP)
|
|
117
|
+
term_rcv_socket.connect(f"tcp://localhost:{port}")
|
|
118
|
+
|
|
119
|
+
return term_rcv_socket
|
|
120
|
+
|
|
121
|
+
def reset_socket(self, socket: zmq.Socket, context: zmq.Context, poller: zmq.Poller) -> zmq.Socket:
|
|
122
|
+
"""Helper method to reset the given socket.
|
|
123
|
+
|
|
124
|
+
This method unregister the socket from the given poller, close the
|
|
125
|
+
socket, and then recreate the socket and register this new socket in the
|
|
126
|
+
poller.
|
|
127
|
+
|
|
128
|
+
Args:
|
|
129
|
+
socket (zmq.Socket): ZMQ socket to reset.
|
|
130
|
+
context (zmq.Context): ZMQ context to use.
|
|
131
|
+
poller (zmq.Poller): ZMQ poller where the socket is registered.
|
|
132
|
+
|
|
133
|
+
Returns:
|
|
134
|
+
zmq.Socket: Initialized and connected socket, ready to use.
|
|
135
|
+
"""
|
|
136
|
+
# Close the existing socket
|
|
137
|
+
poller.unregister(socket)
|
|
138
|
+
socket.close(linger=0)
|
|
139
|
+
|
|
140
|
+
# Recreate the socket
|
|
141
|
+
socket = self.create_socket(context)
|
|
142
|
+
poller.register(socket, zmq.POLLIN)
|
|
143
|
+
|
|
144
|
+
return socket
|
|
145
|
+
|
|
146
|
+
def ping(self, socket: zmq.Socket):
|
|
147
|
+
"""Helper method used for the heartbeat. Also takes care of keeping the
|
|
148
|
+
counter of heartbeats up-to-date.
|
|
149
|
+
|
|
150
|
+
Args:
|
|
151
|
+
socket (zmq.Socket): Socket to use to send the heartbeat.
|
|
152
|
+
"""
|
|
153
|
+
logger.debug("Sending ping...")
|
|
154
|
+
socket.send(PING)
|
|
155
|
+
self.waiting_pong += 1
|
|
156
|
+
|
|
157
|
+
def run(self):
|
|
158
|
+
"""Main method. It will initialize the worker class, and enter an
|
|
159
|
+
infinite loop, waiting for requests. Whenever a request is received, it
|
|
160
|
+
processes it with the code provided in the constructor.
|
|
161
|
+
"""
|
|
162
|
+
# Instanciate the worker
|
|
163
|
+
worker = self.worker_cls(*self.worker_args, **self.worker_kwargs)
|
|
164
|
+
|
|
165
|
+
# Initialize what we need for handling sockets
|
|
166
|
+
context = zmq.Context()
|
|
167
|
+
poller = zmq.Poller()
|
|
168
|
+
|
|
169
|
+
# Create the socket for termination
|
|
170
|
+
term_socket = self.create_term_socket(context)
|
|
171
|
+
poller.register(term_socket, zmq.POLLIN)
|
|
172
|
+
|
|
173
|
+
# Create the socket connecting to the hub
|
|
174
|
+
socket = self.create_socket(context)
|
|
175
|
+
poller.register(socket, zmq.POLLIN)
|
|
176
|
+
|
|
177
|
+
logger.info("Worker ready to roll")
|
|
178
|
+
|
|
179
|
+
# Tell the Hub we are ready
|
|
180
|
+
self.ping(socket)
|
|
181
|
+
|
|
182
|
+
# Indefinitely wait for requests : when we are done with one request,
|
|
183
|
+
# we wait for the next one
|
|
184
|
+
while True:
|
|
185
|
+
logger.debug("Waiting for request...")
|
|
186
|
+
events = dict(poller.poll(self.heartbeat_t * MS))
|
|
187
|
+
|
|
188
|
+
if term_socket in events:
|
|
189
|
+
logger.debug("Termination signal received, shutting down gracefully")
|
|
190
|
+
break
|
|
191
|
+
|
|
192
|
+
if socket in events:
|
|
193
|
+
_, workload = socket.recv_multipart(zmq.NOBLOCK)
|
|
194
|
+
logger.debug("Received something !")
|
|
195
|
+
self.waiting_pong = 0
|
|
196
|
+
else:
|
|
197
|
+
logger.debug(f"Didn't receive anything for {self.heartbeat_t}s ({self.waiting_pong})")
|
|
198
|
+
|
|
199
|
+
if self.waiting_pong >= self.reset_n_miss:
|
|
200
|
+
logger.warning(
|
|
201
|
+
f"The Hub is not answering, even after {self.waiting_pong} missed pings... "
|
|
202
|
+
f"Resetting the socket"
|
|
203
|
+
)
|
|
204
|
+
socket = self.reset_socket(socket=socket, context=context, poller=poller)
|
|
205
|
+
self.waiting_pong = 0
|
|
206
|
+
|
|
207
|
+
# We didn't receive anything for some time, try to ping again
|
|
208
|
+
self.ping(socket)
|
|
209
|
+
continue
|
|
210
|
+
|
|
211
|
+
if workload == PONG:
|
|
212
|
+
# Just a Pong, ignore it
|
|
213
|
+
logger.debug("It was just a pong...")
|
|
214
|
+
continue
|
|
215
|
+
|
|
216
|
+
# From here the Hub sent us an actual request
|
|
217
|
+
req_id, req_args, req_kwargs = msgpack.unpackb(workload)
|
|
218
|
+
logger.debug(f"Request #{req_id} received")
|
|
219
|
+
|
|
220
|
+
# Call worker's code with the request arguments
|
|
221
|
+
try:
|
|
222
|
+
res = worker(*req_args, **req_kwargs)
|
|
223
|
+
except Exception as e:
|
|
224
|
+
logger.warning(f"Exception in user-defined __call__ method : {e.__class__.__name__}({str(e)})")
|
|
225
|
+
socket.send_multipart([b"", msgpack.packb([req_id, CODE_FAILURE, traceback.format_exc()])])
|
|
226
|
+
else:
|
|
227
|
+
logger.debug("Sending back the response")
|
|
228
|
+
socket.send_multipart([b"", msgpack.packb([req_id, CODE_SUCCESS, res])])
|
|
229
|
+
|
|
230
|
+
logger.info("Worker is shut down")
|
|
231
|
+
quit()
|
|
@@ -0,0 +1,230 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: pintel
|
|
3
|
+
Version: 0.3.0.dev0
|
|
4
|
+
Summary: Scale your ML workers asynchronously across processes and machines
|
|
5
|
+
Home-page: https://github.com/astariul/pintel
|
|
6
|
+
Author: Nicolas REMOND
|
|
7
|
+
Author-email: remondnicola@gmail.com
|
|
8
|
+
Classifier: Programming Language :: Python :: 3.8
|
|
9
|
+
Classifier: Operating System :: OS Independent
|
|
10
|
+
Requires-Python: >=3.8
|
|
11
|
+
Description-Content-Type: text/markdown
|
|
12
|
+
License-File: LICENSE
|
|
13
|
+
Requires-Dist: pyzmq<27.0,>=22.3
|
|
14
|
+
Requires-Dist: msgpack~=1.0
|
|
15
|
+
Requires-Dist: loguru~=0.6
|
|
16
|
+
Provides-Extra: test
|
|
17
|
+
Requires-Dist: pytest~=8.0; extra == "test"
|
|
18
|
+
Requires-Dist: pytest-asyncio~=0.18; extra == "test"
|
|
19
|
+
Requires-Dist: pytest-cov~=3.0; extra == "test"
|
|
20
|
+
Requires-Dist: coverage-badge~=1.0; extra == "test"
|
|
21
|
+
Provides-Extra: hook
|
|
22
|
+
Requires-Dist: pre-commit~=4.0; extra == "hook"
|
|
23
|
+
Provides-Extra: lint
|
|
24
|
+
Requires-Dist: isort~=5.9; extra == "lint"
|
|
25
|
+
Requires-Dist: black~=24.1; extra == "lint"
|
|
26
|
+
Requires-Dist: flake518~=1.2; extra == "lint"
|
|
27
|
+
Requires-Dist: darglint~=1.8; extra == "lint"
|
|
28
|
+
Provides-Extra: docs
|
|
29
|
+
Requires-Dist: mkdocs-material~=9.0; extra == "docs"
|
|
30
|
+
Requires-Dist: mkdocstrings[python]~=0.18; extra == "docs"
|
|
31
|
+
Requires-Dist: mike~=2.0; extra == "docs"
|
|
32
|
+
Provides-Extra: ex
|
|
33
|
+
Requires-Dist: fastapi~=0.75; extra == "ex"
|
|
34
|
+
Requires-Dist: uvicorn~=0.17; extra == "ex"
|
|
35
|
+
Requires-Dist: requests~=2.27; extra == "ex"
|
|
36
|
+
Requires-Dist: transformers~=4.17; extra == "ex"
|
|
37
|
+
Provides-Extra: all
|
|
38
|
+
Requires-Dist: pytest~=8.0; extra == "all"
|
|
39
|
+
Requires-Dist: pytest-asyncio~=0.18; extra == "all"
|
|
40
|
+
Requires-Dist: pytest-cov~=3.0; extra == "all"
|
|
41
|
+
Requires-Dist: coverage-badge~=1.0; extra == "all"
|
|
42
|
+
Requires-Dist: pre-commit~=4.0; extra == "all"
|
|
43
|
+
Requires-Dist: isort~=5.9; extra == "all"
|
|
44
|
+
Requires-Dist: black~=24.1; extra == "all"
|
|
45
|
+
Requires-Dist: flake518~=1.2; extra == "all"
|
|
46
|
+
Requires-Dist: darglint~=1.8; extra == "all"
|
|
47
|
+
Requires-Dist: mkdocs-material~=9.0; extra == "all"
|
|
48
|
+
Requires-Dist: mkdocstrings[python]~=0.18; extra == "all"
|
|
49
|
+
Requires-Dist: mike~=2.0; extra == "all"
|
|
50
|
+
Requires-Dist: fastapi~=0.75; extra == "all"
|
|
51
|
+
Requires-Dist: uvicorn~=0.17; extra == "all"
|
|
52
|
+
Requires-Dist: requests~=2.27; extra == "all"
|
|
53
|
+
Requires-Dist: transformers~=4.17; extra == "all"
|
|
54
|
+
Provides-Extra: dev
|
|
55
|
+
Requires-Dist: pytest~=8.0; extra == "dev"
|
|
56
|
+
Requires-Dist: pytest-asyncio~=0.18; extra == "dev"
|
|
57
|
+
Requires-Dist: pytest-cov~=3.0; extra == "dev"
|
|
58
|
+
Requires-Dist: coverage-badge~=1.0; extra == "dev"
|
|
59
|
+
Requires-Dist: pre-commit~=4.0; extra == "dev"
|
|
60
|
+
Requires-Dist: isort~=5.9; extra == "dev"
|
|
61
|
+
Requires-Dist: black~=24.1; extra == "dev"
|
|
62
|
+
Requires-Dist: flake518~=1.2; extra == "dev"
|
|
63
|
+
Requires-Dist: darglint~=1.8; extra == "dev"
|
|
64
|
+
Requires-Dist: mkdocs-material~=9.0; extra == "dev"
|
|
65
|
+
Requires-Dist: mkdocstrings[python]~=0.18; extra == "dev"
|
|
66
|
+
Requires-Dist: mike~=2.0; extra == "dev"
|
|
67
|
+
Dynamic: author
|
|
68
|
+
Dynamic: author-email
|
|
69
|
+
Dynamic: classifier
|
|
70
|
+
Dynamic: description
|
|
71
|
+
Dynamic: description-content-type
|
|
72
|
+
Dynamic: home-page
|
|
73
|
+
Dynamic: license-file
|
|
74
|
+
Dynamic: provides-extra
|
|
75
|
+
Dynamic: requires-dist
|
|
76
|
+
Dynamic: requires-python
|
|
77
|
+
Dynamic: summary
|
|
78
|
+
|
|
79
|
+
<h1 align="center">pintel</h1>
|
|
80
|
+
<p align="center">
|
|
81
|
+
Scale your ML workers asynchronously across processes and machines
|
|
82
|
+
</p>
|
|
83
|
+
|
|
84
|
+
<p align="center">
|
|
85
|
+
<a href="https://github.com/astariul/pintel/releases"><img src="https://img.shields.io/github/release/astariul/pintel.svg" alt="GitHub release" /></a>
|
|
86
|
+
<a href="https://github.com/astariul/pintel/actions/workflows/pytest.yaml"><img src="https://github.com/astariul/pintel/actions/workflows/pytest.yaml/badge.svg" alt="Test status" /></a>
|
|
87
|
+
<a href="https://github.com/astariul/pintel/actions/workflows/lint.yaml"><img src="https://github.com/astariul/pintel/actions/workflows/lint.yaml/badge.svg" alt="Lint status" /></a>
|
|
88
|
+
<img src=".github/badges/coverage.svg" alt="Coverage status" />
|
|
89
|
+
<a href="https://astariul.github.io/pintel"><img src="https://img.shields.io/website?down_message=failing&label=docs&up_color=green&up_message=passing&url=https%3A%2F%2Fastariul.github.io%2Fpintel" alt="Docs" /></a>
|
|
90
|
+
<a href="https://github.com/astariul/pintel/blob/main/LICENSE"><img src="https://img.shields.io/badge/License-MIT-yellow.svg" alt="licence" /></a>
|
|
91
|
+
</p>
|
|
92
|
+
|
|
93
|
+
<p align="center">
|
|
94
|
+
<a href="#description">Description</a> •
|
|
95
|
+
<a href="#install">Install</a> •
|
|
96
|
+
<a href="#usage">Usage</a> •
|
|
97
|
+
<a href="#faq">FAQ</a> •
|
|
98
|
+
<a href="#contribute">Contribute</a>
|
|
99
|
+
<br>
|
|
100
|
+
<a href="https://astariul.github.io/pintel/" target="_blank">Documentation</a>
|
|
101
|
+
</p>
|
|
102
|
+
|
|
103
|
+
|
|
104
|
+
<h2 align="center">Description</h2>
|
|
105
|
+
|
|
106
|
+
**`pintel`** is a python package that helps you scale your ML workers (or any python code) across processes and machines, asynchronously.
|
|
107
|
+
|
|
108
|
+
`pintel` is :
|
|
109
|
+
|
|
110
|
+
* ⚡️ Highly performant
|
|
111
|
+
* 🔀 Asynchronous
|
|
112
|
+
* 🐥 Easy-to-use
|
|
113
|
+
|
|
114
|
+
|
|
115
|
+
<h2 align="center">Install</h2>
|
|
116
|
+
|
|
117
|
+
Install `pintel` by running :
|
|
118
|
+
|
|
119
|
+
|
|
120
|
+
```
|
|
121
|
+
pip install pintel
|
|
122
|
+
```
|
|
123
|
+
|
|
124
|
+
|
|
125
|
+
<h2 align="center">Usage</h2>
|
|
126
|
+
|
|
127
|
+
After defining your awesome model :
|
|
128
|
+
|
|
129
|
+
```python
|
|
130
|
+
import time
|
|
131
|
+
|
|
132
|
+
class MyAwesomeModel:
|
|
133
|
+
def __init__(self, wait_time=0.25):
|
|
134
|
+
super().__init__()
|
|
135
|
+
self.w = wait_time
|
|
136
|
+
|
|
137
|
+
def __call__(self, x):
|
|
138
|
+
time.sleep(self.w)
|
|
139
|
+
return x**2
|
|
140
|
+
```
|
|
141
|
+
|
|
142
|
+
You can simply start a few workers serving the model :
|
|
143
|
+
|
|
144
|
+
```python
|
|
145
|
+
from pintel import Worker
|
|
146
|
+
|
|
147
|
+
for _ in range(4):
|
|
148
|
+
Worker(MyAwesomeModel).start()
|
|
149
|
+
```
|
|
150
|
+
|
|
151
|
+
And send requests through the Hub :
|
|
152
|
+
|
|
153
|
+
```python
|
|
154
|
+
from pintel import Hub
|
|
155
|
+
|
|
156
|
+
hub = Hub()
|
|
157
|
+
|
|
158
|
+
# In an async function
|
|
159
|
+
await hub.request(34)
|
|
160
|
+
```
|
|
161
|
+
|
|
162
|
+
And that's it !
|
|
163
|
+
|
|
164
|
+
---
|
|
165
|
+
|
|
166
|
+
Make sure to check the [documentation](https://astariul.github.io/pintel/usage/) for a more detailed explanation.
|
|
167
|
+
|
|
168
|
+
Or you can run some examples from the `examples/` folder.
|
|
169
|
+
|
|
170
|
+
|
|
171
|
+
<h2 align="center">FAQ</h2>
|
|
172
|
+
|
|
173
|
+
#### ❓ **How `pintel` works ?**
|
|
174
|
+
|
|
175
|
+
`pintel` simply run your model code in separate processes, and send requests to the right process to ensure requests are treated in parallel.
|
|
176
|
+
|
|
177
|
+
`pintel` uses a modified form of [the Paranoid Pirate Pattern from the zmq guide](https://zguide.zeromq.org/docs/chapter4/#Robust-Reliable-Queuing-Paranoid-Pirate-Pattern).
|
|
178
|
+
|
|
179
|
+
#### ❓ **Why the name "pintel" ?**
|
|
180
|
+
|
|
181
|
+
Pintel is one of the pirates of Captain Barbossa's crew on the Black Pearl.
|
|
182
|
+
Since we are using the Paranoid Pirate Pattern, we needed a pirate name !
|
|
183
|
+
|
|
184
|
+
<h2 align="center">Contribute</h2>
|
|
185
|
+
|
|
186
|
+
To contribute, install the package locally, create your own branch, add your code (and tests, and documentation), and open a PR !
|
|
187
|
+
|
|
188
|
+
### Pre-commit hooks
|
|
189
|
+
|
|
190
|
+
Pre-commit hooks are set to check the code added whenever you commit something.
|
|
191
|
+
|
|
192
|
+
If you never ran the hooks before, install it with :
|
|
193
|
+
|
|
194
|
+
```bash
|
|
195
|
+
pre-commit install
|
|
196
|
+
```
|
|
197
|
+
|
|
198
|
+
---
|
|
199
|
+
|
|
200
|
+
Then you can just try to commit your code. If you code does not meet the quality required by linters, it will not be committed. You can just fix your code and try to commit again !
|
|
201
|
+
|
|
202
|
+
---
|
|
203
|
+
|
|
204
|
+
You can manually run the pre-commit hooks with :
|
|
205
|
+
|
|
206
|
+
```bash
|
|
207
|
+
pre-commit run --all-files
|
|
208
|
+
```
|
|
209
|
+
|
|
210
|
+
### Tests
|
|
211
|
+
|
|
212
|
+
When you contribute, you need to make sure all the unit-tests pass. You should also add tests if necessary !
|
|
213
|
+
|
|
214
|
+
You can run the tests with :
|
|
215
|
+
|
|
216
|
+
```bash
|
|
217
|
+
pytest
|
|
218
|
+
```
|
|
219
|
+
|
|
220
|
+
---
|
|
221
|
+
|
|
222
|
+
Pre-commit hooks will not run the tests, but it will automatically update the coverage badge !
|
|
223
|
+
|
|
224
|
+
### Documentation
|
|
225
|
+
|
|
226
|
+
The documentation should be kept up-to-date. You can visualize the documentation locally by running :
|
|
227
|
+
|
|
228
|
+
```bash
|
|
229
|
+
mkdocs serve
|
|
230
|
+
```
|
|
@@ -0,0 +1,8 @@
|
|
|
1
|
+
pintel/__init__.py,sha256=O3Zed-_llzmSEe-eAGjuQ0v76N4PNLp-CbiGvgqDAGY,121
|
|
2
|
+
pintel/hub.py,sha256=wua2GZ-LmnVyy2WW1W4sS_CskLJA6aBHGb1nRPyR-Jg,11288
|
|
3
|
+
pintel/worker.py,sha256=Ipn3NQUgNK51nLoRsJJgDJutqeMyfiN-QqrYgUFxJgk,8476
|
|
4
|
+
pintel-0.3.0.dev0.dist-info/licenses/LICENSE,sha256=efxIJwqRW2jrT96PhExhD9OYd4DjZ9GKfq2UKwog8kc,1065
|
|
5
|
+
pintel-0.3.0.dev0.dist-info/METADATA,sha256=OtZNH3XidsQf3UxB50_5JN8ZSokRP1z68ekTAPtSX-M,7139
|
|
6
|
+
pintel-0.3.0.dev0.dist-info/WHEEL,sha256=SmOxYU7pzNKBqASvQJ7DjX3XGUF92lrGhMb3R6_iiqI,91
|
|
7
|
+
pintel-0.3.0.dev0.dist-info/top_level.txt,sha256=TJn487U_Ud0TaVQFbCER0Ki3fqSan5l6bVevHf54GWg,7
|
|
8
|
+
pintel-0.3.0.dev0.dist-info/RECORD,,
|
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2021 Astariul
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
pintel
|