shared-tensor 0.1.2__tar.gz → 0.2.1__tar.gz

shared_tensor-0.2.1/PKG-INFO

@@ -0,0 +1,177 @@
+Metadata-Version: 2.4
+Name: shared-tensor
+Version: 0.2.1
+Summary: Local endpoint-oriented RPC for same-host same-GPU PyTorch IPC
+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: 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
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: Operating System :: POSIX :: Linux
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Classifier: Topic :: System :: Distributed Computing
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy<2
+Requires-Dist: requests>=2.25.0
+Requires-Dist: torch>=2.2.0
+Provides-Extra: dev
+Requires-Dist: pytest>=6.0; extra == "dev"
+Requires-Dist: pytest-cov>=2.0; extra == "dev"
+Requires-Dist: types-requests>=2.32.0; extra == "dev"
+Requires-Dist: black>=22.0; extra == "dev"
+Requires-Dist: isort>=5.0; extra == "dev"
+Requires-Dist: mypy>=0.950; extra == "dev"
+Requires-Dist: pre-commit>=2.0.0; extra == "dev"
+Requires-Dist: build>=0.8.0; extra == "dev"
+Requires-Dist: twine>=4.0.0; extra == "dev"
+Requires-Dist: ruff>=0.6.0; extra == "dev"
+Provides-Extra: test
+Requires-Dist: pytest>=6.0; extra == "test"
+Requires-Dist: pytest-cov>=2.0; extra == "test"
+Requires-Dist: pytest-asyncio>=0.20.0; extra == "test"
+Provides-Extra: docs
+Requires-Dist: sphinx>=4.0.0; extra == "docs"
+Requires-Dist: sphinx-rtd-theme>=1.0.0; extra == "docs"
+Requires-Dist: myst-parser>=0.18.0; extra == "docs"
+Dynamic: license-file
+
+# Shared Tensor
+
+`shared_tensor` is a localhost-only RPC layer for one thing: passing CUDA `torch.Tensor` and CUDA `torch.nn.Module` objects between processes on the same machine and the same GPU with native PyTorch IPC semantics.
+
+## What It Supports
+
+- same-host, trusted-process deployment
+- same-GPU CUDA object handoff
+- native `torch` tensors and modules
+- explicit endpoint registration
+- sync calls and async task polling
+
+## What It Does Not Support
+
+- CPU tensor transport
+- generic Python object RPC
+- cross-machine transport
+- macOS `mps`
+- silent CPU fallback or implicit device migration
+
+## Payload Contract
+
+Allowed payloads:
+
+- CUDA `torch.Tensor`
+- CUDA `torch.nn.Module`
+- `tuple`, `list`, and `dict[str, ...]` containers built from those values for `args` and `kwargs`
+- empty `args` / `kwargs` through the control path for no-argument calls only
+
+Rejected payloads:
+
+- CPU tensors and CPU modules
+- plain Python values such as `int`, `str`, `dict`, and `list`
+- `mps` tensors and modules
+
+## Install
+
+Use Python `3.10+` and a CUDA-enabled PyTorch build.
+
+```bash
+pip install shared-tensor
+```
+
+For development:
+
+```bash
+conda create -y -n shared-tensor-dev python=3.11
+conda activate shared-tensor-dev
+pip install -e ".[dev,test]"
+```
+
+## Typical Example
+
+Provider process:
+
+```python
+import torch
+
+from shared_tensor import SharedTensorProvider
+
+provider = SharedTensorProvider(execution_mode="server")
+
+
+@provider.share(name="load_model")
+def load_model() -> torch.nn.Module:
+    return torch.nn.Linear(4, 2, device="cuda")
+
+
+@provider.share(name="identity")
+def identity(tensor: torch.Tensor) -> torch.Tensor:
+    return tensor
+```
+
+Run the server:
+
+```bash
+shared-tensor-server --provider my_service:provider --host 127.0.0.1 --port 2537
+```
+
+Consumer process:
+
+```python
+import torch
+
+from shared_tensor import SharedTensorClient
+
+with SharedTensorClient(port=2537) as client:
+    model = client.call("load_model")
+    x = torch.ones(1, 4, device="cuda")
+    y = model(x)
+
+    shared = client.call("identity", x)
+```
+
+## Test Matrix
+
+Default local run:
+
+```bash
+python -m pytest -m "not gpu"
+```
+
+CUDA run:
+
+```bash
+python -m pytest -m gpu
+```
+
+`skipped` means the test was intentionally not run because its precondition was missing. In this repo that usually means a `gpu` test was executed on a machine where `torch.cuda.is_available()` was false. It is not a failure.
+
+Current validation target:
+
+- local non-GPU suite passes
+- H100 CUDA suite passes
+
+## Operational Notes
+
+- This library assumes a trusted same-host environment.
+- The server process must be a separate process from the client when using CUDA IPC.
+- If you need cross-machine transport or CPU object RPC, use a different tool.
+
+## Repo Notes
+
+- `CLAUDE.md` captures repo maintenance rules.
+- `examples/basic_service.py` shows the minimal sync flow.
+- `examples/model_service.py` shows model handoff.

