wool 0.1rc14__tar.gz → 0.1rc15__tar.gz
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.
Potentially problematic release.
This version of wool might be problematic. Click here for more details.
- {wool-0.1rc14 → wool-0.1rc15}/PKG-INFO +2 -2
- {wool-0.1rc14 → wool-0.1rc15}/pyproject.toml +1 -1
- {wool-0.1rc14 → wool-0.1rc15}/wool/__init__.py +0 -23
- wool-0.1rc15/wool/_connection.py +247 -0
- wool-0.1rc15/wool/_context.py +29 -0
- wool-0.1rc15/wool/_loadbalancer.py +213 -0
- {wool-0.1rc14 → wool-0.1rc15}/wool/_protobuf/task_pb2_grpc.py +2 -2
- {wool-0.1rc14 → wool-0.1rc15}/wool/_protobuf/worker_pb2.py +6 -6
- {wool-0.1rc14 → wool-0.1rc15}/wool/_protobuf/worker_pb2.pyi +4 -4
- {wool-0.1rc14 → wool-0.1rc15}/wool/_protobuf/worker_pb2_grpc.py +2 -2
- {wool-0.1rc14 → wool-0.1rc15}/wool/_resource_pool.py +3 -3
- wool-0.1rc15/wool/_undefined.py +11 -0
- {wool-0.1rc14 → wool-0.1rc15}/wool/_work.py +5 -4
- wool-0.1rc15/wool/_worker.py +601 -0
- {wool-0.1rc14 → wool-0.1rc15}/wool/_worker_discovery.py +18 -13
- {wool-0.1rc14 → wool-0.1rc15}/wool/_worker_pool.py +2 -2
- {wool-0.1rc14 → wool-0.1rc15}/wool/_worker_proxy.py +54 -186
- wool-0.1rc15/wool/_worker_service.py +243 -0
- wool-0.1rc14/wool/_worker.py +0 -912
- {wool-0.1rc14 → wool-0.1rc15}/.gitignore +0 -0
- {wool-0.1rc14 → wool-0.1rc15}/README.md +0 -0
- {wool-0.1rc14 → wool-0.1rc15}/wool/_protobuf/__init__.py +0 -0
- {wool-0.1rc14 → wool-0.1rc15}/wool/_protobuf/exception.py +0 -0
- {wool-0.1rc14 → wool-0.1rc15}/wool/_protobuf/task.py +0 -0
- {wool-0.1rc14 → wool-0.1rc15}/wool/_protobuf/task_pb2.py +0 -0
- {wool-0.1rc14 → wool-0.1rc15}/wool/_protobuf/task_pb2.pyi +0 -0
- {wool-0.1rc14 → wool-0.1rc15}/wool/_protobuf/worker.py +0 -0
- {wool-0.1rc14 → wool-0.1rc15}/wool/_typing.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: wool
|
|
3
|
-
Version: 0.
|
|
3
|
+
Version: 0.1rc15
|
|
4
4
|
Summary: A Python framework for distributed multiprocessing.
|
|
5
5
|
Author-email: Conrad Bzura <conrad@wool.io>
|
|
6
6
|
Maintainer-email: maintainers@wool.io
|
|
@@ -222,7 +222,7 @@ Requires-Dist: hypothesis; extra == 'dev'
|
|
|
222
222
|
Requires-Dist: pytest; extra == 'dev'
|
|
223
223
|
Requires-Dist: pytest-asyncio; extra == 'dev'
|
|
224
224
|
Requires-Dist: pytest-cov; extra == 'dev'
|
|
225
|
-
Requires-Dist: pytest-grpc-aio~=0.
|
|
225
|
+
Requires-Dist: pytest-grpc-aio~=0.3.0; extra == 'dev'
|
|
226
226
|
Requires-Dist: pytest-mock; extra == 'dev'
|
|
227
227
|
Requires-Dist: ruff; extra == 'dev'
|
|
228
228
|
Description-Content-Type: text/markdown
|
|
@@ -2,9 +2,6 @@ from contextvars import ContextVar
|
|
|
2
2
|
from importlib.metadata import PackageNotFoundError
|
|
3
3
|
from importlib.metadata import version
|
|
4
4
|
from typing import Final
|
|
5
|
-
from typing import Generic
|
|
6
|
-
from typing import TypeVar
|
|
7
|
-
from typing import cast
|
|
8
5
|
|
|
9
6
|
from tblib import pickling_support
|
|
10
7
|
|
|
@@ -29,26 +26,6 @@ from wool._worker_proxy import WorkerProxy
|
|
|
29
26
|
pickling_support.install()
|
|
30
27
|
|
|
31
28
|
|
|
32
|
-
SENTINEL = object()
|
|
33
|
-
|
|
34
|
-
T = TypeVar("T")
|
|
35
|
-
|
|
36
|
-
|
|
37
|
-
class GlobalVar(Generic[T]):
|
|
38
|
-
def __init__(self, default: T | None = None) -> None:
|
|
39
|
-
self._default = default
|
|
40
|
-
self._value = SENTINEL
|
|
41
|
-
|
|
42
|
-
def get(self) -> T | None:
|
|
43
|
-
if self._value is SENTINEL:
|
|
44
|
-
return self._default
|
|
45
|
-
else:
|
|
46
|
-
return cast(T, self._value)
|
|
47
|
-
|
|
48
|
-
def set(self, value: T):
|
|
49
|
-
self._value = value
|
|
50
|
-
|
|
51
|
-
|
|
52
29
|
try:
|
|
53
30
|
__version__ = version("wool")
|
|
54
31
|
except PackageNotFoundError:
|
|
@@ -0,0 +1,247 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import asyncio
|
|
4
|
+
from typing import AsyncIterator
|
|
5
|
+
from typing import Final
|
|
6
|
+
from typing import Generic
|
|
7
|
+
from typing import TypeAlias
|
|
8
|
+
from typing import TypeVar
|
|
9
|
+
|
|
10
|
+
import cloudpickle
|
|
11
|
+
import grpc.aio
|
|
12
|
+
|
|
13
|
+
from wool import _protobuf as pb
|
|
14
|
+
from wool._work import WoolTask
|
|
15
|
+
|
|
16
|
+
_DispatchCall: TypeAlias = grpc.aio.UnaryStreamCall[pb.task.Task, pb.worker.Response]
|
|
17
|
+
|
|
18
|
+
_T = TypeVar("_T")
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
class _DispatchStream(Generic[_T]):
|
|
22
|
+
"""Async iterator wrapper for streaming task results from workers.
|
|
23
|
+
|
|
24
|
+
Handles iteration over gRPC response streams and deserializes
|
|
25
|
+
task results or raises exceptions received from remote workers.
|
|
26
|
+
|
|
27
|
+
:param call:
|
|
28
|
+
The underlying gRPC response stream.
|
|
29
|
+
"""
|
|
30
|
+
|
|
31
|
+
def __init__(self, call: _DispatchCall):
|
|
32
|
+
self._call = call
|
|
33
|
+
self._iter = aiter(call)
|
|
34
|
+
|
|
35
|
+
def __aiter__(self) -> AsyncIterator[_T]:
|
|
36
|
+
"""Return self as the async iterator."""
|
|
37
|
+
return self
|
|
38
|
+
|
|
39
|
+
async def __anext__(self) -> _T:
|
|
40
|
+
"""Get the next response from the stream.
|
|
41
|
+
|
|
42
|
+
:returns:
|
|
43
|
+
The next task result from the worker.
|
|
44
|
+
:raises StopAsyncIteration:
|
|
45
|
+
When the stream is exhausted.
|
|
46
|
+
:raises UnexpectedResponse:
|
|
47
|
+
If the response is neither a result nor an exception.
|
|
48
|
+
:raises Exception:
|
|
49
|
+
Any exception raised by the task execution is re-raised.
|
|
50
|
+
"""
|
|
51
|
+
try:
|
|
52
|
+
response = await anext(self._iter)
|
|
53
|
+
if response.HasField("result"):
|
|
54
|
+
return cloudpickle.loads(response.result.dump)
|
|
55
|
+
elif response.HasField("exception"):
|
|
56
|
+
raise cloudpickle.loads(response.exception.dump)
|
|
57
|
+
else:
|
|
58
|
+
raise UnexpectedResponse(
|
|
59
|
+
f"Expected 'result' or 'exception' response, "
|
|
60
|
+
f"received '{response.WhichOneof('payload')}'"
|
|
61
|
+
)
|
|
62
|
+
except Exception:
|
|
63
|
+
try:
|
|
64
|
+
self._call.cancel()
|
|
65
|
+
except Exception:
|
|
66
|
+
pass
|
|
67
|
+
raise
|
|
68
|
+
|
|
69
|
+
|
|
70
|
+
# public
|
|
71
|
+
class UnexpectedResponse(Exception):
|
|
72
|
+
"""Raised when a worker returns an unexpected response type.
|
|
73
|
+
|
|
74
|
+
This exception indicates a protocol violation where the worker's
|
|
75
|
+
response doesn't match the expected format (e.g., missing acknowledgment
|
|
76
|
+
or returning an unrecognized payload type).
|
|
77
|
+
"""
|
|
78
|
+
|
|
79
|
+
|
|
80
|
+
# public
|
|
81
|
+
class RpcError(Exception):
|
|
82
|
+
"""Raised when a gRPC call to a worker fails with a non-transient error.
|
|
83
|
+
|
|
84
|
+
Non-transient errors indicate persistent issues with the worker that
|
|
85
|
+
are unlikely to be resolved by retrying (e.g., invalid arguments,
|
|
86
|
+
unimplemented methods, permission denied).
|
|
87
|
+
"""
|
|
88
|
+
|
|
89
|
+
|
|
90
|
+
# public
|
|
91
|
+
class TransientRpcError(RpcError):
|
|
92
|
+
"""Raised when a gRPC call to a worker fails with a transient error.
|
|
93
|
+
|
|
94
|
+
Transient errors indicate temporary issues that may be resolved by
|
|
95
|
+
retrying the operation, such as:
|
|
96
|
+
|
|
97
|
+
- ``UNAVAILABLE``: Worker temporarily unavailable
|
|
98
|
+
- ``DEADLINE_EXCEEDED``: Request took too long
|
|
99
|
+
- ``RESOURCE_EXHAUSTED``: Worker temporarily overloaded
|
|
100
|
+
"""
|
|
101
|
+
|
|
102
|
+
|
|
103
|
+
# public
|
|
104
|
+
class Connection:
|
|
105
|
+
"""Connection to a remote worker service used for dispatching tasks.
|
|
106
|
+
|
|
107
|
+
Maintains a persistent gRPC channel to a single worker and manages
|
|
108
|
+
the channel lifecycle. Provides task dispatch functionality with
|
|
109
|
+
concurrency control.
|
|
110
|
+
|
|
111
|
+
:param target:
|
|
112
|
+
The target URI for the worker service to connect to. Can be specified
|
|
113
|
+
in several formats:
|
|
114
|
+
|
|
115
|
+
- ``host:port`` - DNS name or IP with port (defaults to dns scheme)
|
|
116
|
+
- ``dns://host:port`` - Explicit DNS resolution
|
|
117
|
+
- ``ipv4:address:port`` - IPv4 address
|
|
118
|
+
- ``ipv6:[address]:port`` - IPv6 address (brackets required)
|
|
119
|
+
- ``unix:path`` or ``unix:///path`` - Unix domain socket
|
|
120
|
+
|
|
121
|
+
If no scheme is specified, the dns scheme is used by default.
|
|
122
|
+
Examples: ``localhost:50051``, ``dns://example.com:8080``,
|
|
123
|
+
``ipv4:192.0.2.1:50051``.
|
|
124
|
+
|
|
125
|
+
See https://github.com/grpc/grpc/blob/master/doc/naming.md for
|
|
126
|
+
complete URI format specifications.
|
|
127
|
+
:param limit:
|
|
128
|
+
Maximum number of concurrent task dispatches allowed.
|
|
129
|
+
"""
|
|
130
|
+
|
|
131
|
+
TRANSIENT_ERRORS: Final = {
|
|
132
|
+
grpc.StatusCode.UNAVAILABLE,
|
|
133
|
+
grpc.StatusCode.DEADLINE_EXCEEDED,
|
|
134
|
+
grpc.StatusCode.RESOURCE_EXHAUSTED,
|
|
135
|
+
}
|
|
136
|
+
|
|
137
|
+
def __init__(
|
|
138
|
+
self,
|
|
139
|
+
target: str,
|
|
140
|
+
*,
|
|
141
|
+
limit: int = 100,
|
|
142
|
+
):
|
|
143
|
+
if limit <= 0:
|
|
144
|
+
raise ValueError("Limit must be positive")
|
|
145
|
+
self._channel = grpc.aio.insecure_channel(
|
|
146
|
+
target,
|
|
147
|
+
options=[
|
|
148
|
+
("grpc.keepalive_time_ms", 10000),
|
|
149
|
+
("grpc.keepalive_timeout_ms", 5000),
|
|
150
|
+
("grpc.http2.max_pings_without_data", 0),
|
|
151
|
+
("grpc.http2.min_time_between_pings_ms", 10000),
|
|
152
|
+
("grpc.max_receive_message_length", 100 * 1024 * 1024),
|
|
153
|
+
("grpc.max_send_message_length", 100 * 1024 * 1024),
|
|
154
|
+
],
|
|
155
|
+
)
|
|
156
|
+
self._stub = pb.worker.WorkerStub(self._channel)
|
|
157
|
+
self._semaphore = asyncio.Semaphore(limit)
|
|
158
|
+
|
|
159
|
+
async def dispatch(
|
|
160
|
+
self,
|
|
161
|
+
task: WoolTask,
|
|
162
|
+
*,
|
|
163
|
+
timeout: float | None = None,
|
|
164
|
+
) -> AsyncIterator[pb.task.Result]:
|
|
165
|
+
"""Dispatch a task to the remote worker for execution.
|
|
166
|
+
|
|
167
|
+
Sends the task to the worker via gRPC, waits for acknowledgment,
|
|
168
|
+
and returns an async iterator that streams back results. Respects
|
|
169
|
+
concurrency limits and applies timeout to the dispatch phase only
|
|
170
|
+
(semaphore acquisition and acknowledgment).
|
|
171
|
+
|
|
172
|
+
:param task:
|
|
173
|
+
The :class:`WoolTask` instance to dispatch to the worker.
|
|
174
|
+
:param timeout:
|
|
175
|
+
Timeout in seconds for semaphore acquisition and task
|
|
176
|
+
acknowledgment. If ``None``, no timeout is applied. Does not
|
|
177
|
+
apply to the execution phase.
|
|
178
|
+
:returns:
|
|
179
|
+
An async iterator that yields task results from the worker.
|
|
180
|
+
:raises TransientRpcError:
|
|
181
|
+
If the worker returns a transient RPC error (UNAVAILABLE,
|
|
182
|
+
DEADLINE_EXCEEDED, or RESOURCE_EXHAUSTED).
|
|
183
|
+
:raises RpcError:
|
|
184
|
+
If the worker returns a non-transient RPC error.
|
|
185
|
+
:raises UnexpectedResponse:
|
|
186
|
+
If the worker doesn't acknowledge the task.
|
|
187
|
+
:raises TimeoutError:
|
|
188
|
+
If the timeout is exceeded during dispatch phase.
|
|
189
|
+
:raises ValueError:
|
|
190
|
+
If the timeout value is not positive.
|
|
191
|
+
"""
|
|
192
|
+
if timeout is not None and timeout <= 0:
|
|
193
|
+
raise ValueError("Dispatch timeout must be positive")
|
|
194
|
+
|
|
195
|
+
try:
|
|
196
|
+
call = await self._dispatch(task, timeout)
|
|
197
|
+
except grpc.RpcError as error:
|
|
198
|
+
if error.code() in self.TRANSIENT_ERRORS:
|
|
199
|
+
raise TransientRpcError from error
|
|
200
|
+
else:
|
|
201
|
+
raise RpcError from error
|
|
202
|
+
|
|
203
|
+
return self._execute(call)
|
|
204
|
+
|
|
205
|
+
async def close(self):
|
|
206
|
+
"""Close the connection and release all resources.
|
|
207
|
+
|
|
208
|
+
Gracefully closes the underlying gRPC channel to the remote
|
|
209
|
+
worker and cleans up any associated resources.
|
|
210
|
+
"""
|
|
211
|
+
await self._channel.close()
|
|
212
|
+
|
|
213
|
+
async def _dispatch(self, task, timeout):
|
|
214
|
+
async with asyncio.timeout(timeout):
|
|
215
|
+
await self._semaphore.acquire()
|
|
216
|
+
try:
|
|
217
|
+
call: _DispatchCall = self._stub.dispatch(task.to_protobuf())
|
|
218
|
+
try:
|
|
219
|
+
response = await anext(aiter(call))
|
|
220
|
+
if not response.HasField("ack"):
|
|
221
|
+
raise UnexpectedResponse(
|
|
222
|
+
f"Expected 'ack' response, "
|
|
223
|
+
f"received '{response.WhichOneof('payload')}'"
|
|
224
|
+
)
|
|
225
|
+
except (Exception, asyncio.CancelledError):
|
|
226
|
+
try:
|
|
227
|
+
call.cancel()
|
|
228
|
+
except Exception:
|
|
229
|
+
pass
|
|
230
|
+
raise
|
|
231
|
+
except (Exception, asyncio.CancelledError):
|
|
232
|
+
self._semaphore.release()
|
|
233
|
+
raise
|
|
234
|
+
return call
|
|
235
|
+
|
|
236
|
+
async def _execute(self, call):
|
|
237
|
+
try:
|
|
238
|
+
async for result in _DispatchStream(call):
|
|
239
|
+
yield result
|
|
240
|
+
except (Exception, asyncio.CancelledError):
|
|
241
|
+
try:
|
|
242
|
+
call.cancel()
|
|
243
|
+
except Exception:
|
|
244
|
+
pass
|
|
245
|
+
raise
|
|
246
|
+
finally:
|
|
247
|
+
self._semaphore.release()
|
|
@@ -0,0 +1,29 @@
|
|
|
1
|
+
from contextvars import ContextVar
|
|
2
|
+
from contextvars import Token
|
|
3
|
+
from typing import Final
|
|
4
|
+
|
|
5
|
+
from wool._undefined import Undefined
|
|
6
|
+
from wool._undefined import UndefinedType
|
|
7
|
+
|
|
8
|
+
dispatch_timeout: Final[ContextVar[float | None]] = ContextVar(
|
|
9
|
+
"_dispatch_timeout", default=None
|
|
10
|
+
)
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
# public
|
|
14
|
+
class AppContext:
|
|
15
|
+
_dispatch_timeout: float | None | UndefinedType
|
|
16
|
+
_dispatch_timeout_token: Token | UndefinedType
|
|
17
|
+
|
|
18
|
+
def __init__(self, *, dispatch_timeout: float | None | UndefinedType = Undefined):
|
|
19
|
+
self._dispatch_timeout = dispatch_timeout
|
|
20
|
+
|
|
21
|
+
def __enter__(self):
|
|
22
|
+
if self._dispatch_timeout is not Undefined:
|
|
23
|
+
self._dispatch_timeout_token = dispatch_timeout.set(self._dispatch_timeout)
|
|
24
|
+
else:
|
|
25
|
+
self._dispatch_timeout_token = Undefined
|
|
26
|
+
|
|
27
|
+
def __exit__(self, *_):
|
|
28
|
+
if self._dispatch_timeout_token is not Undefined:
|
|
29
|
+
dispatch_timeout.reset(self._dispatch_timeout_token)
|
|
@@ -0,0 +1,213 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import itertools
|
|
4
|
+
from types import MappingProxyType
|
|
5
|
+
from typing import TYPE_CHECKING
|
|
6
|
+
from typing import AsyncIterator
|
|
7
|
+
from typing import Callable
|
|
8
|
+
from typing import Final
|
|
9
|
+
from typing import Protocol
|
|
10
|
+
from typing import TypeAlias
|
|
11
|
+
from typing import runtime_checkable
|
|
12
|
+
|
|
13
|
+
from wool._connection import Connection
|
|
14
|
+
from wool._connection import TransientRpcError
|
|
15
|
+
from wool._resource_pool import Resource
|
|
16
|
+
from wool._worker_discovery import WorkerInfo
|
|
17
|
+
|
|
18
|
+
if TYPE_CHECKING:
|
|
19
|
+
from wool._work import WoolTask
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
# public
|
|
23
|
+
ConnectionResourceFactory: TypeAlias = Callable[[], Resource[Connection]]
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
# public
|
|
27
|
+
class NoWorkersAvailable(Exception):
|
|
28
|
+
"""Raised when no workers are available for task dispatch.
|
|
29
|
+
|
|
30
|
+
This exception indicates that either no workers exist in the worker pool
|
|
31
|
+
or all available workers have been tried and failed.
|
|
32
|
+
"""
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
# public
|
|
36
|
+
@runtime_checkable
|
|
37
|
+
class LoadBalancerLike(Protocol):
|
|
38
|
+
"""Protocol for load balancers that dispatch tasks to workers.
|
|
39
|
+
|
|
40
|
+
Load balancers implementing this protocol operate on a
|
|
41
|
+
:class:`LoadBalancerContext` to access workers and their connection
|
|
42
|
+
factories. The context provides isolation, allowing a single load balancer
|
|
43
|
+
instance to service multiple worker pools with independent state.
|
|
44
|
+
|
|
45
|
+
The dispatch method accepts a :class:`WoolTask` and returns an async
|
|
46
|
+
iterator that yields task results from the worker.
|
|
47
|
+
"""
|
|
48
|
+
|
|
49
|
+
async def dispatch(
|
|
50
|
+
self,
|
|
51
|
+
task: WoolTask,
|
|
52
|
+
*,
|
|
53
|
+
context: LoadBalancerContext,
|
|
54
|
+
timeout: float | None = None,
|
|
55
|
+
) -> AsyncIterator: ...
|
|
56
|
+
|
|
57
|
+
|
|
58
|
+
# public
|
|
59
|
+
class LoadBalancerContext:
|
|
60
|
+
"""Isolated load balancing context for a single worker pool.
|
|
61
|
+
|
|
62
|
+
Manages workers and their connection resource factories for a specific
|
|
63
|
+
worker pool, enabling load balancer instances to service multiple pools
|
|
64
|
+
with independent state and worker lists.
|
|
65
|
+
"""
|
|
66
|
+
|
|
67
|
+
_workers: Final[dict[WorkerInfo, ConnectionResourceFactory]]
|
|
68
|
+
|
|
69
|
+
def __init__(self):
|
|
70
|
+
self._workers = {}
|
|
71
|
+
|
|
72
|
+
@property
|
|
73
|
+
def workers(self) -> MappingProxyType[WorkerInfo, ConnectionResourceFactory]:
|
|
74
|
+
"""Read-only view of workers in this context.
|
|
75
|
+
|
|
76
|
+
:returns:
|
|
77
|
+
Immutable mapping of worker information to connection resource
|
|
78
|
+
factories. Changes to the underlying context are reflected in
|
|
79
|
+
the returned proxy.
|
|
80
|
+
"""
|
|
81
|
+
return MappingProxyType(self._workers)
|
|
82
|
+
|
|
83
|
+
def add_worker(
|
|
84
|
+
self,
|
|
85
|
+
worker_info: WorkerInfo,
|
|
86
|
+
connection_resource_factory: ConnectionResourceFactory,
|
|
87
|
+
):
|
|
88
|
+
"""Add a worker to this context.
|
|
89
|
+
|
|
90
|
+
:param worker_info:
|
|
91
|
+
Information about the worker to add.
|
|
92
|
+
:param connection_resource_factory:
|
|
93
|
+
Factory function that creates connection resources for this worker.
|
|
94
|
+
"""
|
|
95
|
+
self._workers[worker_info] = connection_resource_factory
|
|
96
|
+
|
|
97
|
+
def update_worker(
|
|
98
|
+
self,
|
|
99
|
+
worker_info: WorkerInfo,
|
|
100
|
+
connection_resource_factory: ConnectionResourceFactory,
|
|
101
|
+
*,
|
|
102
|
+
upsert: bool = False,
|
|
103
|
+
):
|
|
104
|
+
"""Update an existing worker's connection resource factory.
|
|
105
|
+
|
|
106
|
+
:param worker_info:
|
|
107
|
+
Information about the worker to update. If the worker is not
|
|
108
|
+
present in the context, this method does nothing.
|
|
109
|
+
:param connection_resource_factory:
|
|
110
|
+
New factory function that creates connection resources for this
|
|
111
|
+
worker.
|
|
112
|
+
:param upsert:
|
|
113
|
+
Flag indicating whether or not to add the worker if it's not
|
|
114
|
+
already in the context.
|
|
115
|
+
"""
|
|
116
|
+
if upsert or worker_info in self._workers:
|
|
117
|
+
self._workers[worker_info] = connection_resource_factory
|
|
118
|
+
|
|
119
|
+
def remove_worker(self, worker_info: WorkerInfo):
|
|
120
|
+
"""Remove a worker from this context.
|
|
121
|
+
|
|
122
|
+
:param worker_info:
|
|
123
|
+
Information about the worker to remove. If the worker is not
|
|
124
|
+
present in the context, this method does nothing.
|
|
125
|
+
"""
|
|
126
|
+
if worker_info in self._workers:
|
|
127
|
+
del self._workers[worker_info]
|
|
128
|
+
|
|
129
|
+
|
|
130
|
+
# public
|
|
131
|
+
class RoundRobinLoadBalancer(LoadBalancerLike):
|
|
132
|
+
"""Round-robin load balancer for distributing tasks across workers.
|
|
133
|
+
|
|
134
|
+
Distributes tasks evenly across available workers using a simple round-robin
|
|
135
|
+
algorithm. Workers are managed through :class:`LoadBalancerContext` instances
|
|
136
|
+
passed to the dispatch method, enabling a single load balancer to service
|
|
137
|
+
multiple worker pools with independent state.
|
|
138
|
+
|
|
139
|
+
Automatically handles worker failures by trying the next worker in the
|
|
140
|
+
round-robin cycle. Workers that encounter transient errors remain in the
|
|
141
|
+
context, while workers that fail with non-transient errors are removed from
|
|
142
|
+
the context's worker list.
|
|
143
|
+
"""
|
|
144
|
+
|
|
145
|
+
_index: Final[dict[LoadBalancerContext, int]]
|
|
146
|
+
|
|
147
|
+
def __init__(self):
|
|
148
|
+
self._index = {}
|
|
149
|
+
|
|
150
|
+
async def dispatch(
|
|
151
|
+
self,
|
|
152
|
+
task: WoolTask,
|
|
153
|
+
*,
|
|
154
|
+
context: LoadBalancerContext,
|
|
155
|
+
timeout: float | None = None,
|
|
156
|
+
) -> AsyncIterator:
|
|
157
|
+
"""Dispatch a task to the next available worker using round-robin.
|
|
158
|
+
|
|
159
|
+
Tries workers in one round-robin cycle until dispatch succeeds.
|
|
160
|
+
Workers that fail to schedule the task with a non-transient error are
|
|
161
|
+
removed from the context's worker list.
|
|
162
|
+
|
|
163
|
+
:param task:
|
|
164
|
+
The :class:`WoolTask` instance to dispatch to the worker.
|
|
165
|
+
:param context:
|
|
166
|
+
The :class:`LoadBalancerContext` containing workers to dispatch to.
|
|
167
|
+
:param timeout:
|
|
168
|
+
Timeout in seconds for each dispatch attempt. If ``None``, no
|
|
169
|
+
timeout is applied.
|
|
170
|
+
:returns:
|
|
171
|
+
A streaming dispatch result that yields worker responses.
|
|
172
|
+
:raises NoWorkersAvailable:
|
|
173
|
+
If no healthy workers are available to schedule the task.
|
|
174
|
+
"""
|
|
175
|
+
checkpoint = None
|
|
176
|
+
|
|
177
|
+
# Initialize index for this context if not present
|
|
178
|
+
if context not in self._index:
|
|
179
|
+
self._index[context] = 0
|
|
180
|
+
|
|
181
|
+
while context.workers:
|
|
182
|
+
if self._index[context] >= len(context.workers):
|
|
183
|
+
self._index[context] = 0
|
|
184
|
+
|
|
185
|
+
worker_info, connection_resource_factory = next(
|
|
186
|
+
itertools.islice(
|
|
187
|
+
context.workers.items(),
|
|
188
|
+
self._index[context],
|
|
189
|
+
self._index[context] + 1,
|
|
190
|
+
)
|
|
191
|
+
)
|
|
192
|
+
|
|
193
|
+
if checkpoint is None:
|
|
194
|
+
checkpoint = worker_info.uid
|
|
195
|
+
elif worker_info.uid == checkpoint:
|
|
196
|
+
break
|
|
197
|
+
|
|
198
|
+
async with connection_resource_factory() as connection:
|
|
199
|
+
try:
|
|
200
|
+
result = await connection.dispatch(task, timeout=timeout)
|
|
201
|
+
except TransientRpcError:
|
|
202
|
+
self._index[context] = self._index[context] + 1
|
|
203
|
+
continue
|
|
204
|
+
except Exception:
|
|
205
|
+
context.remove_worker(worker_info)
|
|
206
|
+
if worker_info.uid == checkpoint:
|
|
207
|
+
checkpoint = None
|
|
208
|
+
continue
|
|
209
|
+
else:
|
|
210
|
+
self._index[context] = self._index[context] + 1
|
|
211
|
+
return result
|
|
212
|
+
|
|
213
|
+
raise NoWorkersAvailable("No healthy workers available for dispatch")
|
|
@@ -4,7 +4,7 @@ import grpc
|
|
|
4
4
|
import warnings
|
|
5
5
|
|
|
6
6
|
|
|
7
|
-
GRPC_GENERATED_VERSION = '1.
|
|
7
|
+
GRPC_GENERATED_VERSION = '1.76.0'
|
|
8
8
|
GRPC_VERSION = grpc.__version__
|
|
9
9
|
_version_not_supported = False
|
|
10
10
|
|
|
@@ -17,7 +17,7 @@ except ImportError:
|
|
|
17
17
|
if _version_not_supported:
|
|
18
18
|
raise RuntimeError(
|
|
19
19
|
f'The grpc package installed is at version {GRPC_VERSION},'
|
|
20
|
-
+
|
|
20
|
+
+ ' but the generated code in task_pb2_grpc.py depends on'
|
|
21
21
|
+ f' grpcio>={GRPC_GENERATED_VERSION}.'
|
|
22
22
|
+ f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}'
|
|
23
23
|
+ f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.'
|
|
@@ -25,7 +25,7 @@ _sym_db = _symbol_database.Default()
|
|
|
25
25
|
import task_pb2 as task__pb2
|
|
26
26
|
|
|
27
27
|
|
|
28
|
-
DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x0cworker.proto\x12\x15wool._protobuf.worker\x1a\ntask.proto\"\xd1\x01\n\x08Response\x12)\n\x03\x61\x63k\x18\x01 \x01(\x0b\x32\x1a.wool._protobuf.worker.AckH\x00\x12+\n\x04nack\x18\x02 \x01(\x0b\x32\x1b.wool._protobuf.worker.NackH\x00\x12-\n\x06result\x18\x03 \x01(\x0b\x32\x1b.wool._protobuf.task.ResultH\x00\x12\x33\n\texception\x18\x04 \x01(\x0b\x32\x1e.wool._protobuf.task.ExceptionH\x00\x42\t\n\x07payload\"\x05\n\x03\x41\x63k\"\x16\n\x04Nack\x12\x0e\n\x06reason\x18\x01 \x01(\t\"\
|
|
28
|
+
DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x0cworker.proto\x12\x15wool._protobuf.worker\x1a\ntask.proto\"\xd1\x01\n\x08Response\x12)\n\x03\x61\x63k\x18\x01 \x01(\x0b\x32\x1a.wool._protobuf.worker.AckH\x00\x12+\n\x04nack\x18\x02 \x01(\x0b\x32\x1b.wool._protobuf.worker.NackH\x00\x12-\n\x06result\x18\x03 \x01(\x0b\x32\x1b.wool._protobuf.task.ResultH\x00\x12\x33\n\texception\x18\x04 \x01(\x0b\x32\x1e.wool._protobuf.task.ExceptionH\x00\x42\t\n\x07payload\"\x05\n\x03\x41\x63k\"\x16\n\x04Nack\x12\x0e\n\x06reason\x18\x01 \x01(\t\"\x1e\n\x0bStopRequest\x12\x0f\n\x07timeout\x18\x01 \x01(\x02\"\x06\n\x04Void2\x9b\x01\n\x06Worker\x12H\n\x08\x64ispatch\x12\x19.wool._protobuf.task.Task\x1a\x1f.wool._protobuf.worker.Response0\x01\x12G\n\x04stop\x12\".wool._protobuf.worker.StopRequest\x1a\x1b.wool._protobuf.worker.Voidb\x06proto3')
|
|
29
29
|
|
|
30
30
|
_globals = globals()
|
|
31
31
|
_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
|
|
@@ -39,9 +39,9 @@ if not _descriptor._USE_C_DESCRIPTORS:
|
|
|
39
39
|
_globals['_NACK']._serialized_start=270
|
|
40
40
|
_globals['_NACK']._serialized_end=292
|
|
41
41
|
_globals['_STOPREQUEST']._serialized_start=294
|
|
42
|
-
_globals['_STOPREQUEST']._serialized_end=
|
|
43
|
-
_globals['_VOID']._serialized_start=
|
|
44
|
-
_globals['_VOID']._serialized_end=
|
|
45
|
-
_globals['_WORKER']._serialized_start=
|
|
46
|
-
_globals['_WORKER']._serialized_end=
|
|
42
|
+
_globals['_STOPREQUEST']._serialized_end=324
|
|
43
|
+
_globals['_VOID']._serialized_start=326
|
|
44
|
+
_globals['_VOID']._serialized_end=332
|
|
45
|
+
_globals['_WORKER']._serialized_start=335
|
|
46
|
+
_globals['_WORKER']._serialized_end=490
|
|
47
47
|
# @@protoc_insertion_point(module_scope)
|
|
@@ -29,10 +29,10 @@ class Nack(_message.Message):
|
|
|
29
29
|
def __init__(self, reason: _Optional[str] = ...) -> None: ...
|
|
30
30
|
|
|
31
31
|
class StopRequest(_message.Message):
|
|
32
|
-
__slots__ = ("
|
|
33
|
-
|
|
34
|
-
|
|
35
|
-
def __init__(self,
|
|
32
|
+
__slots__ = ("timeout",)
|
|
33
|
+
TIMEOUT_FIELD_NUMBER: _ClassVar[int]
|
|
34
|
+
timeout: float
|
|
35
|
+
def __init__(self, timeout: _Optional[float] = ...) -> None: ...
|
|
36
36
|
|
|
37
37
|
class Void(_message.Message):
|
|
38
38
|
__slots__ = ()
|
|
@@ -6,7 +6,7 @@ import warnings
|
|
|
6
6
|
from . import task_pb2 as task__pb2
|
|
7
7
|
from . import worker_pb2 as worker__pb2
|
|
8
8
|
|
|
9
|
-
GRPC_GENERATED_VERSION = '1.
|
|
9
|
+
GRPC_GENERATED_VERSION = '1.76.0'
|
|
10
10
|
GRPC_VERSION = grpc.__version__
|
|
11
11
|
_version_not_supported = False
|
|
12
12
|
|
|
@@ -19,7 +19,7 @@ except ImportError:
|
|
|
19
19
|
if _version_not_supported:
|
|
20
20
|
raise RuntimeError(
|
|
21
21
|
f'The grpc package installed is at version {GRPC_VERSION},'
|
|
22
|
-
+
|
|
22
|
+
+ ' but the generated code in worker_pb2_grpc.py depends on'
|
|
23
23
|
+ f' grpcio>={GRPC_GENERATED_VERSION}.'
|
|
24
24
|
+ f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}'
|
|
25
25
|
+ f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.'
|
|
@@ -25,7 +25,7 @@ class Resource(Generic[T]):
|
|
|
25
25
|
released again.
|
|
26
26
|
|
|
27
27
|
:param pool:
|
|
28
|
-
The :
|
|
28
|
+
The :class:`ResourcePool` this resource belongs to.
|
|
29
29
|
:param key:
|
|
30
30
|
The cache key for this resource.
|
|
31
31
|
"""
|
|
@@ -185,7 +185,7 @@ class ResourcePool(Generic[T]):
|
|
|
185
185
|
when not concurrently modifying the cache.
|
|
186
186
|
|
|
187
187
|
:returns:
|
|
188
|
-
:
|
|
188
|
+
:class:`ResourcePool.Stats` containing current statistics.
|
|
189
189
|
"""
|
|
190
190
|
pending_cleanup = sum(
|
|
191
191
|
1 for c in self.pending_cleanup.values() if c is not None and not c.done()
|
|
@@ -219,7 +219,7 @@ class ResourcePool(Generic[T]):
|
|
|
219
219
|
:param key:
|
|
220
220
|
The cache key.
|
|
221
221
|
:returns:
|
|
222
|
-
:
|
|
222
|
+
:class:`Resource` that can be awaited or used with 'async with'.
|
|
223
223
|
"""
|
|
224
224
|
return Resource(self, key)
|
|
225
225
|
|