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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: shared-tensor
-Version: 0.2.5
+Version: 0.2.6
 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>
@@ -63,6 +63,7 @@ Supported:
 - sync `call` and task-backed `submit`
 - managed object handles with explicit release
 - server-side caching, `cache_format_key`, and singleflight
+- manual two-process deployment as the primary production path
 - zero-branch auto mode gated by `SHARED_TENSOR_ENABLED=1`
 
 Not supported:
@@ -88,46 +89,58 @@ conda activate shared-tensor-dev
 pip install -e ".[dev,test]"
 ```
 
-## Example: Same Code, Two Processes
+## 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.
 
-See [examples/zero_branch_env.py](./examples/zero_branch_env.py).
+See [examples/model_service.py](./examples/model_service.py) for endpoint definitions.
+
+Server process:
 
 ```python
-import torch
+from shared_tensor import SharedTensorProvider, SharedTensorServer
 
-from shared_tensor import SharedObjectHandle, SharedTensorProvider
+provider = SharedTensorProvider(execution_mode="server")
 
-provider = SharedTensorProvider()
+@provider.share(execution="task", managed=True, concurrency="serialized", cache_format_key="model:{hidden_size}")
+def load_model(hidden_size: int = 4):
+    ...
 
+server = SharedTensorServer(provider)
+server.start(blocking=True)
+```
 
-@provider.share(
-    execution="task",
-    managed=True,
-    concurrency="serialized",
-    cache_format_key="model:{hidden_size}",
-)
-def load_model(hidden_size: int = 4) -> torch.nn.Module:
-    return torch.nn.Linear(hidden_size, 2, device="cuda")
+Client process:
 
+```python
+import torch
+
+from shared_tensor import SharedObjectHandle, SharedTensorClient
 
+client = SharedTensorClient()
 x = torch.ones(1, 4, device="cuda")
-result = load_model(hidden_size=4)
+result = client.call("load_model", hidden_size=4)
 if isinstance(result, SharedObjectHandle):
     with result as handle:
         y = handle.value(x)
-else:
-    y = result(x)
 ```
 
-Server process:
+This keeps the contract explicit:
 
-```bash
-SHARED_TENSOR_ENABLED=1 SHARED_TENSOR_ROLE=server python demo.py
+```text
+server process                      client process
+------------------------------      ------------------------------
+owns CUDA allocations               issues local UDS RPC requests
+executes endpoint functions         reopens CUDA objects via torch IPC
+manages cache and refcounts         releases managed handles explicitly
 ```
 
-Client process with the exact same file:
+## Example: Same Code, Two Processes
+
+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.
 
 ```bash
+SHARED_TENSOR_ENABLED=1 SHARED_TENSOR_ROLE=server python demo.py
 SHARED_TENSOR_ENABLED=1 python demo.py
 ```
 
@@ -138,7 +151,7 @@ same code
 
 server process                      client process
 ------------------------------      ------------------------------
-provider auto-starts UDS daemon     provider builds client wrappers
+provider auto-starts local thread   provider builds client wrappers
 shared function runs locally        shared function becomes RPC call
 CUDA object stays on same GPU       CUDA object is reopened via torch IPC
 ```
@@ -201,19 +214,19 @@ SharedTensorProvider(enabled=None)
 Provider runtime controls:
 
 ```python
-SharedTensorProvider(server_process_start_method="fork")
 SharedTensorProvider(server_startup_timeout=30.0)
 provider.get_runtime_info()
 ```
 
-Use `server_process_start_method="fork"` when you explicitly want POSIX fork behavior.
-Leave it as `None` to let the library choose a safer default for the current entrypoint.
+Non-blocking provider autostart runs the UDS server in a background thread inside the current process.
 
 `execution_mode="auto"` behaves as follows:
 - disabled: local mode
-- enabled + `SHARED_TENSOR_ROLE=server`: auto-start local server and execute endpoints locally
+- enabled + `SHARED_TENSOR_ROLE=server`: auto-start a local background server thread and execute endpoints locally
 - enabled + role unset: build client wrappers
 
+For production deployment, prefer explicit `SharedTensorServer(...).start(blocking=True)` in a dedicated server process.
+
 Socket selection is per CUDA device:
 - base path comes from `SHARED_TENSOR_BASE_PATH` or `/tmp/shared-tensor`
 - runtime socket path is `<base_path>-<device_index>.sock`

{shared_tensor-0.2.5 → shared_tensor-0.2.6}/README.md

@@ -13,6 +13,7 @@ Supported:
 - sync `call` and task-backed `submit`
 - managed object handles with explicit release
 - server-side caching, `cache_format_key`, and singleflight
+- manual two-process deployment as the primary production path
 - zero-branch auto mode gated by `SHARED_TENSOR_ENABLED=1`
 
 Not supported:
@@ -38,46 +39,58 @@ conda activate shared-tensor-dev
 pip install -e ".[dev,test]"
 ```
 
