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

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

@@ -1,13 +1,13 @@
 Metadata-Version: 2.4
 Name: shared-tensor
-Version: 0.2.6
+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
@@ -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"
@@ -75,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
@@ -89,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.6 → 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.6 → shared_tensor-0.2.8}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "shared-tensor"
-version = "0.2.6"
+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"
@@ -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]
@@ -73,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"
 
@@ -89,7 +91,7 @@ shared_tensor = ["*.so", "*.dll", "*.dylib"]
 
 [tool.black]
 line-length = 88
-target-version = ['py310', 'py311', 'py312']
+target-version = ['py39', 'py310', 'py311', 'py312', 'py313']
 include = '\.pyi?$'
 extend-exclude = '''
 /(
@@ -115,7 +117,7 @@ use_parentheses = true
 ensure_newline_before_comments = true
 
 [tool.mypy]
-python_version = "3.10"
+python_version = "3.9"
 warn_return_any = true
 warn_unused_configs = true
 disallow_untyped_defs = true
@@ -180,7 +182,7 @@ exclude_lines = [
 ]
 
 [tool.ruff]
-target-version = "py310"
+target-version = "py39"
 line-length = 88
 
 [tool.ruff.lint]

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

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

{shared_tensor-0.2.6 → 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,9 @@ class TaskInfo:
 class _TaskEntry:
     info: TaskInfo
     future: Future[Any]
+    result_encoding: str | None = None
+    result_payload: bytes | None = None
+    local_result: Any = None
 
 
 class TaskManager:
@@ -86,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,
@@ -97,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")
@@ -139,13 +141,14 @@ class TaskManager:
             )
             return
 
+        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
 
@@ -165,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
@@ -191,6 +197,31 @@ 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 _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:
@@ -207,6 +238,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,
@@ -227,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:
@@ -237,11 +293,26 @@ 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"),
         }
 
+    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:
@@ -266,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.6 → shared_tensor-0.2.8}/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))