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

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

@@ -1,13 +1,13 @@
 Metadata-Version: 2.4
 Name: shared-tensor
-Version: 0.2.7
+Version: 0.2.8
 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>
 License-Expression: Apache-2.0
 Project-URL: Homepage, https://github.com/world-sim-dev/shared-tensor
 Project-URL: Repository, https://github.com/world-sim-dev/shared-tensor
-Project-URL: Documentation, https://github.com/world-sim-dev/shared-tensor/wiki
+Project-URL: Documentation, https://github.com/world-sim-dev/shared-tensor/tree/main/docs
 Project-URL: Bug Reports, https://github.com/world-sim-dev/shared-tensor/issues
 Project-URL: Changelog, https://github.com/world-sim-dev/shared-tensor/releases
 Keywords: gpu,memory,sharing,ipc,inter-process-communication,pytorch,cuda,model-serving,inference,torch,torch-ipc
@@ -77,7 +77,7 @@ Not supported:
 
 ## Install
 
-Use Python `3.10+` and a CUDA-enabled PyTorch build.
+Use Python `3.9+` and a CUDA-enabled PyTorch build.
 
 ```bash
 pip install shared-tensor
@@ -91,6 +91,16 @@ conda activate shared-tensor-dev
 pip install -e ".[dev,test]"
 ```
 
+## Docs
+
+Read the examples first, then the design notes:
+
+- `docs/overview.md`
+- `docs/patterns.md`
+- `docs/architecture.md`
+- `docs/lifecycle.md`
+- `docs/diagrams.md`
+
 ## Example: Manual Two-Process Deployment
 
 Production should prefer two explicitly started processes: one server process that owns CUDA objects, and one or more client processes that reopen them through torch IPC.

{shared_tensor-0.2.7 → shared_tensor-0.2.8}/README.md

@@ -25,7 +25,7 @@ Not supported:
 
 ## Install
 
-Use Python `3.10+` and a CUDA-enabled PyTorch build.
+Use Python `3.9+` and a CUDA-enabled PyTorch build.
 
 ```bash
 pip install shared-tensor
@@ -39,6 +39,16 @@ conda activate shared-tensor-dev
 pip install -e ".[dev,test]"
 ```
 
+## Docs
+
+Read the examples first, then the design notes:
+
+- `docs/overview.md`
+- `docs/patterns.md`
+- `docs/architecture.md`
+- `docs/lifecycle.md`
+- `docs/diagrams.md`
+
 ## Example: Manual Two-Process Deployment
 
 Production should prefer two explicitly started processes: one server process that owns CUDA objects, and one or more client processes that reopen them through torch IPC.

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

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "shared-tensor"
-version = "0.2.7"
+version = "0.2.8"
 description = "Native PyTorch CUDA IPC over Unix Domain Socket for same-host process separation"
 readme = "README.md"
 license = "Apache-2.0"
@@ -75,7 +75,7 @@ docs = [
 [project.urls]
 Homepage = "https://github.com/world-sim-dev/shared-tensor"
 Repository = "https://github.com/world-sim-dev/shared-tensor"
-Documentation = "https://github.com/world-sim-dev/shared-tensor/wiki"
+Documentation = "https://github.com/world-sim-dev/shared-tensor/tree/main/docs"
 "Bug Reports" = "https://github.com/world-sim-dev/shared-tensor/issues"
 Changelog = "https://github.com/world-sim-dev/shared-tensor/releases"
 

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

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

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

@@ -33,8 +33,6 @@ class TaskInfo:
     created_at: float
     started_at: float | None = None
     completed_at: float | None = None
-    result_encoding: str | None = None
-    result_payload: bytes | None = None
     error_type: str | None = None
     error_message: str | None = None
     metadata: dict[str, Any] = field(default_factory=dict)
@@ -47,8 +45,6 @@ class TaskInfo:
             "created_at": self.created_at,
             "started_at": self.started_at,
             "completed_at": self.completed_at,
-            "result_encoding": self.result_encoding,
-            "result_payload": self.result_payload,
             "error_type": self.error_type,
             "error_message": self.error_message,
             "metadata": dict(self.metadata),
@@ -66,6 +62,8 @@ class TaskInfo:
 class _TaskEntry:
     info: TaskInfo
     future: Future[Any]
+    result_encoding: str | None = None
+    result_payload: bytes | None = None
     local_result: Any = None
 
 
@@ -87,6 +85,7 @@ class TaskManager:
         self._last_cleanup = 0.0
         self._lock = RLock()
         self._tasks: dict[str, _TaskEntry] = {}
+        self._accepting_submissions = True
 
     def submit(
         self,
@@ -98,6 +97,8 @@ class TaskManager:
     ) -> TaskInfo:
         self._maybe_cleanup()
         with self._lock:
+            if not self._accepting_submissions:
+                raise SharedTensorTaskError("Task manager is shutting down and is not accepting new tasks")
             self._drop_oldest_finished_tasks_if_needed()
             if len(self._tasks) >= self._max_tasks:
                 raise SharedTensorTaskError("Task capacity exceeded")
@@ -143,12 +144,11 @@ class TaskManager:
         self._store_local_result(task_id, result)
 
         if result is None:
+            self._store_payload(task_id, encoding=None, payload=None, object_id=None)
             self._transition(
                 task_id,
                 status=TaskStatus.COMPLETED,
                 completed_at=time.time(),
-                result_encoding=None,
-                result_payload=None,
             )
             return
 
@@ -168,13 +168,16 @@ class TaskManager:
             )
             return
 