-## Example: Same Code, Two Processes
+## 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.
 
-See [examples/zero_branch_env.py](./examples/zero_branch_env.py).
+See [examples/model_service.py](./examples/model_service.py) for endpoint definitions.
+
+Server process:
 
 ```python
-import torch
+from shared_tensor import SharedTensorProvider, SharedTensorServer
 
-from shared_tensor import SharedObjectHandle, SharedTensorProvider
+provider = SharedTensorProvider(execution_mode="server")
 
-provider = SharedTensorProvider()
+@provider.share(execution="task", managed=True, concurrency="serialized", cache_format_key="model:{hidden_size}")
+def load_model(hidden_size: int = 4):
+    ...
 
+server = SharedTensorServer(provider)
+server.start(blocking=True)
+```
 
-@provider.share(
-    execution="task",
-    managed=True,
-    concurrency="serialized",
-    cache_format_key="model:{hidden_size}",
-)
-def load_model(hidden_size: int = 4) -> torch.nn.Module:
-    return torch.nn.Linear(hidden_size, 2, device="cuda")
+Client process:
 
+```python
+import torch
+
+from shared_tensor import SharedObjectHandle, SharedTensorClient
 
+client = SharedTensorClient()
 x = torch.ones(1, 4, device="cuda")
-result = load_model(hidden_size=4)
+result = client.call("load_model", hidden_size=4)
 if isinstance(result, SharedObjectHandle):
     with result as handle:
         y = handle.value(x)
-else:
-    y = result(x)
 ```
 
-Server process:
+This keeps the contract explicit:
 
-```bash
-SHARED_TENSOR_ENABLED=1 SHARED_TENSOR_ROLE=server python demo.py
+```text
+server process                      client process
+------------------------------      ------------------------------
+owns CUDA allocations               issues local UDS RPC requests
+executes endpoint functions         reopens CUDA objects via torch IPC
+manages cache and refcounts         releases managed handles explicitly
 ```
 
-Client process with the exact same file:
+## Example: Same Code, Two Processes
+
+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.
 
 ```bash
+SHARED_TENSOR_ENABLED=1 SHARED_TENSOR_ROLE=server python demo.py
 SHARED_TENSOR_ENABLED=1 python demo.py
 ```
 
@@ -88,7 +101,7 @@ same code
 
 server process                      client process
 ------------------------------      ------------------------------
-provider auto-starts UDS daemon     provider builds client wrappers
+provider auto-starts local thread   provider builds client wrappers
 shared function runs locally        shared function becomes RPC call
 CUDA object stays on same GPU       CUDA object is reopened via torch IPC
 ```
@@ -151,19 +164,19 @@ SharedTensorProvider(enabled=None)
 Provider runtime controls:
 
 ```python
-SharedTensorProvider(server_process_start_method="fork")
 SharedTensorProvider(server_startup_timeout=30.0)
 provider.get_runtime_info()
 ```
 
-Use `server_process_start_method="fork"` when you explicitly want POSIX fork behavior.
-Leave it as `None` to let the library choose a safer default for the current entrypoint.
+Non-blocking provider autostart runs the UDS server in a background thread inside the current process.
 
 `execution_mode="auto"` behaves as follows:
 - disabled: local mode
-- enabled + `SHARED_TENSOR_ROLE=server`: auto-start local server and execute endpoints locally
+- enabled + `SHARED_TENSOR_ROLE=server`: auto-start a local background server thread and execute endpoints locally
 - enabled + role unset: build client wrappers
 
+For production deployment, prefer explicit `SharedTensorServer(...).start(blocking=True)` in a dedicated server process.
+
 Socket selection is per CUDA device:
 - base path comes from `SHARED_TENSOR_BASE_PATH` or `/tmp/shared-tensor`
 - runtime socket path is `<base_path>-<device_index>.sock`

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

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "shared-tensor"
-version = "0.2.5"
+version = "0.2.6"
 description = "Native PyTorch CUDA IPC over Unix Domain Socket for same-host process separation"
 readme = "README.md"
 license = "Apache-2.0"

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

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

