shared-tensor 0.2.12__tar.gz → 0.2.15__tar.gz

{shared_tensor-0.2.12 → shared_tensor-0.2.15}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: shared-tensor
-Version: 0.2.12
+Version: 0.2.15
 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>
@@ -30,7 +30,6 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: cloudpickle>=3.0.0
 Requires-Dist: numpy<2
-Requires-Dist: torch>=2.1
 Provides-Extra: dev
 Requires-Dist: pytest>=6.0; extra == "dev"
 Requires-Dist: pytest-cov>=2.0; extra == "dev"
@@ -78,9 +77,10 @@ Not supported:
 
 ## Install
 
-Use Python `3.9+` and a CUDA-enabled PyTorch build.
+Use Python `3.9+`. Install a compatible PyTorch build first, then install `shared-tensor`.
 
 ```bash
+pip install torch
 pip install shared-tensor
 ```
 
@@ -92,6 +92,8 @@ conda activate shared-tensor-dev
 pip install -e ".[dev,test]"
 ```
 
+If you want to share Hugging Face `transformers` models, install both `torch` and `transformers` in the server and client environments. `shared-tensor` no longer installs `torch` for you.
+
 ## Docs
 
 Read the examples first, then the design notes:
@@ -150,6 +152,34 @@ executes endpoint functions         reopens CUDA objects via torch IPC
 manages cache and refcounts         releases managed handles explicitly
 ```
 
+## Example: Transformers Models
+
+`shared_tensor` also supports CUDA `transformers.PreTrainedModel` instances.
+
+See:
+- `examples/transformers_two_proc_demo.py`: minimal same-code two-process demo using `AutoModel`-style loading
+- `examples/transformers_mutation_check.py`: proves client-side in-place parameter mutation is visible on the server
+- `examples/transformers_ipc_benchmark.py`: measures reopen latency and client GPU memory delta
+
+Usage:
+
+```bash
+TRANSFORMERS_MODEL_ROOT=/path/to/model-or-hf-cache/models--bert-base-uncased \
+SHARED_TENSOR_ENABLED=1 SHARED_TENSOR_ROLE=server \
+python examples/transformers_two_proc_demo.py
+
+TRANSFORMERS_MODEL_ROOT=/path/to/model-or-hf-cache/models--bert-base-uncased \
+SHARED_TENSOR_ENABLED=1 \
+python examples/transformers_two_proc_demo.py
+```
+
+Notes:
+- `TRANSFORMERS_MODEL_ROOT` may point either to a resolved local model directory or to a Hugging Face cache root like `models--...`; the example resolves the newest snapshot automatically
+- `TRANSFORMERS_AUTO_CLASS` defaults to `AutoModel` and can be overridden to another `Auto*` class that exposes `from_pretrained`
+- for custom `transformers` code paths, the library stages the required module source files before reopening the shared module on the client
+- transport remains same-host same-GPU torch CUDA IPC; the client should not allocate a second full model copy just to reconstruct parameters
+- in a fresh client Python process, the first reopen may still look slow because `transformers` import/module resolution is often much slower than the shared-tensor IPC restore path itself; a second reopen in the same process should be much faster
+
 ## Lifetime And Failure Contract
 
 `shared_tensor` follows native PyTorch CUDA IPC semantics. It does not virtualize or harden producer lifetime.
@@ -346,6 +376,8 @@ Use call-level invalidation when you want to evict one cache key.
 Use endpoint-level invalidation when you want to drop all cached variants for the endpoint.
 Invalidation removes cache lookup entries; it does not guarantee that already-issued client handles remain valid after producer death.
 
+For cached `transformers` model endpoints, keep `cache=True` unless you explicitly want every request to rebuild and re-share the model.
+
 
 ## Handle Health Checks
 