shared_tensor-0.2.1/README.md

@@ -0,0 +1,126 @@
+# Shared Tensor
+
+`shared_tensor` is a localhost-only RPC layer for one thing: passing CUDA `torch.Tensor` and CUDA `torch.nn.Module` objects between processes on the same machine and the same GPU with native PyTorch IPC semantics.
+
+## What It Supports
+
+- same-host, trusted-process deployment
+- same-GPU CUDA object handoff
+- native `torch` tensors and modules
+- explicit endpoint registration
+- sync calls and async task polling
+
+## What It Does Not Support
+
+- CPU tensor transport
+- generic Python object RPC
+- cross-machine transport
+- macOS `mps`
+- silent CPU fallback or implicit device migration
+
+## Payload Contract
+
+Allowed payloads:
+
+- CUDA `torch.Tensor`
+- CUDA `torch.nn.Module`
+- `tuple`, `list`, and `dict[str, ...]` containers built from those values for `args` and `kwargs`
+- empty `args` / `kwargs` through the control path for no-argument calls only
+
+Rejected payloads:
+
+- CPU tensors and CPU modules
+- plain Python values such as `int`, `str`, `dict`, and `list`
+- `mps` tensors and modules
+
+## Install
+
+Use Python `3.10+` and a CUDA-enabled PyTorch build.
+
+```bash
+pip install shared-tensor
+```
+
+For development:
+
+```bash
+conda create -y -n shared-tensor-dev python=3.11
+conda activate shared-tensor-dev
+pip install -e ".[dev,test]"
+```
+
+## Typical Example
+
+Provider process:
+
+```python
+import torch
+
+from shared_tensor import SharedTensorProvider
+
+provider = SharedTensorProvider(execution_mode="server")
+
+
+@provider.share(name="load_model")
+def load_model() -> torch.nn.Module:
+    return torch.nn.Linear(4, 2, device="cuda")
+
+
+@provider.share(name="identity")
+def identity(tensor: torch.Tensor) -> torch.Tensor:
+    return tensor
+```
+
+Run the server:
+
+```bash
+shared-tensor-server --provider my_service:provider --host 127.0.0.1 --port 2537
+```
+
+Consumer process:
+
+```python
+import torch
+
+from shared_tensor import SharedTensorClient
+
+with SharedTensorClient(port=2537) as client:
+    model = client.call("load_model")
+    x = torch.ones(1, 4, device="cuda")
+    y = model(x)
+
+    shared = client.call("identity", x)
+```
+
+## Test Matrix
+
+Default local run:
+
+```bash
+python -m pytest -m "not gpu"
+```
+
+CUDA run:
+
+```bash
+python -m pytest -m gpu
+```
+
+`skipped` means the test was intentionally not run because its precondition was missing. In this repo that usually means a `gpu` test was executed on a machine where `torch.cuda.is_available()` was false. It is not a failure.
+
+Current validation target:
+
+- local non-GPU suite passes
+- H100 CUDA suite passes
+
+## Operational Notes
+
+- This library assumes a trusted same-host environment.
+- The server process must be a separate process from the client when using CUDA IPC.
+- If you need cross-machine transport or CPU object RPC, use a different tool.
+
+## Repo Notes
+
+- `CLAUDE.md` captures repo maintenance rules.
+- `examples/basic_service.py` shows the minimal sync flow.
+- `examples/model_service.py` shows model handoff.

{shared_tensor-0.1.2 → shared_tensor-0.2.1}/pyproject.toml

