shared-tensor 0.2.6__tar.gz → 0.2.7__tar.gz

{shared_tensor-0.2.6 → shared_tensor-0.2.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: shared-tensor
-Version: 0.2.6
+Version: 0.2.7
 Summary: Native PyTorch CUDA IPC over Unix Domain Socket for same-host process separation
 Author-email: Athena Team <contact@world-sim-dev.org>
 Maintainer-email: Athena Team <contact@world-sim-dev.org>
@@ -16,18 +16,20 @@ Classifier: Intended Audience :: Developers
 Classifier: Intended Audience :: Science/Research
 Classifier: Operating System :: POSIX :: Linux
 Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Classifier: Topic :: System :: Distributed Computing
-Requires-Python: >=3.10
+Requires-Python: <3.14,>=3.9
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: cloudpickle>=3.0.0
 Requires-Dist: numpy<2
-Requires-Dist: torch>=2.2.0
+Requires-Dist: torch<2.8,>=2.1
 Provides-Extra: dev
 Requires-Dist: pytest>=6.0; extra == "dev"
 Requires-Dist: pytest-cov>=2.0; extra == "dev"

{shared_tensor-0.2.6 → shared_tensor-0.2.7}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "shared-tensor"
-version = "0.2.6"
+version = "0.2.7"
 description = "Native PyTorch CUDA IPC over Unix Domain Socket for same-host process separation"
 readme = "README.md"
 license = "Apache-2.0"
@@ -33,18 +33,20 @@ classifiers = [
     "Intended Audience :: Science/Research",
     "Operating System :: POSIX :: Linux",
     "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.9",
     "Programming Language :: Python :: 3.10",
     "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
     "Topic :: Scientific/Engineering :: Artificial Intelligence",
     "Topic :: Software Development :: Libraries :: Python Modules",
     "Topic :: System :: Distributed Computing",
 ]
-requires-python = ">=3.10"
+requires-python = ">=3.9,<3.14"
 dependencies = [
     "cloudpickle>=3.0.0",
     "numpy<2",
-    "torch>=2.2.0",
+    "torch>=2.1,<2.8",
 ]
 
 [project.optional-dependencies]
@@ -89,7 +91,7 @@ shared_tensor = ["*.so", "*.dll", "*.dylib"]
 
 [tool.black]
 line-length = 88
-target-version = ['py310', 'py311', 'py312']
+target-version = ['py39', 'py310', 'py311', 'py312', 'py313']
 include = '\.pyi?$'
 extend-exclude = '''
 /(
@@ -115,7 +117,7 @@ use_parentheses = true
 ensure_newline_before_comments = true
 
 [tool.mypy]
-python_version = "3.10"
+python_version = "3.9"
 warn_return_any = true
 warn_unused_configs = true
 disallow_untyped_defs = true
@@ -180,7 +182,7 @@ exclude_lines = [
 ]
 
 [tool.ruff]
-target-version = "py310"
+target-version = "py39"
 line-length = 88
 
 [tool.ruff.lint]

{shared_tensor-0.2.6 → shared_tensor-0.2.7}/shared_tensor/__init__.py

@@ -19,4 +19,4 @@ __all__ = [
     "TaskStatus",
 ]
 
-__version__ = "0.2.6"
+__version__ = "0.2.7"

{shared_tensor-0.2.6 → shared_tensor-0.2.7}/shared_tensor/async_task.py

@@ -66,6 +66,7 @@ class TaskInfo:
 class _TaskEntry:
     info: TaskInfo
     future: Future[Any]
+    local_result: Any = None
 
 
 class TaskManager:
@@ -139,6 +140,8 @@ class TaskManager:
             )
             return
 
+        self._store_local_result(task_id, result)
+
         if result is None:
             self._transition(
                 task_id,
@@ -191,6 +194,13 @@ class TaskManager:
             for key, value in updates.items():
                 setattr(entry.info, key, value)
 
+    def _store_local_result(self, task_id: str, value: Any) -> None:
+        with self._lock:
+            entry = self._tasks.get(task_id)
+            if entry is None:
+                return
+            entry.local_result = value
+
     def get(self, task_id: str) -> TaskInfo:
         self._maybe_cleanup()
         with self._lock:
@@ -207,6 +217,24 @@ class TaskManager:
             return None
         return deserialize_payload(encoding, payload_bytes)
 
+    def result_local(self, task_id: str) -> Any:
+        self._maybe_cleanup()
+        with self._lock:
+            entry = self._tasks.get(task_id)
+            if entry is None:
+                raise SharedTensorTaskError(f"Task '{task_id}' was not found")
+            info = copy.deepcopy(entry.info)
+            value = entry.local_result
+        if info.status == TaskStatus.CANCELLED:
+            raise SharedTensorTaskError(f"Task '{task_id}' was cancelled")
+        if info.status == TaskStatus.FAILED:
+            raise SharedTensorTaskError(info.error_message or f"Task '{task_id}' failed")
+        if info.status != TaskStatus.COMPLETED:
+            raise SharedTensorTaskError(
+                f"Task '{task_id}' is not complete; current status is '{info.status.value}'"
+            )
+        return value
+
     def wait_result_payload(
         self,
         task_id: str,
@@ -242,6 +270,21 @@ class TaskManager:
             "object_id": info.metadata.get("object_id"),
         }
 
+    def wait_result_local(self, task_id: str, timeout: float | None = None) -> Any:
+        self._maybe_cleanup()
+        with self._lock:
+            entry = self._tasks.get(task_id)
+            if entry is None:
+                raise SharedTensorTaskError(f"Task '{task_id}' was not found")
+            future = entry.future
+        try:
+            future.result(timeout=timeout)
+        except FutureTimeoutError as exc:
+            raise SharedTensorTaskError(
+                f"Task '{task_id}' did not complete within {timeout} seconds"
+            ) from exc
+        return self.result_local(task_id)
+
     def cancel(self, task_id: str) -> bool:
         self._maybe_cleanup()
         with self._lock:

{shared_tensor-0.2.6 → shared_tensor-0.2.7}/shared_tensor/client.py

@@ -8,16 +8,25 @@ from dataclasses import dataclass
 from typing import Any, cast
 
 from shared_tensor.errors import (
+    SharedTensorCapabilityError,
     SharedTensorClientError,
+    SharedTensorConfigurationError,
+    SharedTensorError,
+    SharedTensorProviderError,
     SharedTensorProtocolError,
     SharedTensorRemoteError,
+    SharedTensorSerializationError,
+    SharedTensorTaskError,
 )
 from shared_tensor.managed_object import ReleaseHandle, SharedObjectHandle
+from shared_tensor.runtime import get_local_server
 from shared_tensor.transport import recv_message, send_message
+from shared_tensor.async_task import TaskStatus
 from shared_tensor.utils import (
     deserialize_payload,
     resolve_runtime_socket_path,
     serialize_call_payloads,
+    validate_payload_for_transport,
 )
 
 
@@ -50,6 +59,54 @@ class SharedTensorClient:
         self.timeout = timeout
         self.verbose_debug = verbose_debug
 
+    def _local_server(self):
+        return get_local_server(self.socket_path)
+
+    @staticmethod
+    def _remote_error_from_local(exc: SharedTensorError) -> SharedTensorRemoteError:
+        if isinstance(exc, SharedTensorProtocolError):
+            code = 1
+        elif isinstance(exc, SharedTensorProviderError):
+            code = 2
+        elif isinstance(exc, SharedTensorSerializationError):
+            code = 3
+        elif isinstance(exc, SharedTensorCapabilityError):
+            code = 4
+        elif isinstance(exc, SharedTensorTaskError):
+            code = 5
+        elif isinstance(exc, SharedTensorConfigurationError):
+            code = 6
+        else:
+            code = 7
+        return SharedTensorRemoteError(
+            f"Remote error [{code}]: {exc}",
+            code=code,
+            data=None,
+            error_type=type(exc).__name__,
+        )
+
+    def _run_local(self, operation):
+        try:
+            return operation()
+        except SharedTensorError as exc:
+            raise self._remote_error_from_local(exc) from exc
+
+    def _decode_local_result(self, result: Any) -> Any:
+        if result is None:
+            return None
+        value = result.value
+        if value is None:
+            return None
+        validate_payload_for_transport(value, allow_dict_keys=isinstance(value, dict))
+        object_id = result.object_id
+        if object_id is None:
+            return value
+        return SharedObjectHandle(
+            object_id=cast(str, object_id),
+            value=value,
+            _releaser=_ClientReleaser(client=self, object_id=cast(str, object_id)),
+        )
+
     def _send_request(self, request: dict[str, Any]) -> Any:
         method = request.get("method", "<unknown>")
         if self.verbose_debug:
@@ -104,6 +161,13 @@ class SharedTensorClient:
     def call(self, endpoint: str, *args: Any, **kwargs: Any) -> Any:
         if self.verbose_debug:
             logger.debug("Client calling endpoint", extra={"endpoint": endpoint})
+        local_server = self._local_server()
+        if local_server is not None:
+            return self._run_local(
+                lambda: self._decode_local_result(
+                    local_server.call_local_client(endpoint, args=tuple(args), kwargs=dict(kwargs))
+                )
+            )
         encoding, args_payload, kwargs_payload = serialize_call_payloads(tuple(args), dict(kwargs))
         result = self._request(
             "call",
@@ -119,6 +183,19 @@ class SharedTensorClient:
     def submit(self, endpoint: str, *args: Any, **kwargs: Any) -> str:
         if self.verbose_debug:
             logger.debug("Client submitting task", extra={"endpoint": endpoint})
+        local_server = self._local_server()
+        if local_server is not None:
+            return self._run_local(
+                lambda: cast(
+                    str,
+                    local_server._submit_endpoint_task(
+                        endpoint,
+                        local_server.provider.get_endpoint(endpoint),
+                        tuple(args),
+                        dict(kwargs),
+                    ).task_id,
+                )
+            )
         encoding, args_payload, kwargs_payload = serialize_call_payloads(tuple(args), dict(kwargs))
         result = self._request(
             "submit",
@@ -134,18 +211,43 @@ class SharedTensorClient:
     def release(self, object_id: str) -> bool:
         if self.verbose_debug:
             logger.debug("Client releasing managed object", extra={"object_id": object_id})
+        local_server = self._local_server()
+        if local_server is not None:
+            return self._run_local(
+                lambda: bool(local_server._handle_release_object({"object_id": object_id})["released"])
+            )
         result = self._request("release_object", {"object_id": object_id})
         return bool(result["released"])
 
     def release_many(self, object_ids: list[str]) -> dict[str, bool]:
+        local_server = self._local_server()
+        if local_server is not None:
+            return self._run_local(
+                lambda: {
+                    object_id: bool(released)
+                    for object_id, released in local_server._handle_release_objects({"object_ids": object_ids})[
+                        "released"
+                    ].items()
+                }
+            )
         result = self._request("release_objects", {"object_ids": object_ids})
         return {object_id: bool(released) for object_id, released in result["released"].items()}
 
     def get_object_info(self, object_id: str) -> dict[str, Any] | None:
+        local_server = self._local_server()
+        if local_server is not None:
+            return self._run_local(
+                lambda: cast(
+                    dict[str, Any] | None,
+                    local_server._handle_get_object_info({"object_id": object_id}).get("object"),
+                )
+            )
         result = self._request("get_object_info", {"object_id": object_id})
         return cast(dict[str, Any] | None, result.get("object"))
 
     def ping(self) -> bool:
+        if self._local_server() is not None:
+            return True
         try:
             self._request("ping")
         except (SharedTensorClientError, SharedTensorRemoteError):
@@ -153,29 +255,66 @@ class SharedTensorClient:
         return True
 
     def get_server_info(self) -> dict[str, Any]:
+        local_server = self._local_server()
+        if local_server is not None:
+            return self._run_local(lambda: cast(dict[str, Any], local_server._get_server_info()))
         return cast(dict[str, Any], self._request("get_server_info"))
 
     def list_endpoints(self) -> dict[str, Any]:
+        local_server = self._local_server()
+        if local_server is not None:
+            return self._run_local(lambda: cast(dict[str, Any], local_server.provider.list_endpoints()))
         return cast(dict[str, Any], self._request("list_endpoints"))
 
     def get_task_status(self, task_id: str) -> dict[str, Any]:
+        local_server = self._local_server()
+        if local_server is not None:
+            return self._run_local(
+                lambda: cast(dict[str, Any], local_server._task_manager_instance().get(task_id).to_dict())
+            )
         return cast(dict[str, Any], self._request("get_task", {"task_id": task_id}))
 
     def get_task_result(self, task_id: str) -> Any:
+        local_server = self._local_server()
+        if local_server is not None:
+            return self._run_local(
+                lambda: self._decode_local_result(local_server.get_task_result_local(task_id))
+            )
         return self._decode_rpc_payload(self._request("get_task_result", {"task_id": task_id}))
 
     def wait_task(self, task_id: str, timeout: float | None = None) -> dict[str, Any]:
         if self.verbose_debug:
             logger.debug("Client waiting for task", extra={"task_id": task_id, "timeout": timeout})
+        local_server = self._local_server()
+        if local_server is not None:
+            return self._run_local(
+                lambda: cast(dict[str, Any], local_server.wait_task_local(task_id, timeout=timeout))
+            )
         params = {"task_id": task_id}
         if timeout is not None:
             params["timeout"] = timeout
         return cast(dict[str, Any], self._request("wait_task", params))
 
     def cancel_task(self, task_id: str) -> bool:
+        local_server = self._local_server()
+        if local_server is not None:
+            return self._run_local(lambda: bool(local_server._task_manager_instance().cancel(task_id)))
         return bool(self._request("cancel_task", {"task_id": task_id})["cancelled"])
 
     def list_tasks(self, status: str | None = None) -> dict[str, Any]:
+        local_server = self._local_server()
+        if local_server is not None:
+            return self._run_local(
+                lambda: cast(
+                    dict[str, Any],
+                    {
+                        listed_task_id: info.to_dict()
+                        for listed_task_id, info in local_server._task_manager_instance()
+                        .list(status=None if status is None else TaskStatus(status))
+                        .items()
+                    },
+                )
+            )
         params = {"status": status} if status else None
         return cast(dict[str, Any], self._request("list_tasks", params))
 

shared_tensor-0.2.7/shared_tensor/runtime.py

@@ -0,0 +1,30 @@
+"""In-process runtime registry for thread-backed local servers."""
+
+from __future__ import annotations
+
+from threading import RLock
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from shared_tensor.server import SharedTensorServer
+
+
+_LOCK = RLock()
+_SERVERS: dict[str, "SharedTensorServer"] = {}
+
+
+def register_local_server(socket_path: str, server: "SharedTensorServer") -> None:
+    with _LOCK:
+        _SERVERS[socket_path] = server
+
+
+def unregister_local_server(socket_path: str, server: "SharedTensorServer") -> None:
+    with _LOCK:
+        current = _SERVERS.get(socket_path)
+        if current is server:
+            _SERVERS.pop(socket_path, None)
+
+
+def get_local_server(socket_path: str) -> "SharedTensorServer | None":
+    with _LOCK:
+        return _SERVERS.get(socket_path)

{shared_tensor-0.2.6 → shared_tensor-0.2.7}/shared_tensor/server.py

@@ -22,6 +22,7 @@ from shared_tensor.errors import (
 )
 from shared_tensor.managed_object import ManagedObjectRegistry
 from shared_tensor.provider import EndpointDefinition, SharedTensorProvider
+from shared_tensor.runtime import register_local_server, unregister_local_server
 from shared_tensor.transport import recv_message, send_message
 from shared_tensor.utils import (
     CONTROL_ENCODING,
@@ -59,6 +60,12 @@ class _ServerThreadState:
     error: BaseException | None = None
 
 
+@dataclass(slots=True)
+class _EndpointResult:
+    value: Any
+    object_id: str | None = None
+
+
 class SharedTensorServer:
     def __init__(
         self,
@@ -197,22 +204,22 @@ class SharedTensorServer:
     ) -> Any:
         return self._task_manager_instance().submit(
             endpoint,
-            self._execute_endpoint_call,
+            self._execute_endpoint_result,
             (endpoint, definition, args, kwargs),
             {},
-            result_encoder=lambda payload: payload,
+            result_encoder=self._encode_endpoint_result,
         )
 
-    def _execute_endpoint_call(
+    def _execute_endpoint_result(
         self,
         endpoint: str,
         definition: EndpointDefinition,
         args: tuple[Any, ...],
         kwargs: dict[str, Any],
-    ) -> dict[str, Any]:
+    ) -> _EndpointResult:
         cache_key = self._cache_key(endpoint, definition, args, kwargs)
         if cache_key is not None:
-            cached = self._lookup_cached_result(definition, cache_key)
+            cached = self._lookup_cached_result_value(definition, cache_key)
             if cached is not None:
                 if self.verbose_debug:
                     logger.debug("Server cache hit", extra={"endpoint": endpoint, "cache_key": cache_key})
@@ -224,20 +231,15 @@ class SharedTensorServer:
             if self.verbose_debug and owner:
                 logger.debug("Server created singleflight entry", extra={"endpoint": endpoint, "cache_key": inflight_key})
             if not owner:
-                if self.verbose_debug:
-                    logger.debug("Server joined singleflight entry", extra={"endpoint": endpoint, "cache_key": inflight_key})
-                if definition.managed:
-                    payload = future.result()
-                    object_id = payload.get("object_id")
-                    if object_id is not None:
-                        self._managed_objects.add_ref(object_id)
-                    return payload
-                return future.result()
+                result = future.result()
+                if definition.managed and result.object_id is not None:
+                    self._managed_objects.add_ref(result.object_id)
+                return result
         else:
             future = None
 
         try:
-            encoded = self._run_endpoint_under_policy(endpoint, definition, args, kwargs, cache_key)
+            result = self._run_endpoint_under_policy(endpoint, definition, args, kwargs, cache_key)
         except Exception as exc:
             if future is not None:
                 future.set_exception(exc)
@@ -245,9 +247,20 @@ class SharedTensorServer:
             raise
 
         if future is not None:
-            future.set_result(encoded)
+            future.set_result(result)
             self._release_inflight(inflight_key, future)
-        return encoded
+        return result
+
+    def _execute_endpoint_call(
+        self,
+        endpoint: str,
+        definition: EndpointDefinition,
+        args: tuple[Any, ...],
+        kwargs: dict[str, Any],
+    ) -> dict[str, Any]:
+        return self._encode_endpoint_result(
+            self._execute_endpoint_result(endpoint, definition, args, kwargs)
+        )
 
     def _run_endpoint_under_policy(
         self,
@@ -256,11 +269,11 @@ class SharedTensorServer:
         args: tuple[Any, ...],
         kwargs: dict[str, Any],
         cache_key: str | None,
-    ) -> dict[str, Any]:
+    ) -> _EndpointResult:
         if definition.concurrency == "serialized":
             lock = self._endpoint_lock(endpoint)
             with lock:
-                cached = self._lookup_cached_result(definition, cache_key)
+                cached = self._lookup_cached_result_value(definition, cache_key)
                 if cached is not None:
                     return cached
                 return self._materialize_endpoint_result(endpoint, definition, args, kwargs, cache_key)
@@ -273,17 +286,15 @@ class SharedTensorServer:
         args: tuple[Any, ...],
         kwargs: dict[str, Any],
         cache_key: str | None,
-    ) -> dict[str, Any]:
+    ) -> _EndpointResult:
         if definition.managed:
             return self._materialize_managed_result(endpoint, definition, args, kwargs, cache_key)
         value = definition.func(*args, **kwargs)
         if self.verbose_debug:
             logger.debug("Server executed direct endpoint", extra={"endpoint": endpoint})
-        result = self._encode_result(value)
         if cache_key is not None:
-            self._cache[cache_key] = result
             self._local_cache[cache_key] = value
-        return result
+        return _EndpointResult(value=value)
 
     def _materialize_managed_result(
         self,
@@ -292,24 +303,24 @@ class SharedTensorServer:
         args: tuple[Any, ...],
         kwargs: dict[str, Any],
         cache_key: str | None,
-    ) -> dict[str, Any]:
+    ) -> _EndpointResult:
         if cache_key is not None:
             cached = self._managed_objects.get_cached(cache_key)
             if cached is not None:
                 self._managed_objects.add_ref(cached.object_id)
-                return self._encode_result(cached.value, object_id=cached.object_id)
+                return _EndpointResult(value=cached.value, object_id=cached.object_id)
 
         result = definition.func(*args, **kwargs)
         if self.verbose_debug:
             logger.debug("Server created managed object", extra={"endpoint": endpoint, "cache_key": cache_key})
         entry = self._managed_objects.register(endpoint=endpoint, value=result, cache_key=cache_key)
-        return self._encode_result(entry.value, object_id=entry.object_id)
+        return _EndpointResult(value=entry.value, object_id=entry.object_id)
 
-    def _lookup_cached_result(
+    def _lookup_cached_result_value(
         self,
         definition: EndpointDefinition,
         cache_key: str | None,
-    ) -> dict[str, Any] | None:
+    ) -> _EndpointResult | None:
         if cache_key is None:
             return None
         if definition.managed:
@@ -317,16 +328,52 @@ class SharedTensorServer:
             if cached is None:
                 return None
             self._managed_objects.add_ref(cached.object_id)
-            return self._encode_result(cached.value, object_id=cached.object_id)
-        cached = self._cache.get(cache_key)
-        if cached is not None:
-            return cached
+            return _EndpointResult(value=cached.value, object_id=cached.object_id)
         local_value = self._local_cache.get(cache_key)
         if local_value is None:
             return None
-        encoded = self._encode_result(local_value)
-        self._cache[cache_key] = encoded
-        return encoded
+        return _EndpointResult(value=local_value)
+
+    def call_local_client(
+        self,
+        endpoint: str,
+        *,
+        args: tuple[Any, ...] = (),
+        kwargs: dict[str, Any] | None = None,
+    ) -> _EndpointResult | None:
+        definition = self.provider.get_endpoint(endpoint)
+        resolved_kwargs = kwargs or {}
+        if definition.execution == "task":
+            task_info = self._submit_endpoint_task(endpoint, definition, args, resolved_kwargs)
+            return self.wait_task_result_local(task_info.task_id)
+        return self._execute_endpoint_result(endpoint, definition, args, resolved_kwargs)
+
+    def get_task_result_local(self, task_id: str) -> _EndpointResult | None:
+        result = self._task_manager_instance().result_local(task_id)
+        if result is None:
+            return None
+        return result
+
+    def wait_task_result_local(self, task_id: str, timeout: float | None = None) -> _EndpointResult | None:
+        result = self._task_manager_instance().wait_result_local(task_id, timeout=timeout)
+        if result is None:
+            return None
+        return result
+
+    def wait_task_local(self, task_id: str, timeout: float | None = None) -> dict[str, Any]:
+        try:
+            self._task_manager_instance().wait_result_local(task_id, timeout=timeout)
+        except SharedTensorTaskError:
+            info = self._task_manager_instance().get(task_id)
+            if info.status in {TaskStatus.PENDING, TaskStatus.RUNNING}:
+                return info.to_dict()
+            raise
+        return self._task_manager_instance().get(task_id).to_dict()
+
+    def encode_local_result(self, result: _EndpointResult | None) -> dict[str, Any]:
+        if result is None:
+            return {"encoding": None, "payload_bytes": None, "object_id": None}
+        return self._encode_endpoint_result(result)
 
     def invoke_local(
         self,
@@ -455,6 +502,9 @@ class SharedTensorServer:
         encoding, payload = serialize_payload(value)
         return {"encoding": encoding, "payload_bytes": payload, "object_id": object_id}
 
+    def _encode_endpoint_result(self, result: _EndpointResult) -> dict[str, Any]:
+        return self._encode_result(result.value, object_id=result.object_id)
+
     def _task_manager_instance(self) -> TaskManager:
         if self._task_manager is None:
             self._task_manager = TaskManager(
@@ -560,6 +610,7 @@ class SharedTensorServer:
             self.listener = listener
             self.running = True
             self.started_at = time.time()
+            register_local_server(self.socket_path, self)
             if started_event is not None:
                 started_event.set()
             while self.running:
@@ -632,6 +683,7 @@ class SharedTensorServer:
         self._local_cache.clear()
         self._inflight.clear()
         self._endpoint_locks.clear()
+        unregister_local_server(self.socket_path, self)
         unlink_socket_path(self.socket_path)
 
     def __enter__(self) -> SharedTensorServer:

{shared_tensor-0.2.6 → shared_tensor-0.2.7}/shared_tensor.egg-info/SOURCES.txt

@@ -10,6 +10,7 @@ shared_tensor/client.py
 shared_tensor/errors.py
 shared_tensor/managed_object.py
 shared_tensor/provider.py
+shared_tensor/runtime.py
 shared_tensor/server.py
 shared_tensor/transport.py
 shared_tensor/utils.py