@@ -371,6 +403,22 @@ info = provider.get_runtime_info()
 # execution_mode, server_socket_path, server_running, server_ready, server_info...
 ```
 
+## Client Retry And Timeout Defaults
+
+The client now retries initial connection setup for up to `60s` when the server socket is not ready yet, covering the common server-startup race where the client starts slightly earlier.
+
+Default request timeout is now `600s` for:
+- `SharedTensorClient`
+- `AsyncSharedTensorClient`
+- `SharedTensorProvider`
+
+You can still override these per instance:
+
+```python
+client = SharedTensorClient(timeout=120.0)
+provider = SharedTensorProvider(timeout=120.0)
+```
+
 ## Testing
 
 Default suite:

{shared_tensor-0.2.12 → shared_tensor-0.2.15}/README.md

@@ -25,9 +25,10 @@ Not supported:
 
 ## Install
 
-Use Python `3.9+` and a CUDA-enabled PyTorch build.
+Use Python `3.9+`. Install a compatible PyTorch build first, then install `shared-tensor`.
 
 ```bash
+pip install torch
 pip install shared-tensor
 ```
 
@@ -39,6 +40,8 @@ conda activate shared-tensor-dev
 pip install -e ".[dev,test]"
 ```
 
+If you want to share Hugging Face `transformers` models, install both `torch` and `transformers` in the server and client environments. `shared-tensor` no longer installs `torch` for you.
+
 ## Docs
 
 Read the examples first, then the design notes:
@@ -97,6 +100,34 @@ executes endpoint functions         reopens CUDA objects via torch IPC
 manages cache and refcounts         releases managed handles explicitly
 ```
 
+## Example: Transformers Models
+
+`shared_tensor` also supports CUDA `transformers.PreTrainedModel` instances.
+
+See:
+- `examples/transformers_two_proc_demo.py`: minimal same-code two-process demo using `AutoModel`-style loading
+- `examples/transformers_mutation_check.py`: proves client-side in-place parameter mutation is visible on the server
+- `examples/transformers_ipc_benchmark.py`: measures reopen latency and client GPU memory delta
+
+Usage:
+
+```bash
+TRANSFORMERS_MODEL_ROOT=/path/to/model-or-hf-cache/models--bert-base-uncased \
+SHARED_TENSOR_ENABLED=1 SHARED_TENSOR_ROLE=server \
+python examples/transformers_two_proc_demo.py
+
+TRANSFORMERS_MODEL_ROOT=/path/to/model-or-hf-cache/models--bert-base-uncased \
+SHARED_TENSOR_ENABLED=1 \
+python examples/transformers_two_proc_demo.py
+```
+
+Notes:
+- `TRANSFORMERS_MODEL_ROOT` may point either to a resolved local model directory or to a Hugging Face cache root like `models--...`; the example resolves the newest snapshot automatically
+- `TRANSFORMERS_AUTO_CLASS` defaults to `AutoModel` and can be overridden to another `Auto*` class that exposes `from_pretrained`
+- for custom `transformers` code paths, the library stages the required module source files before reopening the shared module on the client
+- transport remains same-host same-GPU torch CUDA IPC; the client should not allocate a second full model copy just to reconstruct parameters
+- in a fresh client Python process, the first reopen may still look slow because `transformers` import/module resolution is often much slower than the shared-tensor IPC restore path itself; a second reopen in the same process should be much faster
+
 ## Lifetime And Failure Contract
 
 `shared_tensor` follows native PyTorch CUDA IPC semantics. It does not virtualize or harden producer lifetime.
@@ -293,6 +324,8 @@ Use call-level invalidation when you want to evict one cache key.
 Use endpoint-level invalidation when you want to drop all cached variants for the endpoint.
 Invalidation removes cache lookup entries; it does not guarantee that already-issued client handles remain valid after producer death.
 
+For cached `transformers` model endpoints, keep `cache=True` unless you explicitly want every request to rebuild and re-share the model.
+
 
 ## Handle Health Checks
 
@@ -318,6 +351,22 @@ info = provider.get_runtime_info()
 # execution_mode, server_socket_path, server_running, server_ready, server_info...
 ```
 
+## Client Retry And Timeout Defaults
+
+The client now retries initial connection setup for up to `60s` when the server socket is not ready yet, covering the common server-startup race where the client starts slightly earlier.
+
+Default request timeout is now `600s` for:
+- `SharedTensorClient`
+- `AsyncSharedTensorClient`
+- `SharedTensorProvider`
+
+You can still override these per instance:
+
+```python
+client = SharedTensorClient(timeout=120.0)
+provider = SharedTensorProvider(timeout=120.0)
+```
+
 ## Testing
 
 Default suite:

{shared_tensor-0.2.12 → shared_tensor-0.2.15}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "shared-tensor"
-version = "0.2.12"
+version = "0.2.15"
 description = "Native PyTorch CUDA IPC over Unix Domain Socket for same-host process separation"
 readme = "README.md"
 license = "MIT"
@@ -47,7 +47,6 @@ requires-python = ">=3.9"
 dependencies = [
     "cloudpickle>=3.0.0",
     "numpy<2",
-    "torch>=2.1",
 ]
 
 [project.optional-dependencies]

{shared_tensor-0.2.12 → shared_tensor-0.2.15}/shared_tensor/__init__.py

@@ -19,4 +19,4 @@ __all__ = [
     "TaskStatus",
 ]
 
-__version__ = "0.2.12"
+__version__ = "0.2.15"

{shared_tensor-0.2.12 → shared_tensor-0.2.15}/shared_tensor/async_client.py

@@ -24,7 +24,7 @@ class AsyncSharedTensorClient:
         poll_interval: float = 1.0,
         *,
         device_index: int | None = None,
-        timeout: float = 30.0,
+        timeout: float = 600.0,
     ) -> None:
         self.poll_interval = poll_interval
         self.timeout = timeout

{shared_tensor-0.2.12 → shared_tensor-0.2.15}/shared_tensor/client.py

@@ -3,7 +3,9 @@
 from __future__ import annotations
 
 import logging
+import errno
 import socket
+import time
 from dataclasses import dataclass
 from typing import Any, cast
 
@@ -32,6 +34,9 @@ from shared_tensor.utils import (
 
 logger = logging.getLogger(__name__)
 
+CONNECT_RETRY_TIMEOUT_SECONDS = 60.0
+CONNECT_RETRY_INTERVAL_SECONDS = 0.5
+
 
 @dataclass(slots=True)
 class _ClientReleaser(ReleaseHandle):
@@ -53,7 +58,7 @@ class SharedTensorClient:
         base_path: str = "/tmp/shared-tensor",
         *,
         device_index: int | None = None,
-        timeout: float = 30.0,
+        timeout: float = 600.0,
         verbose_debug: bool = False,
     ) -> None:
         self.base_path = base_path
@@ -117,12 +122,35 @@ class SharedTensorClient:
         method = request.get("method", "<unknown>")
         if self.verbose_debug:
             logger.debug("Client sending request", extra={"method": method, "socket_path": self.socket_path})
+        started_at = time.monotonic()
         try:
-            with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as sock:
-                sock.settimeout(self.timeout)
-                sock.connect(self.socket_path)
-                send_message(sock, request)
-                response = recv_message(sock)
+            while True:
+                try:
+                    with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as sock:
+                        sock.settimeout(self.timeout)
+                        sock.connect(self.socket_path)
+                        send_message(sock, request)
+                        response = recv_message(sock)
+                    break
+                except OSError as exc:
+                    retryable = exc.errno in {
+                        errno.ENOENT,
+                        errno.ECONNREFUSED,
+                    }
+                    elapsed = time.monotonic() - started_at
+                    if not retryable or elapsed >= CONNECT_RETRY_TIMEOUT_SECONDS:
+                        raise
+                    if self.verbose_debug:
+                        logger.info(
+                            "Client retrying request connection",
+                            extra={
+                                "method": method,
+                                "socket_path": self.socket_path,
+                                "elapsed": round(elapsed, 3),
+                                "error": str(exc),
+                            },
+                        )
+                    time.sleep(CONNECT_RETRY_INTERVAL_SECONDS)
         except TimeoutError as exc:
             if self.verbose_debug:
                 logger.warning("Client request timed out", extra={"method": method, "socket_path": self.socket_path})

{shared_tensor-0.2.12 → shared_tensor-0.2.15}/shared_tensor/provider.py

@@ -91,7 +91,7 @@ class SharedTensorProvider:
         *,
         enabled: bool | None = None,
         device_index: int | None = None,