{shared_tensor-0.2.5 → shared_tensor-0.2.6}/shared_tensor/provider.py

@@ -92,7 +92,6 @@ class SharedTensorProvider:
         device_index: int | None = None,
         timeout: float = 30.0,
         execution_mode: str = "auto",
-        server_process_start_method: str | None = None,
         server_startup_timeout: float = 30.0,
         verbose_debug: bool = False,
     ) -> None:
@@ -106,7 +105,6 @@ class SharedTensorProvider:
         self.timeout = timeout
         self.execution_mode = resolved_mode
         self.auto_mode = auto_mode
-        self.server_process_start_method = server_process_start_method
         self.server_startup_timeout = server_startup_timeout
         self.verbose_debug = verbose_debug
         self._client: Any | None = None
@@ -165,9 +163,6 @@ class SharedTensorProvider:
         if self._should_autostart_server():
             self._restart_autostart_server()
 
-        if self.execution_mode == "server":
-            return func
-
         @wraps(func)
         def wrapper(*args: Any, **kwargs: Any) -> Any:
             return self.call(endpoint_name, *args, **kwargs)
@@ -215,7 +210,11 @@ class SharedTensorProvider:
     def call(self, endpoint: str, *args: Any, **kwargs: Any) -> Any:
         if self.verbose_debug:
             logger.debug("Provider dispatching call", extra={"endpoint": endpoint, "mode": self.execution_mode})
-        if self.execution_mode in {"server", "local"}:
+        if self.execution_mode == "server":
+            if self._server is not None and hasattr(self._server, "invoke_local"):
+                return self._server.invoke_local(endpoint, args=args, kwargs=kwargs)
+            return self.invoke_local(endpoint, args=args, kwargs=kwargs)
+        if self.execution_mode == "local":
             return self.invoke_local(endpoint, args=args, kwargs=kwargs)
         return self._get_client().call(endpoint, *args, **kwargs)
 
@@ -370,7 +369,6 @@ class SharedTensorProvider:
                 "Provider restarting autostart server",
                 extra={
                     "socket_path": resolve_runtime_socket_path(self.base_path, self.device_index),
-                    "process_start_method": self.server_process_start_method,
                 },
             )
         if self._server is not None:
@@ -378,7 +376,6 @@ class SharedTensorProvider:
         self._server = SharedTensorServer(
             self,
             socket_path=resolve_runtime_socket_path(self.base_path, self.device_index),
-            process_start_method=self.server_process_start_method,
             startup_timeout=self.server_startup_timeout,
             verbose_debug=self.verbose_debug,
         )

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

@@ -2,16 +2,13 @@
 
 from __future__ import annotations
 
-import cloudpickle
 import logging
-import multiprocessing as mp
 import os
-import sys
 import socket
 import threading
 import time
 from concurrent.futures import Future
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Any
 
 from shared_tensor.async_task import TaskManager, TaskStatus
