shared-tensor 0.2.9__tar.gz → 0.2.11__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.
- shared_tensor-0.2.11/LICENSE +21 -0
- {shared_tensor-0.2.9 → shared_tensor-0.2.11}/PKG-INFO +69 -4
- {shared_tensor-0.2.9 → shared_tensor-0.2.11}/README.md +67 -2
- {shared_tensor-0.2.9 → shared_tensor-0.2.11}/pyproject.toml +2 -2
- {shared_tensor-0.2.9 → shared_tensor-0.2.11}/shared_tensor/__init__.py +3 -3
- {shared_tensor-0.2.9 → shared_tensor-0.2.11}/shared_tensor/async_client.py +33 -0
- {shared_tensor-0.2.9 → shared_tensor-0.2.11}/shared_tensor/client.py +62 -2
- {shared_tensor-0.2.9 → shared_tensor-0.2.11}/shared_tensor/errors.py +18 -0
- {shared_tensor-0.2.9 → shared_tensor-0.2.11}/shared_tensor/managed_object.py +50 -1
- {shared_tensor-0.2.9 → shared_tensor-0.2.11}/shared_tensor/provider.py +34 -0
- {shared_tensor-0.2.9 → shared_tensor-0.2.11}/shared_tensor/server.py +117 -10
- {shared_tensor-0.2.9 → shared_tensor-0.2.11}/shared_tensor/utils.py +70 -4
- {shared_tensor-0.2.9 → shared_tensor-0.2.11}/shared_tensor.egg-info/SOURCES.txt +0 -1
- shared_tensor-0.2.9/LICENSE +0 -181
- shared_tensor-0.2.9/shared_tensor/async_provider.py +0 -97
- {shared_tensor-0.2.9 → shared_tensor-0.2.11}/MANIFEST.in +0 -0
- {shared_tensor-0.2.9 → shared_tensor-0.2.11}/setup.cfg +0 -0
- {shared_tensor-0.2.9 → shared_tensor-0.2.11}/shared_tensor/async_task.py +0 -0
- {shared_tensor-0.2.9 → shared_tensor-0.2.11}/shared_tensor/runtime.py +0 -0
- {shared_tensor-0.2.9 → shared_tensor-0.2.11}/shared_tensor/transport.py +0 -0
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2024 Athena Team
|
|
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.
|
|
@@ -1,10 +1,10 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: shared-tensor
|
|
3
|
-
Version: 0.2.
|
|
3
|
+
Version: 0.2.11
|
|
4
4
|
Summary: Native PyTorch CUDA IPC over Unix Domain Socket for same-host process separation
|
|
5
5
|
Author-email: Athena Team <contact@world-sim-dev.org>
|
|
6
6
|
Maintainer-email: Athena Team <contact@world-sim-dev.org>
|
|
7
|
-
License-Expression:
|
|
7
|
+
License-Expression: MIT
|
|
8
8
|
Project-URL: Homepage, https://github.com/world-sim-dev/shared-tensor
|
|
9
9
|
Project-URL: Repository, https://github.com/world-sim-dev/shared-tensor
|
|
10
10
|
Project-URL: Documentation, https://github.com/world-sim-dev/shared-tensor/tree/main/docs
|
|
@@ -65,7 +65,7 @@ Supported:
|
|
|
65
65
|
- explicit endpoint registration
|
|
66
66
|
- sync `call` and task-backed `submit`
|
|
67
67
|
- managed object handles with explicit release
|
|
68
|
-
- server-side caching, `cache_format_key`, and
|
|
68
|
+
- server-side caching, `cache_format_key`, singleflight, and explicit cache invalidation
|
|
69
69
|
- manual two-process deployment as the primary production path
|
|
70
70
|
- zero-branch auto mode gated by `SHARED_TENSOR_ENABLED=1`
|
|
71
71
|
|
|
@@ -108,6 +108,8 @@ Production should prefer two explicitly started processes: one server process th
|
|
|
108
108
|
|
|
109
109
|
See [examples/model_service.py](./examples/model_service.py) for endpoint definitions.
|
|
110
110
|
|
|
111
|
+
The server-oriented example modules construct providers with explicit `execution_mode="server"` so importing the module already reflects the intended deployment role.
|
|
112
|
+
|
|
111
113
|
Server process:
|
|
112
114
|
|
|
113
115
|
```python
|
|
@@ -154,6 +156,7 @@ manages cache and refcounts releases managed handles explicitly
|
|
|
154
156
|
|
|
155
157
|
Core assumption:
|
|
156
158
|
- the server process that owns the original CUDA allocation must stay alive while clients are still using reopened CUDA tensors or modules
|
|
159
|
+
- handle health checks can detect some stale-object conditions, but they do not remove the producer-liveness requirement
|
|
157
160
|
|
|
158
161
|
If the server exits, crashes, or is killed before the client is done with the shared CUDA object, behavior is no longer guaranteed by this library. Depending on PyTorch and CUDA runtime state, the client may see CUDA runtime errors, invalid resource handle failures, broken module execution, or process-level instability.
|
|
159
162
|
|
|
@@ -168,6 +171,11 @@ Treat producer liveness as a hard requirement, not a soft optimization.
|
|
|
168
171
|
|
|
169
172
|
See [examples/zero_branch_env.py](./examples/zero_branch_env.py). This is a convenience mode for environments that want one file and environment-controlled behavior.
|
|
170
173
|
|
|
174
|
+
Resolution rule:
|
|
175
|
+
- `SHARED_TENSOR_ENABLED` unset or false: provider stays local
|
|
176
|
+
- `SHARED_TENSOR_ENABLED=1` and `SHARED_TENSOR_ROLE=server`: provider resolves to server and auto-starts the thread-backed local server
|
|
177
|
+
- `SHARED_TENSOR_ENABLED=1` and role unset or `client`: provider resolves to client
|
|
178
|
+
|
|
171
179
|
```bash
|
|
172
180
|
SHARED_TENSOR_ENABLED=1 SHARED_TENSOR_ROLE=server python demo.py
|
|
173
181
|
SHARED_TENSOR_ENABLED=1 python demo.py
|
|
@@ -185,6 +193,27 @@ shared function runs locally shared function becomes RPC call
|
|
|
185
193
|
CUDA object stays on same GPU CUDA object is reopened via torch IPC
|
|
186
194
|
```
|
|
187
195
|
|
|
196
|
+
## Example: Task Submission And Wait
|
|
197
|
+
|
|
198
|
+
See [examples/async_service.py](./examples/async_service.py).
|
|
199
|
+
|
|
200
|
+
```python
|
|
201
|
+
from shared_tensor import AsyncSharedTensorClient, SharedTensorProvider
|
|
202
|
+
|
|
203
|
+
provider = SharedTensorProvider(execution_mode="server")
|
|
204
|
+
|
|
205
|
+
@provider.share(execution="task")
|
|
206
|
+
def build_delayed_model(delay: float = 0.1):
|
|
207
|
+
...
|
|
208
|
+
|
|
209
|
+
client = AsyncSharedTensorClient()
|
|
210
|
+
task_id = client.submit("build_delayed_model", delay=0.1)
|
|
211
|
+
model = client.wait_for_task(task_id, timeout=30)
|
|
212
|
+
```
|
|
213
|
+
|
|
214
|
+
Use `SharedTensorProvider(execution="task")` for task-backed endpoints.
|
|
215
|
+
Use `AsyncSharedTensorClient` when you want a task-oriented waiting interface.
|
|
216
|
+
|
|
188
217
|
## Example: Reusable Model Registry
|
|
189
218
|
|
|
190
219
|
See [examples/model_service.py](./examples/model_service.py).
|
|
@@ -295,11 +324,47 @@ handle.release()
|
|
|
295
324
|
```
|
|
296
325
|
|
|
297
326
|
Use managed mode for cached models or other reusable long-lived CUDA objects.
|
|
327
|
+
Managed object introspection now includes `created_at` and `last_accessed_at` timestamps through `get_object_info()`.
|
|
328
|
+
|
|
329
|
+
## Cache Invalidation
|
|
330
|
+
|
|
331
|
+
The library now exposes explicit cache invalidation instead of forcing process restarts when a cached object becomes stale.
|
|
332
|
+
|
|
333
|
+
```python
|
|
334
|
+
provider.invalidate_call_cache("load_model", hidden_size=4096)
|
|
335
|
+
provider.invalidate_endpoint_cache("load_model")
|
|
336
|
+
```
|
|
337
|
+
|
|
338
|
+
Client-side equivalents are also available:
|
|
339
|
+
|
|
340
|
+
```python
|
|
341
|
+
client.invalidate_call_cache("load_model", hidden_size=4096)
|
|
342
|
+
client.invalidate_endpoint_cache("load_model")
|
|
343
|
+
```
|
|
344
|
+
|
|
345
|
+
Use call-level invalidation when you want to evict one cache key.
|
|
346
|
+
Use endpoint-level invalidation when you want to drop all cached variants for the endpoint.
|
|
347
|
+
Invalidation removes cache lookup entries; it does not guarantee that already-issued client handles remain valid after producer death.
|
|
348
|
+
|
|
349
|
+
|
|
350
|
+
## Handle Health Checks
|
|
351
|
+
|
|
352
|
+
Managed handles now carry the producer `server_id` and support lightweight liveness probes:
|
|
353
|
+
|
|
354
|
+
```python
|
|
355
|
+
handle = client.call("load_model", hidden_size=4096)
|
|
356
|
+
info = handle.get_object_info()
|
|
357
|
+
client.ensure_handle_live(handle)
|
|
358
|
+
```
|
|
359
|
+
|
|
360
|
+
If the producer no longer owns the object, `client.ensure_handle_live(handle)` raises `SharedTensorStaleHandleError`.
|
|
361
|
+
This is still advisory, not a durability guarantee: it helps detect stale handles earlier, but it cannot make producer death safe.
|
|
298
362
|
|
|
299
363
|
## Runtime Introspection
|
|
300
364
|
|
|
301
|
-
`client.get_server_info()` now returns readiness and process metadata in addition to endpoint and capability data.
|
|
365
|
+
`client.get_server_info()` now returns readiness, stable `server_id`, cache/task counters, and process metadata in addition to endpoint and capability data.
|
|
302
366
|
In client mode, `provider.get_runtime_info()` wraps that into a provider-oriented view.
|
|
367
|
+
`AsyncSharedTensorClient` exposes the same runtime, cache invalidation, release, and handle-health helper methods as `SharedTensorClient`; the async surface is task-oriented, not capability-reduced.
|
|
303
368
|
|
|
304
369
|
```python
|
|
305
370
|
info = provider.get_runtime_info()
|
|
@@ -12,7 +12,7 @@ Supported:
|
|
|
12
12
|
- explicit endpoint registration
|
|
13
13
|
- sync `call` and task-backed `submit`
|
|
14
14
|
- managed object handles with explicit release
|
|
15
|
-
- server-side caching, `cache_format_key`, and
|
|
15
|
+
- server-side caching, `cache_format_key`, singleflight, and explicit cache invalidation
|
|
16
16
|
- manual two-process deployment as the primary production path
|
|
17
17
|
- zero-branch auto mode gated by `SHARED_TENSOR_ENABLED=1`
|
|
18
18
|
|
|
@@ -55,6 +55,8 @@ Production should prefer two explicitly started processes: one server process th
|
|
|
55
55
|
|
|
56
56
|
See [examples/model_service.py](./examples/model_service.py) for endpoint definitions.
|
|
57
57
|
|
|
58
|
+
The server-oriented example modules construct providers with explicit `execution_mode="server"` so importing the module already reflects the intended deployment role.
|
|
59
|
+
|
|
58
60
|
Server process:
|
|
59
61
|
|
|
60
62
|
```python
|
|
@@ -101,6 +103,7 @@ manages cache and refcounts releases managed handles explicitly
|
|
|
101
103
|
|
|
102
104
|
Core assumption:
|
|
103
105
|
- the server process that owns the original CUDA allocation must stay alive while clients are still using reopened CUDA tensors or modules
|
|
106
|
+
- handle health checks can detect some stale-object conditions, but they do not remove the producer-liveness requirement
|
|
104
107
|
|
|
105
108
|
If the server exits, crashes, or is killed before the client is done with the shared CUDA object, behavior is no longer guaranteed by this library. Depending on PyTorch and CUDA runtime state, the client may see CUDA runtime errors, invalid resource handle failures, broken module execution, or process-level instability.
|
|
106
109
|
|
|
@@ -115,6 +118,11 @@ Treat producer liveness as a hard requirement, not a soft optimization.
|
|
|
115
118
|
|
|
116
119
|
See [examples/zero_branch_env.py](./examples/zero_branch_env.py). This is a convenience mode for environments that want one file and environment-controlled behavior.
|
|
117
120
|
|
|
121
|
+
Resolution rule:
|
|
122
|
+
- `SHARED_TENSOR_ENABLED` unset or false: provider stays local
|
|
123
|
+
- `SHARED_TENSOR_ENABLED=1` and `SHARED_TENSOR_ROLE=server`: provider resolves to server and auto-starts the thread-backed local server
|
|
124
|
+
- `SHARED_TENSOR_ENABLED=1` and role unset or `client`: provider resolves to client
|
|
125
|
+
|
|
118
126
|
```bash
|
|
119
127
|
SHARED_TENSOR_ENABLED=1 SHARED_TENSOR_ROLE=server python demo.py
|
|
120
128
|
SHARED_TENSOR_ENABLED=1 python demo.py
|
|
@@ -132,6 +140,27 @@ shared function runs locally shared function becomes RPC call
|
|
|
132
140
|
CUDA object stays on same GPU CUDA object is reopened via torch IPC
|
|
133
141
|
```
|
|
134
142
|
|
|
143
|
+
## Example: Task Submission And Wait
|
|
144
|
+
|
|
145
|
+
See [examples/async_service.py](./examples/async_service.py).
|
|
146
|
+
|
|
147
|
+
```python
|
|
148
|
+
from shared_tensor import AsyncSharedTensorClient, SharedTensorProvider
|
|
149
|
+
|
|
150
|
+
provider = SharedTensorProvider(execution_mode="server")
|
|
151
|
+
|
|
152
|
+
@provider.share(execution="task")
|
|
153
|
+
def build_delayed_model(delay: float = 0.1):
|
|
154
|
+
...
|
|
155
|
+
|
|
156
|
+
client = AsyncSharedTensorClient()
|
|
157
|
+
task_id = client.submit("build_delayed_model", delay=0.1)
|
|
158
|
+
model = client.wait_for_task(task_id, timeout=30)
|
|
159
|
+
```
|
|
160
|
+
|
|
161
|
+
Use `SharedTensorProvider(execution="task")` for task-backed endpoints.
|
|
162
|
+
Use `AsyncSharedTensorClient` when you want a task-oriented waiting interface.
|
|
163
|
+
|
|
135
164
|
## Example: Reusable Model Registry
|
|
136
165
|
|
|
137
166
|
See [examples/model_service.py](./examples/model_service.py).
|
|
@@ -242,11 +271,47 @@ handle.release()
|
|
|
242
271
|
```
|
|
243
272
|
|
|
244
273
|
Use managed mode for cached models or other reusable long-lived CUDA objects.
|
|
274
|
+
Managed object introspection now includes `created_at` and `last_accessed_at` timestamps through `get_object_info()`.
|
|
275
|
+
|
|
276
|
+
## Cache Invalidation
|
|
277
|
+
|
|
278
|
+
The library now exposes explicit cache invalidation instead of forcing process restarts when a cached object becomes stale.
|
|
279
|
+
|
|
280
|
+
```python
|
|
281
|
+
provider.invalidate_call_cache("load_model", hidden_size=4096)
|
|
282
|
+
provider.invalidate_endpoint_cache("load_model")
|
|
283
|
+
```
|
|
284
|
+
|
|
285
|
+
Client-side equivalents are also available:
|
|
286
|
+
|
|
287
|
+
```python
|
|
288
|
+
client.invalidate_call_cache("load_model", hidden_size=4096)
|
|
289
|
+
client.invalidate_endpoint_cache("load_model")
|
|
290
|
+
```
|
|
291
|
+
|
|
292
|
+
Use call-level invalidation when you want to evict one cache key.
|
|
293
|
+
Use endpoint-level invalidation when you want to drop all cached variants for the endpoint.
|
|
294
|
+
Invalidation removes cache lookup entries; it does not guarantee that already-issued client handles remain valid after producer death.
|
|
295
|
+
|
|
296
|
+
|
|
297
|
+
## Handle Health Checks
|
|
298
|
+
|
|
299
|
+
Managed handles now carry the producer `server_id` and support lightweight liveness probes:
|
|
300
|
+
|
|
301
|
+
```python
|
|
302
|
+
handle = client.call("load_model", hidden_size=4096)
|
|
303
|
+
info = handle.get_object_info()
|
|
304
|
+
client.ensure_handle_live(handle)
|
|
305
|
+
```
|
|
306
|
+
|
|
307
|
+
If the producer no longer owns the object, `client.ensure_handle_live(handle)` raises `SharedTensorStaleHandleError`.
|
|
308
|
+
This is still advisory, not a durability guarantee: it helps detect stale handles earlier, but it cannot make producer death safe.
|
|
245
309
|
|
|
246
310
|
## Runtime Introspection
|
|
247
311
|
|
|
248
|
-
`client.get_server_info()` now returns readiness and process metadata in addition to endpoint and capability data.
|
|
312
|
+
`client.get_server_info()` now returns readiness, stable `server_id`, cache/task counters, and process metadata in addition to endpoint and capability data.
|
|
249
313
|
In client mode, `provider.get_runtime_info()` wraps that into a provider-oriented view.
|
|
314
|
+
`AsyncSharedTensorClient` exposes the same runtime, cache invalidation, release, and handle-health helper methods as `SharedTensorClient`; the async surface is task-oriented, not capability-reduced.
|
|
250
315
|
|
|
251
316
|
```python
|
|
252
317
|
info = provider.get_runtime_info()
|
|
@@ -4,10 +4,10 @@ build-backend = "setuptools.build_meta"
|
|
|
4
4
|
|
|
5
5
|
[project]
|
|
6
6
|
name = "shared-tensor"
|
|
7
|
-
version = "0.2.
|
|
7
|
+
version = "0.2.11"
|
|
8
8
|
description = "Native PyTorch CUDA IPC over Unix Domain Socket for same-host process separation"
|
|
9
9
|
readme = "README.md"
|
|
10
|
-
license = "
|
|
10
|
+
license = "MIT"
|
|
11
11
|
authors = [
|
|
12
12
|
{name = "Athena Team", email = "contact@world-sim-dev.org"}
|
|
13
13
|
]
|
|
@@ -1,22 +1,22 @@
|
|
|
1
1
|
"""shared_tensor: same-host same-GPU PyTorch CUDA IPC over local UDS RPC."""
|
|
2
2
|
|
|
3
3
|
from shared_tensor.async_client import AsyncSharedTensorClient
|
|
4
|
-
from shared_tensor.async_provider import AsyncSharedTensorProvider
|
|
5
4
|
from shared_tensor.async_task import TaskInfo, TaskStatus
|
|
6
5
|
from shared_tensor.client import SharedTensorClient
|
|
6
|
+
from shared_tensor.errors import SharedTensorStaleHandleError
|
|
7
7
|
from shared_tensor.managed_object import SharedObjectHandle
|
|
8
8
|
from shared_tensor.provider import SharedTensorProvider
|
|
9
9
|
from shared_tensor.server import SharedTensorServer
|
|
10
10
|
|
|
11
11
|
__all__ = [
|
|
12
12
|
"AsyncSharedTensorClient",
|
|
13
|
-
"AsyncSharedTensorProvider",
|
|
14
13
|
"SharedTensorClient",
|
|
15
14
|
"SharedObjectHandle",
|
|
15
|
+
"SharedTensorStaleHandleError",
|
|
16
16
|
"SharedTensorProvider",
|
|
17
17
|
"SharedTensorServer",
|
|
18
18
|
"TaskInfo",
|
|
19
19
|
"TaskStatus",
|
|
20
20
|
]
|
|
21
21
|
|
|
22
|
-
__version__ = "0.2.
|
|
22
|
+
__version__ = "0.2.11"
|
|
@@ -10,6 +10,7 @@ from typing import Any, cast
|
|
|
10
10
|
from shared_tensor.async_task import TaskInfo, TaskStatus
|
|
11
11
|
from shared_tensor.client import SharedTensorClient
|
|
12
12
|
from shared_tensor.errors import SharedTensorRemoteError, SharedTensorTaskError
|
|
13
|
+
from shared_tensor.managed_object import SharedObjectHandle
|
|
13
14
|
|
|
14
15
|
|
|
15
16
|
logger = logging.getLogger(__name__)
|
|
@@ -51,6 +52,38 @@ class AsyncSharedTensorClient:
|
|
|
51
52
|
def get_task_result(self, task_id: str) -> Any:
|
|
52
53
|
return self.result(task_id)
|
|
53
54
|
|
|
55
|
+
def ping(self) -> bool:
|
|
56
|
+
return self._client.ping()
|
|
57
|
+
|
|
58
|
+
def get_server_info(self) -> dict[str, Any]:
|
|
59
|
+
return self._client.get_server_info()
|
|
60
|
+
|
|
61
|
+
def list_endpoints(self) -> dict[str, Any]:
|
|
62
|
+
return self._client.list_endpoints()
|
|
63
|
+
|
|
64
|
+
def release(self, object_id: str) -> bool:
|
|
65
|
+
return self._client.release(object_id)
|
|
66
|
+
|
|
67
|
+
def release_many(self, object_ids: list[str]) -> dict[str, bool]:
|
|
68
|
+
return self._client.release_many(object_ids)
|
|
69
|
+
|
|
70
|
+
def get_object_info(self, object_id: str) -> dict[str, Any] | None:
|
|
71
|
+
return self._client.get_object_info(object_id)
|
|
72
|
+
|
|
73
|
+
def ensure_handle_live(
|
|
74
|
+
self,
|
|
75
|
+
handle: SharedObjectHandle[Any],
|
|
76
|
+
*,
|
|
77
|
+
refresh: bool = True,
|
|
78
|
+
) -> dict[str, Any]:
|
|
79
|
+
return self._client.ensure_handle_live(handle, refresh=refresh)
|
|
80
|
+
|
|
81
|
+
def invalidate_call_cache(self, endpoint: str, *args: Any, **kwargs: Any) -> bool:
|
|
82
|
+
return self._client.invalidate_call_cache(endpoint, *args, **kwargs)
|
|
83
|
+
|
|
84
|
+
def invalidate_endpoint_cache(self, endpoint: str) -> int:
|
|
85
|
+
return self._client.invalidate_endpoint_cache(endpoint)
|
|
86
|
+
|
|
54
87
|
def wait(
|
|
55
88
|
self,
|
|
56
89
|
task_id: str,
|
|
@@ -7,6 +7,7 @@ import socket
|
|
|
7
7
|
from dataclasses import dataclass
|
|
8
8
|
from typing import Any, cast
|
|
9
9
|
|
|
10
|
+
from shared_tensor.async_task import TaskStatus
|
|
10
11
|
from shared_tensor.errors import (
|
|
11
12
|
SharedTensorCapabilityError,
|
|
12
13
|
SharedTensorClientError,
|
|
@@ -16,12 +17,12 @@ from shared_tensor.errors import (
|
|
|
16
17
|
SharedTensorProtocolError,
|
|
17
18
|
SharedTensorRemoteError,
|
|
18
19
|
SharedTensorSerializationError,
|
|
20
|
+
SharedTensorStaleHandleError,
|
|
19
21
|
SharedTensorTaskError,
|
|
20
22
|
)
|
|
21
23
|
from shared_tensor.managed_object import ReleaseHandle, SharedObjectHandle
|
|
22
24
|
from shared_tensor.runtime import get_local_server
|
|
23
25
|
from shared_tensor.transport import recv_message, send_message
|
|
24
|
-
from shared_tensor.async_task import TaskStatus
|
|
25
26
|
from shared_tensor.utils import (
|
|
26
27
|
deserialize_payload,
|
|
27
28
|
resolve_runtime_socket_path,
|
|
@@ -29,7 +30,6 @@ from shared_tensor.utils import (
|
|
|
29
30
|
validate_payload_for_transport,
|
|
30
31
|
)
|
|
31
32
|
|
|
32
|
-
|
|
33
33
|
logger = logging.getLogger(__name__)
|
|
34
34
|
|
|
35
35
|
|
|
@@ -41,6 +41,9 @@ class _ClientReleaser(ReleaseHandle):
|
|
|
41
41
|
def release(self) -> bool:
|
|
42
42
|
return self.client.release(self.object_id)
|
|
43
43
|
|
|
44
|
+
def get_object_info(self) -> dict[str, Any] | None:
|
|
45
|
+
return self.client.get_object_info(self.object_id)
|
|
46
|
+
|
|
44
47
|
|
|
45
48
|
class SharedTensorClient:
|
|
46
49
|
"""UDS client for endpoint-oriented local RPC execution."""
|
|
@@ -76,6 +79,8 @@ class SharedTensorClient:
|
|
|
76
79
|
code = 5
|
|
77
80
|
elif isinstance(exc, SharedTensorConfigurationError):
|
|
78
81
|
code = 6
|
|
82
|
+
elif isinstance(exc, SharedTensorStaleHandleError):
|
|
83
|
+
code = 8
|
|
79
84
|
else:
|
|
80
85
|
code = 7
|
|
81
86
|
return SharedTensorRemoteError(
|
|
@@ -105,6 +110,7 @@ class SharedTensorClient:
|
|
|
105
110
|
object_id=cast(str, object_id),
|
|
106
111
|
value=value,
|
|
107
112
|
_releaser=_ClientReleaser(client=self, object_id=cast(str, object_id)),
|
|
113
|
+
server_id=self._infer_server_id(),
|
|
108
114
|
)
|
|
109
115
|
|
|
110
116
|
def _send_request(self, request: dict[str, Any]) -> Any:
|
|
@@ -158,6 +164,15 @@ class SharedTensorClient:
|
|
|
158
164
|
def _request(self, method: str, params: dict[str, Any] | None = None) -> Any:
|
|
159
165
|
return self._send_request({"method": method, "params": params or {}})
|
|
160
166
|
|
|
167
|
+
def _infer_server_id(self) -> str | None:
|
|
168
|
+
local_server = self._local_server()
|
|
169
|
+
if local_server is not None:
|
|
170
|
+
return cast(str | None, getattr(local_server, "server_id", None))
|
|
171
|
+
try:
|
|
172
|
+
return cast(str | None, self.get_server_info().get("server_id"))
|
|
173
|
+
except (SharedTensorClientError, SharedTensorRemoteError, SharedTensorProtocolError):
|
|
174
|
+
return None
|
|
175
|
+
|
|
161
176
|
def call(self, endpoint: str, *args: Any, **kwargs: Any) -> Any:
|
|
162
177
|
if self.verbose_debug:
|
|
163
178
|
logger.debug("Client calling endpoint", extra={"endpoint": endpoint})
|
|
@@ -245,6 +260,25 @@ class SharedTensorClient:
|
|
|
245
260
|
result = self._request("get_object_info", {"object_id": object_id})
|
|
246
261
|
return cast(dict[str, Any] | None, result.get("object"))
|
|
247
262
|
|
|
263
|
+
def ensure_handle_live(self, handle: SharedObjectHandle[Any], *, refresh: bool = True) -> dict[str, Any]:
|
|
264
|
+
info = handle.get_object_info(refresh=refresh)
|
|
265
|
+
if info is None:
|
|
266
|
+
raise SharedTensorStaleHandleError(
|
|
267
|
+
f"Managed object '{handle.object_id}' is no longer registered on the producer",
|
|
268
|
+
object_id=handle.object_id,
|
|
269
|
+
server_id=handle.server_id,
|
|
270
|
+
reason="object_missing",
|
|
271
|
+
)
|
|
272
|
+
observed_server_id = cast(str | None, info.get("server_id"))
|
|
273
|
+
if handle.server_id is not None and observed_server_id is not None and observed_server_id != handle.server_id:
|
|
274
|
+
raise SharedTensorStaleHandleError(
|
|
275
|
+
f"Managed object '{handle.object_id}' belongs to server '{handle.server_id}' but producer now reports '{observed_server_id}'",
|
|
276
|
+
object_id=handle.object_id,
|
|
277
|
+
server_id=handle.server_id,
|
|
278
|
+
reason="server_mismatch",
|
|
279
|
+
)
|
|
280
|
+
return info
|
|
281
|
+
|
|
248
282
|
def ping(self) -> bool:
|
|
249
283
|
if self._local_server() is not None:
|
|
250
284
|
return True
|
|
@@ -260,6 +294,31 @@ class SharedTensorClient:
|
|
|
260
294
|
return self._run_local(lambda: cast(dict[str, Any], local_server._get_server_info()))
|
|
261
295
|
return cast(dict[str, Any], self._request("get_server_info"))
|
|
262
296
|
|
|
297
|
+
def invalidate_call_cache(self, endpoint: str, *args: Any, **kwargs: Any) -> bool:
|
|
298
|
+
local_server = self._local_server()
|
|
299
|
+
if local_server is not None:
|
|
300
|
+
return self._run_local(
|
|
301
|
+
lambda: bool(local_server.invalidate_call_cache(endpoint, args=tuple(args), kwargs=dict(kwargs)))
|
|
302
|
+
)
|
|
303
|
+
encoding, args_payload, kwargs_payload = serialize_call_payloads(tuple(args), dict(kwargs))
|
|
304
|
+
result = self._request(
|
|
305
|
+
"invalidate_call_cache",
|
|
306
|
+
{
|
|
307
|
+
"endpoint": endpoint,
|
|
308
|
+
"args_bytes": args_payload,
|
|
309
|
+
"kwargs_bytes": kwargs_payload,
|
|
310
|
+
"encoding": encoding,
|
|
311
|
+
},
|
|
312
|
+
)
|
|
313
|
+
return bool(result["invalidated"])
|
|
314
|
+
|
|
315
|
+
def invalidate_endpoint_cache(self, endpoint: str) -> int:
|
|
316
|
+
local_server = self._local_server()
|
|
317
|
+
if local_server is not None:
|
|
318
|
+
return self._run_local(lambda: int(local_server.invalidate_endpoint_cache(endpoint)))
|
|
319
|
+
result = self._request("invalidate_endpoint_cache", {"endpoint": endpoint})
|
|
320
|
+
return int(result["invalidated"])
|
|
321
|
+
|
|
263
322
|
def list_endpoints(self) -> dict[str, Any]:
|
|
264
323
|
local_server = self._local_server()
|
|
265
324
|
if local_server is not None:
|
|
@@ -344,4 +403,5 @@ class SharedTensorClient:
|
|
|
344
403
|
object_id=cast(str, object_id),
|
|
345
404
|
value=value,
|
|
346
405
|
_releaser=_ClientReleaser(client=self, object_id=cast(str, object_id)),
|
|
406
|
+
server_id=cast(str | None, result.get("server_id")),
|
|
347
407
|
)
|
|
@@ -15,6 +15,7 @@ __all__ = [
|
|
|
15
15
|
"SharedTensorClientError",
|
|
16
16
|
"SharedTensorServerError",
|
|
17
17
|
"SharedTensorProviderError",
|
|
18
|
+
"SharedTensorStaleHandleError",
|
|
18
19
|
]
|
|
19
20
|
|
|
20
21
|
|
|
@@ -69,3 +70,20 @@ class SharedTensorServerError(SharedTensorError):
|
|
|
69
70
|
|
|
70
71
|
class SharedTensorProviderError(SharedTensorError):
|
|
71
72
|
"""Raised for provider registration or invocation problems."""
|
|
73
|
+
|
|
74
|
+
|
|
75
|
+
class SharedTensorStaleHandleError(SharedTensorError):
|
|
76
|
+
"""Raised when a managed handle can no longer be trusted."""
|
|
77
|
+
|
|
78
|
+
def __init__(
|
|
79
|
+
self,
|
|
80
|
+
message: str,
|
|
81
|
+
*,
|
|
82
|
+
object_id: str | None = None,
|
|
83
|
+
server_id: str | None = None,
|
|
84
|
+
reason: str | None = None,
|
|
85
|
+
) -> None:
|
|
86
|
+
super().__init__(message)
|
|
87
|
+
self.object_id = object_id
|
|
88
|
+
self.server_id = server_id
|
|
89
|
+
self.reason = reason
|
|
@@ -2,8 +2,9 @@
|
|
|
2
2
|
|
|
3
3
|
from __future__ import annotations
|
|
4
4
|
|
|
5
|
+
import time
|
|
5
6
|
import uuid
|
|
6
|
-
from dataclasses import dataclass
|
|
7
|
+
from dataclasses import dataclass, field
|
|
7
8
|
from threading import RLock
|
|
8
9
|
from typing import Any, Generic, TypeVar
|
|
9
10
|
|
|
@@ -17,6 +18,8 @@ class ManagedObjectEntry:
|
|
|
17
18
|
endpoint: str
|
|
18
19
|
cache_key: str | None
|
|
19
20
|
refcount: int = 1
|
|
21
|
+
created_at: float = 0.0
|
|
22
|
+
last_accessed_at: float = 0.0
|
|
20
23
|
|
|
21
24
|
|
|
22
25
|
@dataclass(slots=True)
|
|
@@ -42,15 +45,19 @@ class ManagedObjectRegistry:
|
|
|
42
45
|
if entry is None:
|
|
43
46
|
self._cache_index.pop(cache_key, None)
|
|
44
47
|
return None
|
|
48
|
+
entry.last_accessed_at = time.time()
|
|
45
49
|
return entry
|
|
46
50
|
|
|
47
51
|
def register(self, *, endpoint: str, value: Any, cache_key: str | None) -> ManagedObjectEntry:
|
|
48
52
|
with self._lock:
|
|
53
|
+
now = time.time()
|
|
49
54
|
entry = ManagedObjectEntry(
|
|
50
55
|
object_id=uuid.uuid4().hex,
|
|
51
56
|
value=value,
|
|
52
57
|
endpoint=endpoint,
|
|
53
58
|
cache_key=cache_key,
|
|
59
|
+
created_at=now,
|
|
60
|
+
last_accessed_at=now,
|
|
54
61
|
)
|
|
55
62
|
self._entries[entry.object_id] = entry
|
|
56
63
|
if cache_key is not None:
|
|
@@ -67,6 +74,7 @@ class ManagedObjectRegistry:
|
|
|
67
74
|
if entry is None:
|
|
68
75
|
return None
|
|
69
76
|
entry.refcount += 1
|
|
77
|
+
entry.last_accessed_at = time.time()
|
|
70
78
|
return entry
|
|
71
79
|
|
|
72
80
|
def release(self, object_id: str) -> ManagedReleaseResult:
|
|
@@ -100,6 +108,31 @@ class ManagedObjectRegistry:
|
|
|
100
108
|
"endpoint": entry.endpoint,
|
|
101
109
|
"cache_key": entry.cache_key,
|
|
102
110
|
"refcount": entry.refcount,
|
|
111
|
+
"created_at": entry.created_at,
|
|
112
|
+
"last_accessed_at": entry.last_accessed_at,
|
|
113
|
+
}
|
|
114
|
+
|
|
115
|
+
def invalidate_cache_key(self, cache_key: str) -> bool:
|
|
116
|
+
with self._lock:
|
|
117
|
+
object_id = self._cache_index.pop(cache_key, None)
|
|
118
|
+
return object_id is not None
|
|
119
|
+
|
|
120
|
+
def invalidate_endpoint(self, endpoint: str) -> int:
|
|
121
|
+
with self._lock:
|
|
122
|
+
keys = [
|
|
123
|
+
cache_key
|
|
124
|
+
for cache_key, object_id in self._cache_index.items()
|
|
125
|
+
if (entry := self._entries.get(object_id)) is not None and entry.endpoint == endpoint
|
|
126
|
+
]
|
|
127
|
+
for cache_key in keys:
|
|
128
|
+
self._cache_index.pop(cache_key, None)
|
|
129
|
+
return len(keys)
|
|
130
|
+
|
|
131
|
+
def stats(self) -> dict[str, int]:
|
|
132
|
+
with self._lock:
|
|
133
|
+
return {
|
|
134
|
+
"objects": len(self._entries),
|
|
135
|
+
"cached_objects": len(self._cache_index),
|
|
103
136
|
}
|
|
104
137
|
|
|
105
138
|
def clear(self) -> None:
|
|
@@ -112,6 +145,9 @@ class ReleaseHandle:
|
|
|
112
145
|
def release(self) -> bool: # pragma: no cover - protocol surface only
|
|
113
146
|
raise NotImplementedError
|
|
114
147
|
|
|
148
|
+
def get_object_info(self) -> dict[str, Any] | None: # pragma: no cover - protocol surface only
|
|
149
|
+
raise NotImplementedError
|
|
150
|
+
|
|
115
151
|
|
|
116
152
|
@dataclass(slots=True)
|
|
117
153
|
class SharedObjectHandle(Generic[T]):
|
|
@@ -119,6 +155,8 @@ class SharedObjectHandle(Generic[T]):
|
|
|
119
155
|
value: T
|
|
120
156
|
_releaser: ReleaseHandle
|
|
121
157
|
released: bool = False
|
|
158
|
+
server_id: str | None = None
|
|
159
|
+
_metadata_cache: dict[str, Any] | None = field(default=None, init=False, repr=False)
|
|
122
160
|
|
|
123
161
|
def release(self) -> bool:
|
|
124
162
|
if self.released:
|
|
@@ -126,8 +164,19 @@ class SharedObjectHandle(Generic[T]):
|
|
|
126
164
|
released = self._releaser.release()
|
|
127
165
|
if released:
|
|
128
166
|
self.released = True
|
|
167
|
+
self._metadata_cache = None
|
|
129
168
|
return released
|
|
130
169
|
|
|
170
|
+
def get_object_info(self, *, refresh: bool = False) -> dict[str, Any] | None:
|
|
171
|
+
if self.released:
|
|
172
|
+
return None
|
|
173
|
+
if self._metadata_cache is None or refresh:
|
|
174
|
+
self._metadata_cache = self._releaser.get_object_info()
|
|
175
|
+
return None if self._metadata_cache is None else dict(self._metadata_cache)
|
|
176
|
+
|
|
177
|
+
def is_stale(self) -> bool:
|
|
178
|
+
return self.get_object_info(refresh=True) is None
|
|
179
|
+
|
|
131
180
|
def __enter__(self) -> SharedObjectHandle[T]:
|
|
132
181
|
return self
|
|
133
182
|
|