-        timeout: float = 30.0,
+        timeout: float = 600.0,
         execution_mode: str = "auto",
         server_startup_timeout: float = 30.0,
         verbose_debug: bool = False,

{shared_tensor-0.2.12 → shared_tensor-0.2.15}/shared_tensor/utils.py

@@ -9,6 +9,8 @@ import io
 import multiprocessing.reduction as mp_reduction
 import os
 import pickle
+import sys
+import tempfile
 from pathlib import Path
 from collections.abc import Callable
 from typing import Any, cast
@@ -32,6 +34,7 @@ except ImportError:  # pragma: no cover - exercised when transformers is not ins
 
 TORCH_MODULE: Any | None = _torch
 TRANSFORMERS_PRETRAINED_MODEL: Any | None = _transformers_pretrained_model
+_TRANSFORMERS_CLASS_CACHE: dict[tuple[str, str], Any] = {}
 
 
 CONTROL_ENCODING = "control_pickle"
@@ -174,16 +177,21 @@ def _serialize_transformers_model(module: Any) -> bytes:
     config = getattr(module, "config", None)
     if config is None:
         raise SharedTensorSerializationError("transformers model is missing a config")
-    if not hasattr(config, "to_dict"):
-        raise SharedTensorSerializationError("transformers model config must provide to_dict()")
+    module_sources = _collect_transformers_module_sources(type(module).__module__, type(config).__module__)
+    parameters = dict(module.named_parameters(remove_duplicate=False))
+    buffers = dict(module.named_buffers(remove_duplicate=False))
     payload = {
+        "module_sources": module_sources,
+        "config_bytes": pickle.dumps(config, protocol=pickle.HIGHEST_PROTOCOL),
         "model_module": type(module).__module__,
-        "model_name": type(module).__name__,
-        "config_module": type(config).__module__,
-        "config_name": type(config).__name__,
-        "config_dict": config.to_dict(),
-        "training": bool(module.training),
-        "state_bytes": _torch_serialize(module.state_dict()),
+        "model_qualname": type(module).__qualname__,
+        "state_bytes": _torch_serialize(
+            {
+                "parameters": parameters,
+                "buffers": buffers,
+            }
+        ),
+        "training": bool(getattr(module, "training", False)),
     }
     return pickle.dumps(payload, protocol=pickle.HIGHEST_PROTOCOL)
 