@@ -41,11 +38,27 @@ from shared_tensor.utils import (
 logger = logging.getLogger(__name__)
 
 
+def _server_version() -> str:
+    try:
+        from shared_tensor import __version__
+    except ImportError:
+        return "unknown"
+    return __version__
+
+
 @dataclass(slots=True)
 class _InFlightCall:
     future: Future[dict[str, Any]]
 
 
+@dataclass(slots=True)
+class _ServerThreadState:
+    thread: threading.Thread
+    ready: threading.Event = field(default_factory=threading.Event)
+    stopped: threading.Event = field(default_factory=threading.Event)
+    error: BaseException | None = None
+
+
 class SharedTensorServer:
     def __init__(
         self,
@@ -72,6 +85,7 @@ class SharedTensorServer:
         self.startup_timeout = startup_timeout
         self.listener: socket.socket | None = None
         self.server_process: Any | None = None
+        self.server_thread: _ServerThreadState | None = None
         self._resolved_process_start_method: str | None = None
         self.running = False
         self.started_at: float | None = None
@@ -81,10 +95,13 @@ class SharedTensorServer:
         }
         self._task_manager: TaskManager | None = None
         self._cache: dict[str, dict[str, Any]] = {}
+        self._local_cache: dict[str, Any] = {}
         self._managed_objects = ManagedObjectRegistry()
         self._inflight: dict[str, _InFlightCall] = {}
         self._endpoint_locks: dict[str, threading.Lock] = {}
         self._coordination_lock = threading.RLock()
+        if getattr(self.provider, "_server", None) is None:
+            self.provider._server = self
 
     def process_request(self, request: dict[str, Any]) -> dict[str, Any]:
         if self.verbose_debug:
@@ -265,6 +282,7 @@ class SharedTensorServer:
         result = self._encode_result(value)
         if cache_key is not None:
             self._cache[cache_key] = result
+            self._local_cache[cache_key] = value
         return result
 
     def _materialize_managed_result(
@@ -300,7 +318,44 @@ class SharedTensorServer:
                 return None
             self._managed_objects.add_ref(cached.object_id)
             return self._encode_result(cached.value, object_id=cached.object_id)
-        return self._cache.get(cache_key)
+        cached = self._cache.get(cache_key)
+        if cached is not None:
+            return cached
+        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
+
+    def invoke_local(
+        self,
+        endpoint: str,
+        *,
+        args: tuple[Any, ...] = (),
+        kwargs: dict[str, Any] | None = None,
+    ) -> Any:
+        definition = self.provider.get_endpoint(endpoint)
+        resolved_kwargs = kwargs or {}
+        cache_key = self._cache_key(endpoint, definition, args, resolved_kwargs)
+        if definition.managed:
+            if cache_key is not None:
+                cached = self._managed_objects.get_cached(cache_key)
+                if cached is not None:
+                    return cached.value
+            value = definition.func(*args, **resolved_kwargs)
+            if cache_key is not None:
+                existing = self._managed_objects.get_cached(cache_key)
+                if existing is not None:
+                    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]
+        value = definition.func(*args, **resolved_kwargs)
+        if cache_key is not None:
+            self._local_cache[cache_key] = value
+        return value
 
     def _cache_key(
         self,
@@ -426,7 +481,7 @@ class SharedTensorServer:
         uptime = 0.0 if self.started_at is None else time.time() - self.started_at
         return {
             "server": "SharedTensorServer",
-            "version": "0.2.4",
+            "version": _server_version(),
             "socket_path": self.socket_path,
             "uptime": uptime,
             "running": self.running,
@@ -448,101 +503,65 @@ class SharedTensorServer:
             "data": None,
         }
 
-    def _resolve_process_start_method(self) -> str:
-        if self.process_start_method is not None:
-            allowed = set(mp.get_all_start_methods())
-            if self.process_start_method not in allowed:
-                raise SharedTensorConfigurationError(
-                    f"Unsupported process_start_method '{self.process_start_method}'"
-                )
-            return self.process_start_method
-        if os.name != "posix":
-            return "spawn"
-        try:
-            import torch
-        except ImportError:
-            torch = None
-        if torch is not None and torch.cuda.is_available() and torch.cuda.is_initialized():
-            return "spawn"
-        if not hasattr(sys.modules.get("__main__"), "__file__"):
-            return "fork"
-        return "spawn"
-
     def start(self, blocking: bool = True) -> None:
         if self.verbose_debug:
             logger.info("Server starting", extra={"socket_path": self.socket_path, "blocking": blocking})
-        if self.running:
+        if self.running or self.server_thread is not None:
             raise SharedTensorConfigurationError("Server is already running")
         if blocking:
             self._resolved_process_start_method = None
             self._serve_forever()
             return
-        if os.name != "posix":
+        if self.process_start_method is not None:
             raise SharedTensorConfigurationError(
-                "Non-blocking shared_tensor servers require POSIX multiprocessing support"
+                "process_start_method is not supported for thread-backed non-blocking servers"
             )
-        start_method = self._resolve_process_start_method()
-        payload = cloudpickle.dumps(self.provider)
-        process = mp.get_context(start_method).Process(
-            target=self._serve_forever_from_payload,
-            args=(
-                payload,
-                self.socket_path,
-                self.max_request_bytes,
-                self.max_workers,
-                self.result_ttl,
-                self.verbose_debug,
-                start_method,
-            ),
-            name=f"shared-tensor-daemon:{self.socket_path}",
+        thread = threading.Thread(
+            target=self._serve_forever_in_thread,
+            name=f"shared-tensor-server:{self.socket_path}",
+            daemon=True,
         )
-        process.start()
-        if self.verbose_debug:
-            logger.info(
-                "Server spawned background process",
-                extra={"socket_path": self.socket_path, "pid": process.pid, "start_method": start_method},
-            )
-        self.server_process = process
-        self._resolved_process_start_method = start_method
-        self.running = True
-        self.started_at = time.time()
+        state = _ServerThreadState(thread=thread)
+        self.server_thread = state
+        self._resolved_process_start_method = "thread"
+        thread.start()
+        if not state.ready.wait(timeout=self.startup_timeout):
+            self.stop()
+            raise TimeoutError(f"Timed out waiting for server socket {self.socket_path}")
+        if state.error is not None:
+            error = state.error
+            self.stop()
+            raise SharedTensorConfigurationError(
+                f"Failed to start background server thread for {self.socket_path}: {error}"
+            ) from error
 
-    @staticmethod
-    def _serve_forever_from_payload(
-        payload: bytes,
-        socket_path: str,
-        max_request_bytes: int,
-        max_workers: int,
-        result_ttl: float,
-        verbose_debug: bool,
-        process_start_method: str | None,
-    ) -> None:
-        SharedTensorServer._configure_cuda_runtime()
-        provider = cloudpickle.loads(payload)
-        server = SharedTensorServer(
-            provider,
-            socket_path=socket_path,
-            max_request_bytes=max_request_bytes,
-            max_workers=max_workers,
-            result_ttl=result_ttl,
-            process_start_method=process_start_method,
-            verbose_debug=verbose_debug,
-        )
-        server._resolved_process_start_method = process_start_method
-        server._serve_forever()
+    def _serve_forever_in_thread(self) -> None:
+        state = self.server_thread
+        if state is None:
+            return
+        try:
+            self._serve_forever(started_event=state.ready)
+        except BaseException as exc:  # noqa: BLE001
+            state.error = exc
+            state.ready.set()
+            raise
+        finally:
+            state.stopped.set()
 
-    def _serve_forever(self) -> None:
+    def _serve_forever(self, *, started_event: threading.Event | None = None) -> None:
         self._configure_cuda_runtime()
         unlink_socket_path(self.socket_path)
         listener = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
-        listener.bind(self.socket_path)
-        listener.listen()
-        if self.verbose_debug:
-            logger.info("Server listening", extra={"socket_path": self.socket_path})
-        self.listener = listener
-        self.running = True
-        self.started_at = time.time()
         try:
+            listener.bind(self.socket_path)
+            listener.listen()
+            if self.verbose_debug:
+                logger.info("Server listening", extra={"socket_path": self.socket_path})
+            self.listener = listener
+            self.running = True
+            self.started_at = time.time()
+            if started_event is not None:
+                started_event.set()
             while self.running:
                 try:
                     conn, _ = listener.accept()
@@ -553,6 +572,8 @@ class SharedTensorServer:
                 thread = threading.Thread(target=self._handle_connection, args=(conn,), daemon=True)
                 thread.start()
         finally:
+            if started_event is not None and not started_event.is_set():
+                started_event.set()
             self._shutdown_local_resources()
 
     def _handle_connection(self, conn: socket.socket) -> None:
@@ -586,24 +607,20 @@ class SharedTensorServer:
     def stop(self) -> None:
         if self.verbose_debug:
             logger.info("Server stopping", extra={"socket_path": self.socket_path})
-        if not self.running:
-            unlink_socket_path(self.socket_path)
-            return
         self.running = False
-        if self.server_process is not None:
-            self.server_process.terminate()
-            self.server_process.join(timeout=5)
-            if self.server_process.is_alive():
-                self.server_process.kill()
-                self.server_process.join(timeout=5)
-            self.server_process = None
-            unlink_socket_path(self.socket_path)
-            return
         if self.listener is not None:
             self.listener.close()
-        self._shutdown_local_resources()
+        state = self.server_thread
+        if state is not None and state.thread.is_alive() and threading.current_thread() is not state.thread:
+            state.stopped.wait(timeout=5)
+            state.thread.join(timeout=5)
+        self.server_thread = None
+        self.server_process = None
+        if self.listener is None:
+            unlink_socket_path(self.socket_path)
 
     def _shutdown_local_resources(self) -> None:
+        self.running = False
         if self.listener is not None:
             self.listener.close()
             self.listener = None
@@ -612,6 +629,7 @@ class SharedTensorServer:
             self._task_manager = None
         self._managed_objects.clear()
         self._cache.clear()
+        self._local_cache.clear()
         self._inflight.clear()
         self._endpoint_locks.clear()
         unlink_socket_path(self.socket_path)