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.
@@ -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.9
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: Apache-2.0
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 singleflight
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 singleflight
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.9"
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 = "Apache-2.0"
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.9"
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