@@ -4,8 +4,8 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "shared-tensor"
-version = "0.1.2"
-description = "A library for sharing GPU memory objects across processes using IPC mechanisms"
+version = "0.2.1"
+description = "Local endpoint-oriented RPC for same-host same-GPU PyTorch IPC"
 readme = "README.md"
 license = "Apache-2.0"
 authors = [
@@ -16,16 +16,16 @@ maintainers = [
 ]
 keywords = [
     "gpu",
-    "memory", 
+    "memory",
     "sharing",
     "ipc",
     "inter-process-communication",
     "pytorch",
-    "tensorflow",
     "cuda",
     "model-serving",
     "inference",
-    "distributed-computing"
+    "torch",
+    "torch-ipc"
 ]
 classifiers = [
     "Development Status :: 3 - Alpha",
@@ -33,38 +33,36 @@ classifiers = [
     "Intended Audience :: Science/Research",
     "Operating System :: POSIX :: Linux",
     "Programming Language :: Python :: 3",
-    "Programming Language :: Python :: 3.8",
-    "Programming Language :: Python :: 3.9",
     "Programming Language :: Python :: 3.10",
     "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3.12",
     "Topic :: Scientific/Engineering :: Artificial Intelligence",
     "Topic :: Software Development :: Libraries :: Python Modules",
-    "Topic :: System :: Hardware :: Symmetric Multi-processing",
+    "Topic :: System :: Distributed Computing",
 ]
-requires-python = ">=3.8"
+requires-python = ">=3.10"
 dependencies = [
-    "torch>=1.12.0",
-    "numpy>=1.20.0",
+    "numpy<2",
     "requests>=2.25.0",
+    "torch>=2.2.0",
 ]
 
 [project.optional-dependencies]
 dev = [
     "pytest>=6.0",
     "pytest-cov>=2.0",
+    "types-requests>=2.32.0",
     "black>=22.0",
-    "flake8>=4.0",
     "isort>=5.0",
     "mypy>=0.950",
     "pre-commit>=2.0.0",
     "build>=0.8.0",
     "twine>=4.0.0",
+    "ruff>=0.6.0",
 ]
 test = [
     "pytest>=6.0",
     "pytest-cov>=2.0",
-    "pytest-benchmark>=3.0",
     "pytest-asyncio>=0.20.0",
 ]
 docs = [
@@ -93,10 +91,9 @@ exclude = ["tests*", "examples*", "docs*"]
 [tool.setuptools.package-data]
 shared_tensor = ["*.so", "*.dll", "*.dylib"]
 
-# Black configuration
 [tool.black]
 line-length = 88
-target-version = ['py38', 'py39', 'py310', 'py311']
+target-version = ['py310', 'py311', 'py312']
 include = '\.pyi?$'
 extend-exclude = '''
 /(
@@ -112,7 +109,6 @@ extend-exclude = '''
 )/
 '''
 
-# isort configuration
 [tool.isort]
 profile = "black"
 multi_line_output = 3
@@ -122,9 +118,8 @@ force_grid_wrap = 0
 use_parentheses = true
 ensure_newline_before_comments = true
 
-# mypy configuration
 [tool.mypy]
-python_version = "3.8"
+python_version = "3.10"
 warn_return_any = true
 warn_unused_configs = true
 disallow_untyped_defs = true
@@ -141,11 +136,10 @@ strict_equality = true
 [[tool.mypy.overrides]]
 module = [
     "torch.*",
-    "numpy.*",
+    "requests.*",
 ]
 ignore_missing_imports = true
 
-# pytest configuration
 [tool.pytest.ini_options]
 minversion = "6.0"
 addopts = [
@@ -168,7 +162,6 @@ markers = [
     "unit: marks tests as unit tests",
 ]
 
-# Coverage configuration
 [tool.coverage.run]
 source = ["shared_tensor"]
 omit = [
@@ -191,40 +184,17 @@ exclude_lines = [
     "@(abc\\.)?abstractmethod",
 ]
 
-# flake8 configuration (in setup.cfg format within pyproject.toml comments)
-# [flake8]
-# max-line-length = 88
-# extend-ignore = E203, E266, E501, W503
-# max-complexity = 10
-# select = B,C,E,F,W,T4,B9
-
-# Ruff configuration (modern alternative to flake8)
 [tool.ruff]
-target-version = "py38"
+target-version = "py310"
 line-length = 88
-select = [
-    "E",  # pycodestyle errors
-    "W",  # pycodestyle warnings
-    "F",  # pyflakes
-    "I",  # isort
-    "B",  # flake8-bugbear
-    "C4", # flake8-comprehensions
-    "UP", # pyupgrade
-]
-ignore = [
-    "E501",  # line too long, handled by black
-    "B008",  # do not perform function calls in argument defaults
-    "C901",  # too complex
-]
 
-[tool.ruff.per-file-ignores]
+[tool.ruff.lint]
+select = ["E", "W", "F", "I", "B", "C4", "UP"]
+ignore = ["E501", "B008", "C901"]
+
+[tool.ruff.lint.per-file-ignores]
 "__init__.py" = ["F401"]
 "tests/**/*" = ["B011", "B018"]
 
-[tool.ruff.isort]
+[tool.ruff.lint.isort]
 known-first-party = ["shared_tensor"]
-
-# Bandit security linter configuration
-[tool.bandit]
-exclude_dirs = ["tests", "examples"]
-skips = ["B101", "B601"]

{shared_tensor-0.1.2 → shared_tensor-0.2.1}/shared_tensor/__init__.py

@@ -1,27 +1,20 @@
-"""
-Shared Tensor Library
+"""shared_tensor: local endpoint-oriented RPC for Python and PyTorch."""
 
-A library for sharing GPU memory objects across processes using IPC mechanisms.
-Enables model and inference engine separation architecture using JSON-RPC 2.0 protocol.
-"""
-
-from shared_tensor.provider import SharedTensorProvider
+from shared_tensor.async_client import AsyncSharedTensorClient
+from shared_tensor.async_provider import AsyncSharedTensorProvider
+from shared_tensor.async_task import TaskInfo, TaskStatus
 from shared_tensor.client import SharedTensorClient
+from shared_tensor.provider import SharedTensorProvider
 from shared_tensor.server import SharedTensorServer
-from shared_tensor.async_provider import AsyncSharedTensorProvider
-from shared_tensor.async_client import AsyncSharedTensorClient
-from shared_tensor.async_task import TaskStatus, TaskInfo
 
-__version__ = "0.1.0"
-__author__ = "Athena Team"
-
-# Export main functionality
 __all__ = [
-    "SharedTensorProvider", 
+    "AsyncSharedTensorClient",
+    "AsyncSharedTensorProvider",
     "SharedTensorClient",
+    "SharedTensorProvider",
     "SharedTensorServer",
-    "AsyncSharedTensorProvider",
-    "AsyncSharedTensorClient", 
-    "TaskStatus",
     "TaskInfo",
-]
+    "TaskStatus",
+]
+
+__version__ = "0.2.1"

shared_tensor-0.2.1/shared_tensor/async_client.py

@@ -0,0 +1,166 @@
+"""Async task-oriented client facade built on top of :mod:`shared_tensor.client`."""
+
+from __future__ import annotations
+
+import time
+from collections.abc import Callable
+from typing import Any, cast
+
+from shared_tensor.async_task import TaskInfo, TaskStatus
+from shared_tensor.client import SharedTensorClient
+from shared_tensor.errors import SharedTensorTaskError
+from shared_tensor.utils import resolve_legacy_endpoint_name, serialize_call_payloads
+
+
+class AsyncSharedTensorClient:
+    def __init__(
+        self,
+        port: int = 2537,
+        verbose_debug: bool = False,
+        poll_interval: float = 1.0,
+        *,
+        host: str = "127.0.0.1",
+        timeout: float = 30.0,
+    ) -> None:
+        self.poll_interval = poll_interval
+        self._client = SharedTensorClient(
+            port=port,
+            host=host,
+            timeout=timeout,
+            verbose_debug=verbose_debug,
+        )
+
+    def submit(self, endpoint: str, *args: Any, **kwargs: Any) -> str:
+        encoding, args_payload, kwargs_payload = serialize_call_payloads(tuple(args), dict(kwargs))
+        result = self._client._request(
+            "submit",
+            {
+                "endpoint": endpoint,
+                "args_hex": args_payload.hex(),
+                "kwargs_hex": kwargs_payload.hex(),
+                "encoding": encoding,
+            },
+        )
+        return cast(str, result["task_id"])
+
+    def submit_task(
+        self,
+        function_path: str,
+        args: tuple[Any, ...] = (),
+        kwargs: dict[str, Any] | None = None,
+        options: dict[str, Any] | None = None,
+    ) -> str:
+        del options
+        endpoint = resolve_legacy_endpoint_name(function_path)
+        return self.submit(endpoint, *(args or ()), **(kwargs or {}))
+
+    def status(self, task_id: str) -> TaskInfo:
+        return TaskInfo.from_dict(self._client._request("get_task", {"task_id": task_id}))
+
+    def get_task_status(self, task_id: str) -> TaskInfo:
+        return self.status(task_id)
+
+    def result(self, task_id: str) -> Any:
+        result = self._client._request("get_task_result", {"task_id": task_id})
+        return SharedTensorClient._decode_rpc_payload(result)
+
+    def get_task_result(self, task_id: str) -> Any:
+        return self.result(task_id)
+
+    def wait(
+        self,
+        task_id: str,
+        timeout: float | None = None,
+        callback: Callable[[TaskInfo], None] | None = None,
+    ) -> Any:
+        started = time.time()
+        while True:
+            info = self.status(task_id)
+            if callback is not None:
+                callback(info)
+            if info.status == TaskStatus.COMPLETED:
+                return self.result(task_id)
+            if info.status == TaskStatus.FAILED:
+                raise SharedTensorTaskError(info.error_message or f"Task '{task_id}' failed")
+            if info.status == TaskStatus.CANCELLED:
+                raise SharedTensorTaskError(f"Task '{task_id}' was cancelled")
+            if timeout is not None and time.time() - started > timeout:
+                raise SharedTensorTaskError(
+                    f"Task '{task_id}' did not complete within {timeout} seconds"
+                )
+            time.sleep(self.poll_interval)
+
+    def wait_for_task(
+        self,
+        task_id: str,
+        timeout: float | None = None,
+        callback: Callable[[TaskInfo], None] | None = None,
+    ) -> Any:
+        return self.wait(task_id, timeout=timeout, callback=callback)
+
+    def execute_function_async(
+        self,
+        function_path: str,
+        args: tuple[Any, ...] = (),
+        kwargs: dict[str, Any] | None = None,
+        options: dict[str, Any] | None = None,
+        wait: bool = True,
+        timeout: float | None = None,
+        callback: Callable[[TaskInfo], None] | None = None,
+    ) -> Any:
+        del options
+        task_id = self.submit_task(function_path, args=args, kwargs=kwargs)
+        if not wait:
+            return task_id
+        return self.wait(task_id, timeout=timeout, callback=callback)
+
+    def cancel(self, task_id: str) -> bool:
+        return bool(self._client._request("cancel_task", {"task_id": task_id})["cancelled"])
+
+    def cancel_task(self, task_id: str) -> bool:
+        return self.cancel(task_id)
+
+    def list_tasks(self, status: str | None = None) -> dict[str, TaskInfo]:
+        params = {"status": status} if status else None
+        result = self._client._request("list_tasks", params)
+        return {task_id: TaskInfo.from_dict(data) for task_id, data in result.items()}
+
+    def close(self) -> None:
+        self._client.close()
+
+    def __enter__(self) -> AsyncSharedTensorClient:
+        return self
+
+    def __exit__(self, exc_type: object, exc_val: object, exc_tb: object) -> None:
+        self.close()
+
+
+def execute_remote_function_async(
+    function_path: str,
+    args: tuple[Any, ...] = (),
+    kwargs: dict[str, Any] | None = None,
+    options: dict[str, Any] | None = None,
+    *,
+    server_port: int = 2537,
+    host: str = "127.0.0.1",
+    verbose_debug: bool = False,
+    poll_interval: float = 1.0,
+    wait: bool = True,
+    timeout: float | None = None,
+    callback: Callable[[TaskInfo], None] | None = None,
+) -> Any:
+    with AsyncSharedTensorClient(
+        port=server_port,
+        host=host,
+        verbose_debug=verbose_debug,
+        poll_interval=poll_interval,
+    ) as client:
+        return client.execute_function_async(
+            function_path,
+            args=args,
+            kwargs=kwargs,
+            options=options,
+            wait=wait,
+            timeout=timeout,
+            callback=callback,
+        )