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

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

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

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

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

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

@@ -34,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"
@@ -177,9 +178,20 @@ def _serialize_transformers_model(module: Any) -> bytes:
     if config is None:
         raise SharedTensorSerializationError("transformers model is missing a config")
     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,
-        "module_bytes": _torch_serialize(module),
+        "config_bytes": pickle.dumps(config, protocol=pickle.HIGHEST_PROTOCOL),
+        "model_module": type(module).__module__,
+        "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)
 
@@ -307,17 +319,128 @@ def _import_transformers_module(module_name: str, module_sources: dict[str, dict
         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 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)
     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)
-    module_bytes = encoded.get("module_bytes")
-    if not isinstance(module_bytes, (bytes, bytearray)):
-        raise SharedTensorSerializationError("transformers module payload is invalid")
-    model = pickle.loads(module_bytes)
+    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):
+    if callable(eval_fn) and not training:
         eval_fn()
     _validate_module_device(model)
     return model