+        self._store_payload(
+            task_id,
+            encoding=payload["encoding"],
+            payload=payload["payload_bytes"],
+            object_id=payload.get("object_id"),
+        )
         self._transition(
             task_id,
             status=TaskStatus.COMPLETED,
             completed_at=time.time(),
-            result_encoding=payload["encoding"],
-            result_payload=payload["payload_bytes"],
-            metadata={"object_id": payload.get("object_id")},
         )
 
     @staticmethod
@@ -201,6 +204,24 @@ class TaskManager:
                 return
             entry.local_result = value
 
+    def _store_payload(
+        self,
+        task_id: str,
+        *,
+        encoding: str | None,
+        payload: bytes | None,
+        object_id: str | None,
+    ) -> None:
+        with self._lock:
+            entry = self._tasks.get(task_id)
+            if entry is None:
+                return
+            entry.result_encoding = encoding
+            entry.result_payload = payload
+            metadata = dict(entry.info.metadata)
+            metadata["object_id"] = object_id
+            entry.info.metadata = metadata
+
     def get(self, task_id: str) -> TaskInfo:
         self._maybe_cleanup()
         with self._lock:
@@ -255,7 +276,14 @@ class TaskManager:
         return self.result_payload(task_id)
 
     def result_payload(self, task_id: str) -> dict[str, str | bytes | None]:
-        info = self.get(task_id)
+        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)
+            encoding = entry.result_encoding
+            payload = entry.result_payload
         if info.status == TaskStatus.CANCELLED:
             raise SharedTensorTaskError(f"Task '{task_id}' was cancelled")
         if info.status == TaskStatus.FAILED:
@@ -265,8 +293,8 @@ class TaskManager:
                 f"Task '{task_id}' is not complete; current status is '{info.status.value}'"
             )
         return {
-            "encoding": info.result_encoding,
-            "payload_bytes": info.result_payload,
+            "encoding": encoding,
+            "payload_bytes": payload,
             "object_id": info.metadata.get("object_id"),
         }
 
@@ -309,8 +337,10 @@ class TaskManager:
             }
         return items
 
-    def shutdown(self, *, wait: bool = True) -> None:
-        self._executor.shutdown(wait=wait, cancel_futures=True)
+    def shutdown(self, *, wait: bool = True, cancel_futures: bool = True) -> None:
+        with self._lock:
+            self._accepting_submissions = False
+        self._executor.shutdown(wait=wait, cancel_futures=cancel_futures)
 
     def _maybe_cleanup(self) -> None:
         now = time.time()

shared_tensor-0.2.8/shared_tensor/managed_object.py

