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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: shared-tensor
-Version: 0.2.5
+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"
@@ -63,6 +65,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 +91,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 +153,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 +216,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.7}/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.7}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "shared-tensor"
-version = "0.2.5"
+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.5 → shared_tensor-0.2.7}/shared_tensor/__init__.py

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

{shared_tensor-0.2.5 → 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.5 → 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.5 → shared_tensor-0.2.7}/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.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.5 → shared_tensor-0.2.7}/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
@@ -25,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,
@@ -41,11 +39,33 @@ 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
+
+
+@dataclass(slots=True)
+class _EndpointResult:
+    value: Any
+    object_id: str | None = None
+
+
 class SharedTensorServer:
     def __init__(
         self,
@@ -72,6 +92,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 +102,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:
@@ -180,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})
@@ -207,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)
@@ -228,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,
@@ -239,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)
@@ -256,16 +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
-        return result
+            self._local_cache[cache_key] = value
+        return _EndpointResult(value=value)
 
     def _materialize_managed_result(
         self,
@@ -274,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:
@@ -299,8 +328,81 @@ 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)
-        return self._cache.get(cache_key)
+            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
+        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,
+        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,
@@ -400,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(
@@ -426,7 +531,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 +553,66 @@ 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()
+            register_local_server(self.socket_path, self)
+            if started_event is not None:
+                started_event.set()
             while self.running:
                 try:
                     conn, _ = listener.accept()
@@ -553,6 +623,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 +658,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,8 +680,10 @@ class SharedTensorServer:
             self._task_manager = None
         self._managed_objects.clear()
         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)
 
     def __enter__(self) -> SharedTensorServer:

{shared_tensor-0.2.5 → 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