@@ -240,37 +248,200 @@ def serialize_call_payloads(
         raise SharedTensorSerializationError(f"Failed to serialize call payloads: {exc}") from exc
 
 
-def _load_module_state_dict(module: Any, state_dict: dict[str, Any]) -> None:
-    load_state_dict = module.load_state_dict
-    if "assign" in inspect.signature(load_state_dict).parameters:
-        incompatible = load_state_dict(state_dict, strict=False, assign=True)
-    else:  # pragma: no cover - compatibility fallback
-        incompatible = load_state_dict(state_dict, strict=False)
-    missing = list(getattr(incompatible, "missing_keys", []))
-    unexpected = list(getattr(incompatible, "unexpected_keys", []))
-    if missing or unexpected:
+def _bundle_module_files(module_name: str) -> dict[str, str] | None:
+    module = sys.modules.get(module_name)
+    if module is None:
+        module = importlib.import_module(module_name)
+    module_file = getattr(module, "__file__", None)
+    if not module_file:
+        return None
+    module_path = Path(module_file).resolve()
+    if module_path.suffix != ".py":
+        return None
+    transformers_root = None
+    if TRANSFORMERS_PRETRAINED_MODEL is not None:
+        transformers_file = getattr(sys.modules.get("transformers"), "__file__", None)
+        if transformers_file is not None:
+            transformers_root = Path(transformers_file).resolve().parent
+    if transformers_root is not None:
+        try:
+            module_path.relative_to(transformers_root)
+            if not module_name.startswith("transformers_modules."):
+                return None
+        except ValueError:
+            pass
+    files = {module_path.name: module_path.read_text(encoding="utf-8")}
+    get_relative_import_files = getattr(importlib.import_module("transformers.dynamic_module_utils"), "get_relative_import_files", None)
+    if callable(get_relative_import_files):
+        for needed in get_relative_import_files(str(module_path)):
+            needed_path = Path(needed).resolve()
+            files[needed_path.name] = needed_path.read_text(encoding="utf-8")
+    return files
+
+
+def _collect_transformers_module_sources(*module_names: str) -> dict[str, dict[str, str]]:
+    bundles: dict[str, dict[str, str]] = {}
+    for module_name in module_names:
+        bundle = _bundle_module_files(module_name)
+        if bundle:
+            bundles[module_name] = bundle
+    return bundles
+
+
+def _stage_transformers_module_sources(module_sources: dict[str, dict[str, str]]) -> None:
+    if not module_sources:
+        return
+    staging_root = Path(tempfile.gettempdir()) / "shared-tensor-transformers-modules"
+    staging_root.mkdir(parents=True, exist_ok=True)
+    root_str = str(staging_root)
+    if root_str not in sys.path:
+        sys.path.insert(0, root_str)
+    for module_name, files in module_sources.items():
+        parts = module_name.split(".")
+        package_dir = staging_root.joinpath(*parts[:-1])
+        package_dir.mkdir(parents=True, exist_ok=True)
+        current = staging_root
+        for part in parts[:-1]:
+            current = current / part
+            init_file = current / "__init__.py"
+            if not init_file.exists():
+                init_file.write_text("", encoding="utf-8")
+        for filename, content in files.items():
+            (package_dir / filename).write_text(content, encoding="utf-8")
+    importlib.invalidate_caches()
+
+
+def _import_transformers_module(module_name: str, module_sources: dict[str, dict[str, str]]) -> Any:
+    try:
+        return importlib.import_module(module_name)
+    except ModuleNotFoundError:
+        _stage_transformers_module_sources(module_sources)
+        return importlib.import_module(module_name)
+
+
+def _resolve_module_qualname(module_name: str, qualname: str) -> Any:
+    cache_key = (module_name, qualname)
+    cached = _TRANSFORMERS_CLASS_CACHE.get(cache_key)
+    if cached is not None:
+        return cached
+    module = importlib.import_module(module_name)
+    current: Any = module
+    for part in qualname.split("."):
+        current = getattr(current, part)
+    _TRANSFORMERS_CLASS_CACHE[cache_key] = current
+    return current
+
+
+def _build_transformers_meta_shell(model_cls: Any, config: Any) -> Any:
+    if TORCH_MODULE is None:
+        raise SharedTensorSerializationError("PyTorch is required for transformers deserialization")
+    try:
+        with TORCH_MODULE.device("meta"):
+            model = model_cls(config)
+    except Exception as exc:
         raise SharedTensorSerializationError(
-            f"Failed to reconstruct module state: missing={missing[:10]} unexpected={unexpected[:10]}"
+            f"Failed to construct transformers model shell on meta device: {exc}"
+        ) from exc
+    return model
+
+
+def _resolve_parent_module(root: Any, parent_name: str) -> Any:
+    if not parent_name:
+        return root
+    get_submodule = getattr(root, "get_submodule", None)
+    if callable(get_submodule):
+        return get_submodule(parent_name)
+    current = root
+    for part in parent_name.split("."):
+        current = getattr(current, part)
+    return current
+
+
+def _rebind_named_parameter(module: Any, name: str, value: Any) -> None:
+    parent_name, _, local_name = name.rpartition(".")
+    parent = _resolve_parent_module(module, parent_name)
+    parent._parameters[local_name] = value
+
+
+def _rebind_named_buffer(module: Any, name: str, value: Any) -> None:
+    parent_name, _, local_name = name.rpartition(".")
+    parent = _resolve_parent_module(module, parent_name)
+    parent._buffers[local_name] = value
+
+
+def _validate_transformers_state_layout(module: Any, parameters: dict[str, Any], buffers: dict[str, Any]) -> None:
+    expected_parameters = dict(module.named_parameters(remove_duplicate=False))
+    expected_buffers = dict(module.named_buffers(remove_duplicate=False))
+    if set(expected_parameters) != set(parameters):
+        raise SharedTensorSerializationError(
+            "Transformers parameter layout mismatch between serialized model and client shell"
+        )
+    if set(expected_buffers) != set(buffers):
+        raise SharedTensorSerializationError(
+            "Transformers buffer layout mismatch between serialized model and client shell"
         )
+    for name, expected in expected_parameters.items():
+        actual = parameters[name]
+        if not isinstance(actual, TORCH_MODULE.nn.Parameter):
+            raise SharedTensorSerializationError(
+                f"Transformers parameter '{name}' payload is not a torch.nn.Parameter"
+            )
+        if actual.shape != expected.shape or actual.dtype != expected.dtype:
+            raise SharedTensorSerializationError(
+                f"Transformers parameter '{name}' shape or dtype does not match the client shell"
+            )
+    for name, expected in expected_buffers.items():
+        actual = buffers[name]
+        if not isinstance(actual, TORCH_MODULE.Tensor):
+            raise SharedTensorSerializationError(
+                f"Transformers buffer '{name}' payload is not a torch.Tensor"
+            )
+        if actual.shape != expected.shape or actual.dtype != expected.dtype:
+            raise SharedTensorSerializationError(
+                f"Transformers buffer '{name}' shape or dtype does not match the client shell"
+            )
+
+
+def _restore_transformers_state(module: Any, state: dict[str, Any]) -> Any:
+    parameters = cast(dict[str, Any], state.get("parameters", {}))
+    buffers = cast(dict[str, Any], state.get("buffers", {}))
+    _validate_transformers_state_layout(module, parameters, buffers)
+    for name, value in parameters.items():
+        _rebind_named_parameter(module, name, value)
+    for name, value in buffers.items():
+        _rebind_named_buffer(module, name, value)
+    return module
 
 
 def _deserialize_transformers_model(payload: bytes) -> Any:
     encoded = pickle.loads(payload)
-    model_module = importlib.import_module(encoded["model_module"])
-    model_cls = getattr(model_module, encoded["model_name"])
-    config_module = importlib.import_module(encoded["config_module"])
-    config_cls = getattr(config_module, encoded["config_name"])
-    if hasattr(config_cls, "from_dict"):
-        config = config_cls.from_dict(encoded["config_dict"])
-    else:  # pragma: no cover - defensive
-        config = config_cls(**encoded["config_dict"])
-    state_dict = pickle.loads(encoded["state_bytes"])
-    model = model_cls(config)
-    first_tensor = next(iter(state_dict.values()), None)
-    if TORCH_MODULE is not None and isinstance(first_tensor, TORCH_MODULE.Tensor):
-        model = model.to(first_tensor.device)
-    _load_module_state_dict(model, state_dict)
-    model.train(bool(encoded["training"]))
+    module_sources = cast(dict[str, dict[str, str]], encoded.get("module_sources", {}))
+    for module_name in module_sources:
+        _import_transformers_module(module_name, module_sources)
+    config_bytes = encoded.get("config_bytes")
+    state_bytes = encoded.get("state_bytes")
+    model_module = encoded.get("model_module")
+    model_qualname = encoded.get("model_qualname")
+    if not isinstance(config_bytes, (bytes, bytearray)):
+        raise SharedTensorSerializationError("transformers config payload is invalid")
+    if not isinstance(state_bytes, (bytes, bytearray)):
+        raise SharedTensorSerializationError("transformers state payload is invalid")
+    if not isinstance(model_module, str) or not model_module:
+        raise SharedTensorSerializationError("transformers model module is invalid")
+    if not isinstance(model_qualname, str) or not model_qualname:
+        raise SharedTensorSerializationError("transformers model qualname is invalid")
+    config = pickle.loads(config_bytes)
+    model_cls = _resolve_module_qualname(model_module, model_qualname)
+    model = _build_transformers_meta_shell(model_cls, config)
+    state = pickle.loads(state_bytes)
+    model = _restore_transformers_state(model, cast(dict[str, Any], state))
+    train_fn = getattr(model, "train", None)
+    training = bool(encoded.get("training", False))
+    if callable(train_fn):
+        train_fn(training)
+    eval_fn = getattr(model, "eval", None)
+    if callable(eval_fn) and not training:
+        eval_fn()
     _validate_module_device(model)
     return model