@@ -0,0 +1,135 @@
+"""Managed remote object handles and registry state."""
+
+from __future__ import annotations
+
+import uuid
+from dataclasses import dataclass
+from threading import RLock
+from typing import Any, Generic, TypeVar
+
+T = TypeVar("T")
+
+
+@dataclass(slots=True)
+class ManagedObjectEntry:
+    object_id: str
+    value: Any
+    endpoint: str
+    cache_key: str | None
+    refcount: int = 1
+
+
+@dataclass(slots=True)
+class ManagedReleaseResult:
+    released: bool
+    destroyed: bool
+    refcount: int
+    cache_key: str | None
+
+
+class ManagedObjectRegistry:
+    def __init__(self) -> None:
+        self._entries: dict[str, ManagedObjectEntry] = {}
+        self._cache_index: dict[str, str] = {}
+        self._lock = RLock()
+
+    def get_cached(self, cache_key: str) -> ManagedObjectEntry | None:
+        with self._lock:
+            object_id = self._cache_index.get(cache_key)
+            if object_id is None:
+                return None
+            entry = self._entries.get(object_id)
+            if entry is None:
+                self._cache_index.pop(cache_key, None)
+                return None
+            return entry
+
+    def register(self, *, endpoint: str, value: Any, cache_key: str | None) -> ManagedObjectEntry:
+        with self._lock:
+            entry = ManagedObjectEntry(
+                object_id=uuid.uuid4().hex,
+                value=value,
+                endpoint=endpoint,
+                cache_key=cache_key,
+            )
+            self._entries[entry.object_id] = entry
+            if cache_key is not None:
+                self._cache_index[cache_key] = entry.object_id
+            return entry
+
+    def get(self, object_id: str) -> ManagedObjectEntry | None:
+        with self._lock:
+            return self._entries.get(object_id)
+
+    def add_ref(self, object_id: str) -> ManagedObjectEntry | None:
+        with self._lock:
+            entry = self._entries.get(object_id)
+            if entry is None:
+                return None
+            entry.refcount += 1
+            return entry
+
+    def release(self, object_id: str) -> ManagedReleaseResult:
+        with self._lock:
+            entry = self._entries.get(object_id)
+            if entry is None:
+                return ManagedReleaseResult(released=False, destroyed=False, refcount=0, cache_key=None)
+
+            entry.refcount -= 1
+            destroyed = entry.refcount <= 0
+            cache_key = entry.cache_key
+            refcount = max(entry.refcount, 0)
+            if destroyed:
+                self._entries.pop(object_id, None)
+                if cache_key is not None and self._cache_index.get(cache_key) == object_id:
+                    self._cache_index.pop(cache_key, None)
+            return ManagedReleaseResult(
+                released=True,
+                destroyed=destroyed,
+                refcount=refcount,
+                cache_key=cache_key,
+            )
+
+    def info(self, object_id: str) -> dict[str, Any] | None:
+        with self._lock:
+            entry = self._entries.get(object_id)
+            if entry is None:
+                return None
+            return {
+                "object_id": entry.object_id,
+                "endpoint": entry.endpoint,
+                "cache_key": entry.cache_key,
+                "refcount": entry.refcount,
+            }
+
+    def clear(self) -> None:
+        with self._lock:
+            self._entries.clear()
+            self._cache_index.clear()
+
+
+class ReleaseHandle:
+    def release(self) -> bool:  # pragma: no cover - protocol surface only
+        raise NotImplementedError
+
+
+@dataclass(slots=True)
+class SharedObjectHandle(Generic[T]):
+    object_id: str
+    value: T
+    _releaser: ReleaseHandle
+    released: bool = False
+
+    def release(self) -> bool:
+        if self.released:
+            return False
+        released = self._releaser.release()
+        if released:
+            self.released = True
+        return released
+
+    def __enter__(self) -> SharedObjectHandle[T]:
+        return self
+
+    def __exit__(self, exc_type: object, exc_val: object, exc_tb: object) -> None:
+        self.release()

{shared_tensor-0.2.7 → shared_tensor-0.2.8}/shared_tensor/provider.py

@@ -8,6 +8,7 @@ import os
 from collections.abc import Callable
 from dataclasses import dataclass
 from functools import wraps
+from threading import RLock
 from typing import Any, Literal
 
 from shared_tensor.errors import (
@@ -113,7 +114,9 @@ class SharedTensorProvider:
         self._cache: dict[str, Any] = {}
         self._endpoints: dict[str, EndpointDefinition] = {}
         self._registered_functions = self._endpoints
-        atexit.register(self.close)
+        self._lock = RLock()
+        self._atexit_registered = False
+        self._register_atexit_once()
 
     def register(
         self,
@@ -129,25 +132,26 @@ class SharedTensorProvider:
     ) -> Callable[..., Any]:
         _validate_endpoint_options(execution=execution, concurrency=concurrency)
         endpoint_name = func.__name__
-        if endpoint_name in self._endpoints:
-            raise SharedTensorProviderError(f"Endpoint '{endpoint_name}' is already registered")
+        with self._lock:
+            if endpoint_name in self._endpoints:
+                raise SharedTensorProviderError(f"Endpoint '{endpoint_name}' is already registered")
 
-        resolved_cache_format_key = (
-            func.__qualname__ if cache_format_key is None else cache_format_key
-        )
+            resolved_cache_format_key = (
+                func.__qualname__ if cache_format_key is None else cache_format_key
+            )
 
-        definition = EndpointDefinition(
-            name=endpoint_name,
-            func=func,
-            cache=cache,
-            cache_format_key=resolved_cache_format_key,
-            managed=managed,
-            async_default_wait=async_default_wait,
-            execution=execution,
-            concurrency=concurrency,
-            singleflight=singleflight,
-        )
-        self._endpoints[endpoint_name] = definition
+            definition = EndpointDefinition(
+                name=endpoint_name,
+                func=func,
+                cache=cache,
+                cache_format_key=resolved_cache_format_key,
+                managed=managed,
+                async_default_wait=async_default_wait,
+                execution=execution,
+                concurrency=concurrency,
+                singleflight=singleflight,
+            )
+            self._endpoints[endpoint_name] = definition
         if self.verbose_debug:
             logger.debug(
                 "Provider registered endpoint",
@@ -161,7 +165,7 @@ class SharedTensorProvider:
             )
 
         if self._should_autostart_server():
-            self._restart_autostart_server()
+            self._ensure_autostart_server()
 
         @wraps(func)
         def wrapper(*args: Any, **kwargs: Any) -> Any:
@@ -272,11 +276,13 @@ class SharedTensorProvider:
             return definition.func(*args, **resolved_kwargs)
 
         cache_key = self._cache_key_for(endpoint, definition, args, resolved_kwargs)
-        if cache_key in self._cache:
-            return self._cache[cache_key]
+        with self._lock:
+            if cache_key in self._cache:
+                return self._cache[cache_key]
 
         result = definition.func(*args, **resolved_kwargs)
-        self._cache[cache_key] = result
+        with self._lock:
+            self._cache[cache_key] = result
         return result
 
     def get_endpoint(self, endpoint: str) -> EndpointDefinition:
@@ -309,18 +315,19 @@ class SharedTensorProvider:
             self._async_client.close()
             self._async_client = None
         if self._server is not None:
-            self._server.stop()
+            self._server.stop(wait_for_tasks=True)
             self._server = None
 
     def get_runtime_info(self) -> dict[str, Any]:
         if self.execution_mode in {"server", "local"}:
+            server = self._server
             return {
                 "execution_mode": self.execution_mode,
                 "auto_mode": self.auto_mode,
                 "base_path": self.base_path,
                 "device_index": self.device_index,
                 "server_socket_path": resolve_runtime_socket_path(self.base_path, self.device_index),
-                "server_running": self._server is not None,
+                "server_running": bool(server is not None and getattr(server, "running", True)),
             }
         server_info = self._get_client().get_server_info()
         return {
@@ -361,18 +368,18 @@ class SharedTensorProvider:
     def _should_autostart_server(self) -> bool:
         return self.auto_mode and self.execution_mode == "server"
 
-    def _restart_autostart_server(self) -> None:
+    def _ensure_autostart_server(self) -> None:
         from shared_tensor.server import SharedTensorServer
 
+        if self._server is not None:
+            return
         if self.verbose_debug:
             logger.debug(
-                "Provider restarting autostart server",
+                "Provider starting autostart server",
                 extra={
                     "socket_path": resolve_runtime_socket_path(self.base_path, self.device_index),
                 },
             )
-        if self._server is not None:
-            self._server.stop()
         self._server = SharedTensorServer(
             self,
             socket_path=resolve_runtime_socket_path(self.base_path, self.device_index),
@@ -395,3 +402,9 @@ class SharedTensorProvider:
             func=definition.func,
             cache_format_key=definition.cache_format_key,
         )
+
+    def _register_atexit_once(self) -> None:
+        if self._atexit_registered:
+            return
+        atexit.register(self.close)
+        self._atexit_registered = True

{shared_tensor-0.2.7 → shared_tensor-0.2.8}/shared_tensor/runtime.py

@@ -5,6 +5,8 @@ from __future__ import annotations
 from threading import RLock
 from typing import TYPE_CHECKING
 
+from shared_tensor.errors import SharedTensorConfigurationError
+
 if TYPE_CHECKING:
     from shared_tensor.server import SharedTensorServer
 
@@ -15,6 +17,11 @@ _SERVERS: dict[str, "SharedTensorServer"] = {}
 
 def register_local_server(socket_path: str, server: "SharedTensorServer") -> None:
     with _LOCK:
+        current = _SERVERS.get(socket_path)
+        if current is not None and current is not server:
+            raise SharedTensorConfigurationError(
+                f"Local runtime socket '{socket_path}' is already registered by another server"
+            )
         _SERVERS[socket_path] = server
 
 

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

@@ -39,6 +39,24 @@ from shared_tensor.utils import (
 logger = logging.getLogger(__name__)
 
 
+class _ConnectionExecutor:
+    def __init__(self, *, max_workers: int) -> None:
+        self._semaphore = threading.BoundedSemaphore(max_workers)
+
+    def submit(self, func, *args, **kwargs) -> threading.Thread:
+        self._semaphore.acquire()
+
+        def runner() -> None:
+            try:
+                func(*args, **kwargs)
+            finally:
+                self._semaphore.release()
+
+        thread = threading.Thread(target=runner, daemon=True)
+        thread.start()
+        return thread
+
+
 def _server_version() -> str:
     try:
         from shared_tensor import __version__
@@ -49,7 +67,7 @@ def _server_version() -> str:
 
 @dataclass(slots=True)
 class _InFlightCall:
-    future: Future[dict[str, Any]]
+    future: Future
 
 
 @dataclass(slots=True)
@@ -107,6 +125,8 @@ class SharedTensorServer:
         self._inflight: dict[str, _InFlightCall] = {}
         self._endpoint_locks: dict[str, threading.Lock] = {}
         self._coordination_lock = threading.RLock()
+        self._connection_executor = _ConnectionExecutor(max_workers=max_workers)
+        self._accepting_requests = True
         if getattr(self.provider, "_server", None) is None:
             self.provider._server = self
 
@@ -114,6 +134,8 @@ class SharedTensorServer:
         if self.verbose_debug:
             logger.debug("Server processing request", extra={"method": request.get("method")})
         try:
+            if not self._accepting_requests:
+                raise SharedTensorConfigurationError("Server is stopping and not accepting new requests")
             method = request.get("method")
             if not isinstance(method, str) or not method:
                 raise SharedTensorProtocolError("Missing required field 'method'")
@@ -293,7 +315,8 @@ class SharedTensorServer:
         if self.verbose_debug:
             logger.debug("Server executed direct endpoint", extra={"endpoint": endpoint})
         if cache_key is not None:
-            self._local_cache[cache_key] = value
+            with self._coordination_lock:
+                self._local_cache[cache_key] = value
         return _EndpointResult(value=value)
 
     def _materialize_managed_result(
@@ -329,7 +352,8 @@ class SharedTensorServer:
                 return None
             self._managed_objects.add_ref(cached.object_id)
             return _EndpointResult(value=cached.value, object_id=cached.object_id)
-        local_value = self._local_cache.get(cache_key)
+        with self._coordination_lock:
+            local_value = self._local_cache.get(cache_key)
         if local_value is None:
             return None
         return _EndpointResult(value=local_value)
@@ -397,11 +421,14 @@ class SharedTensorServer:
                     return existing.value
                 self._managed_objects.register(endpoint=endpoint, value=value, cache_key=cache_key)
             return value
-        if cache_key is not None and cache_key in self._local_cache:
-            return self._local_cache[cache_key]
+        if cache_key is not None:
+            with self._coordination_lock:
+                if cache_key in self._local_cache:
+                    return self._local_cache[cache_key]
         value = definition.func(*args, **resolved_kwargs)
         if cache_key is not None:
-            self._local_cache[cache_key] = value
+            with self._coordination_lock:
+                self._local_cache[cache_key] = value
         return value
 
     def _cache_key(
@@ -421,16 +448,16 @@ class SharedTensorServer:
             cache_format_key=definition.cache_format_key,
         )
 
-    def _acquire_inflight(self, inflight_key: str) -> tuple[Future[dict[str, Any]], bool]:
+    def _acquire_inflight(self, inflight_key: str) -> tuple[Future, bool]:
         with self._coordination_lock:
             inflight = self._inflight.get(inflight_key)
             if inflight is not None:
                 return inflight.future, False
-            future: Future[dict[str, Any]] = Future()
+            future = Future()
             self._inflight[inflight_key] = _InFlightCall(future=future)
             return future, True
 
-    def _release_inflight(self, inflight_key: str | None, future: Future[dict[str, Any]]) -> None:
+    def _release_inflight(self, inflight_key: str | None, future: Future) -> None:
         if inflight_key is None:
             return
         with self._coordination_lock:
@@ -535,7 +562,7 @@ class SharedTensorServer:
             "socket_path": self.socket_path,
             "uptime": uptime,
             "running": self.running,
-            "ready": self.running and self.listener is not None,
+            "ready": self.running and self.listener is not None and self._accepting_requests,
             "pid": os.getpid(),
             "ppid": os.getppid(),
             "device_index": resolve_device_index(self.provider.device_index),
@@ -558,6 +585,7 @@ class SharedTensorServer:
             logger.info("Server starting", extra={"socket_path": self.socket_path, "blocking": blocking})
         if self.running or self.server_thread is not None:
             raise SharedTensorConfigurationError("Server is already running")
+        self._accepting_requests = True
         if blocking:
             self._resolved_process_start_method = None
             self._serve_forever()
@@ -576,11 +604,11 @@ class SharedTensorServer:
         self._resolved_process_start_method = "thread"
         thread.start()
         if not state.ready.wait(timeout=self.startup_timeout):
-            self.stop()
+            self.stop(wait_for_tasks=False)
             raise TimeoutError(f"Timed out waiting for server socket {self.socket_path}")
         if state.error is not None:
             error = state.error
-            self.stop()
+            self.stop(wait_for_tasks=False)
             raise SharedTensorConfigurationError(
                 f"Failed to start background server thread for {self.socket_path}: {error}"
             ) from error
@@ -620,12 +648,11 @@ class SharedTensorServer:
                     if self.running:
                         raise
                     break
-                thread = threading.Thread(target=self._handle_connection, args=(conn,), daemon=True)
-                thread.start()
+                self._connection_executor.submit(self._handle_connection, conn)
         finally:
             if started_event is not None and not started_event.is_set():
                 started_event.set()
-            self._shutdown_local_resources()
+            self._shutdown_local_resources(wait_for_tasks=True)
 
     def _handle_connection(self, conn: socket.socket) -> None:
         with conn:
@@ -655,9 +682,10 @@ class SharedTensorServer:
         if 0 <= local_rank < torch.cuda.device_count():
             torch.cuda.set_device(local_rank)
 
-    def stop(self) -> None:
+    def stop(self, *, wait_for_tasks: bool = True) -> None:
         if self.verbose_debug:
             logger.info("Server stopping", extra={"socket_path": self.socket_path})
+        self._accepting_requests = False
         self.running = False
         if self.listener is not None:
             self.listener.close()
@@ -668,21 +696,23 @@ class SharedTensorServer:
         self.server_thread = None
         self.server_process = None
         if self.listener is None:
-            unlink_socket_path(self.socket_path)
+            self._shutdown_local_resources(wait_for_tasks=wait_for_tasks)
 
-    def _shutdown_local_resources(self) -> None:
+    def _shutdown_local_resources(self, *, wait_for_tasks: bool) -> None:
+        self._accepting_requests = False
         self.running = False
         if self.listener is not None:
             self.listener.close()
             self.listener = None
         if self._task_manager is not None:
-            self._task_manager.shutdown(wait=False)
+            self._task_manager.shutdown(wait=wait_for_tasks, cancel_futures=not wait_for_tasks)
             self._task_manager = None
         self._managed_objects.clear()
-        self._cache.clear()
-        self._local_cache.clear()
-        self._inflight.clear()
-        self._endpoint_locks.clear()
+        with self._coordination_lock:
+            self._cache.clear()
+            self._local_cache.clear()
+            self._inflight.clear()
+            self._endpoint_locks.clear()
         unregister_local_server(self.socket_path, self)
         unlink_socket_path(self.socket_path)
 

shared_tensor-0.2.7/shared_tensor/managed_object.py

@@ -1,126 +0,0 @@
-"""Managed remote object handles and registry state."""
-
-from __future__ import annotations
-
-import uuid
-from dataclasses import dataclass
-from typing import Any, Generic, TypeVar
-
-T = TypeVar("T")
-
-
-@dataclass(slots=True)
-class ManagedObjectEntry:
-    object_id: str
-    value: Any
-    endpoint: str
-    cache_key: str | None
-    refcount: int = 1
-
-
-@dataclass(slots=True)
-class ManagedReleaseResult:
-    released: bool
-    destroyed: bool
-    refcount: int
-    cache_key: str | None
-
-
-class ManagedObjectRegistry:
-    def __init__(self) -> None:
-        self._entries: dict[str, ManagedObjectEntry] = {}
-        self._cache_index: dict[str, str] = {}
-
-    def get_cached(self, cache_key: str) -> ManagedObjectEntry | None:
-        object_id = self._cache_index.get(cache_key)
-        if object_id is None:
-            return None
-        entry = self._entries.get(object_id)
-        if entry is None:
-            self._cache_index.pop(cache_key, None)
-            return None
-        return entry
-
-    def register(self, *, endpoint: str, value: Any, cache_key: str | None) -> ManagedObjectEntry:
-        entry = ManagedObjectEntry(
-            object_id=uuid.uuid4().hex,
-            value=value,
-            endpoint=endpoint,
-            cache_key=cache_key,
-        )
-        self._entries[entry.object_id] = entry
-        if cache_key is not None:
-            self._cache_index[cache_key] = entry.object_id
-        return entry
-
-    def get(self, object_id: str) -> ManagedObjectEntry | None:
-        return self._entries.get(object_id)
-
-    def add_ref(self, object_id: str) -> ManagedObjectEntry | None:
-        entry = self._entries.get(object_id)
-        if entry is None:
-            return None
-        entry.refcount += 1
-        return entry
-
-    def release(self, object_id: str) -> ManagedReleaseResult:
-        entry = self._entries.get(object_id)
-        if entry is None:
-            return ManagedReleaseResult(released=False, destroyed=False, refcount=0, cache_key=None)
-
-        entry.refcount -= 1
-        destroyed = entry.refcount <= 0
-        cache_key = entry.cache_key
-        refcount = max(entry.refcount, 0)
-        if destroyed:
-            self._entries.pop(object_id, None)
-            if cache_key is not None and self._cache_index.get(cache_key) == object_id:
-                self._cache_index.pop(cache_key, None)
-        return ManagedReleaseResult(
-            released=True,
-            destroyed=destroyed,
-            refcount=refcount,
-            cache_key=cache_key,
-        )
-
-    def info(self, object_id: str) -> dict[str, Any] | None:
-        entry = self._entries.get(object_id)
-        if entry is None:
-            return None
-        return {
-            "object_id": entry.object_id,
-            "endpoint": entry.endpoint,
-            "cache_key": entry.cache_key,
-            "refcount": entry.refcount,
-        }
-
-    def clear(self) -> None:
-        self._entries.clear()
-        self._cache_index.clear()
-
-
-class ReleaseHandle:
-    def release(self) -> bool:  # pragma: no cover - protocol surface only
-        raise NotImplementedError
-
-
-@dataclass(slots=True)
-class SharedObjectHandle(Generic[T]):
-    object_id: str
-    value: T
-    _releaser: ReleaseHandle
-    released: bool = False
-
-    def release(self) -> bool:
-        if self.released:
-            return False
-        released = self._releaser.release()
-        if released:
-            self.released = True
-        return released
-
-    def __enter__(self) -> SharedObjectHandle[T]:
-        return self
-
-    def __exit__(self, exc_type: object, exc_val: object, exc_tb: object) -> None:
-        self.release()