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

shared_tensor-0.2.12/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Athena Team
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

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

@@ -1,10 +1,10 @@
 Metadata-Version: 2.4
 Name: shared-tensor
-Version: 0.2.10
+Version: 0.2.12
 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
+License-Expression: MIT
 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/tree/main/docs

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

@@ -4,10 +4,10 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "shared-tensor"
-version = "0.2.10"
+version = "0.2.12"
 description = "Native PyTorch CUDA IPC over Unix Domain Socket for same-host process separation"
 readme = "README.md"
-license = "Apache-2.0"
+license = "MIT"
 authors = [
     {name = "Athena Team", email = "contact@world-sim-dev.org"}
 ]

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

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

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

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import hashlib
+import importlib
 import inspect
 import io
 import multiprocessing.reduction as mp_reduction
@@ -23,12 +24,19 @@ try:
 except ImportError:  # pragma: no cover - exercised only without torch installed.
     _torch = cast(Any, None)
 
+try:
+    from transformers import PreTrainedModel as _transformers_pretrained_model
+except ImportError:  # pragma: no cover - exercised when transformers is not installed.
+    _transformers_pretrained_model = cast(Any, None)
+
 
 TORCH_MODULE: Any | None = _torch
+TRANSFORMERS_PRETRAINED_MODEL: Any | None = _transformers_pretrained_model
 
 
 CONTROL_ENCODING = "control_pickle"
 TORCH_ENCODING = "torch_cuda_ipc"
+TRANSFORMERS_MODEL_ENCODING = "torch_transformers_pretrained_model"
 SHARED_TENSOR_BASE_PATH_ENV = "SHARED_TENSOR_BASE_PATH"
 DEFAULT_SOCKET_BASE_PATH = "/tmp/shared-tensor"
 _EMPTY_TUPLE = ()
@@ -162,17 +170,41 @@ def _torch_serialize(obj: Any) -> bytes:
     return buffer.getvalue()
 
 
-def _serialize_ipc_payload(obj: Any) -> bytes:
+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()")
+    payload = {
+        "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()),
+    }
+    return pickle.dumps(payload, protocol=pickle.HIGHEST_PROTOCOL)
+
+
+def _serialize_ipc_payload(obj: Any) -> tuple[str, bytes]:
+    if (
+        TORCH_MODULE is not None
+        and TRANSFORMERS_PRETRAINED_MODEL is not None
+        and isinstance(obj, TRANSFORMERS_PRETRAINED_MODEL)
+    ):
+        return TRANSFORMERS_MODEL_ENCODING, _serialize_transformers_model(obj)
     if TORCH_MODULE is not None and isinstance(obj, (TORCH_MODULE.Tensor, TORCH_MODULE.nn.Module)):
-        return _torch_serialize(obj)
-    return pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL)
+        return TORCH_ENCODING, _torch_serialize(obj)
+    return CONTROL_ENCODING, pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL)
 
 
 def serialize_payload(obj: Any) -> tuple[str, bytes]:
     """Serialize a CUDA torch payload using PyTorch's IPC-aware pickler."""
     try:
         _validate_torch_payload(obj, allow_dict_keys=isinstance(obj, dict))
-        return TORCH_ENCODING, _serialize_ipc_payload(obj)
+        return _serialize_ipc_payload(obj)
     except SharedTensorCapabilityError:
         raise
     except Exception as exc:  # pragma: no cover - defensive
@@ -208,6 +240,41 @@ 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:
+        raise SharedTensorSerializationError(
+            f"Failed to reconstruct module state: missing={missing[:10]} unexpected={unexpected[:10]}"
+        )
+
+
+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"]))
+    _validate_module_device(model)
+    return model
+
+
 def deserialize_payload(encoding: str, data: bytes | str) -> Any:
     """Deserialize a payload produced by :func:`serialize_payload`."""
     payload = bytes.fromhex(data) if isinstance(data, str) else data
@@ -216,6 +283,8 @@ def deserialize_payload(encoding: str, data: bytes | str) -> Any:
             return pickle.loads(payload)
         if encoding == TORCH_ENCODING:
             return pickle.loads(payload)
+        if encoding == TRANSFORMERS_MODEL_ENCODING:
+            return _deserialize_transformers_model(payload)
         raise SharedTensorSerializationError(f"Unsupported payload encoding: {encoding}")
     except SharedTensorCapabilityError:
         raise

shared_tensor-0.2.10/LICENSE

@@ -1,181 +0,0 @@
-Apache License
-Version 2.0, January 2004
-http://www.apache.org/licenses/
-
-TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
-
-1. Definitions.
-
-"License" shall mean the terms and conditions for use, reproduction,
-and distribution as defined by Sections 1 through 9 of this document.
-
-"Licensor" shall mean the copyright owner or entity granting the License.
-
-"Legal Entity" shall mean the union of the acting entity and all
-other entities that control, are controlled by, or are under common
-control with that entity. For the purposes of this definition,
-"control" means (i) the power, direct or indirect, to cause the
-direction or management of such entity, whether by contract or
-otherwise, or (ii) ownership of fifty percent (50%) or more of the
-outstanding shares, or (iii) beneficial ownership of such entity.
-
-"You" (or "Your") shall mean an individual or Legal Entity
-exercising permissions granted by this License.
-
-"Source" form shall mean the preferred form for making modifications,
-including but not limited to software source code, documentation
-source, and configuration files.
-
-"Object" form shall mean any form resulting from mechanical
-transformation or translation of a Source form, including but
-not limited to compiled object code, generated documentation,
-and conversions to other media types.
-
-"Work" shall mean the work of authorship, whether in Source or
-Object form, made available under the License, as indicated by a
-copyright notice that is included in or attached to the work
-(which shall not include communications that are clearly marked or
-otherwise designated in writing by the copyright owner as "Not a Contribution").
-
-"Contribution" shall mean any work of authorship, including
-the original version of the Work and any modifications or additions
-to that Work or Derivative Works thereof, that is intentionally
-submitted to Licensor for inclusion in the Work by the copyright owner
-or by an individual or Legal Entity authorized to submit on behalf of
-the copyright owner. For the purposes of this definition, "submitted"
-means any form of electronic, verbal, or written communication sent
-to the Licensor or its representatives, including but not limited to
-communication on electronic mailing lists, source code control
-systems, and issue tracking systems that are managed by, or on behalf of,
-the Licensor for the purpose of discussing and improving the Work, but
-excluding communication that is conspicuously marked or otherwise
-designated in writing by the copyright owner as "Not a Contribution."
-
-2. Grant of Copyright License. Subject to the terms and conditions of
-this License, each Contributor hereby grants to You a perpetual,
-worldwide, non-exclusive, no-charge, royalty-free, irrevocable
-copyright license to use, reproduce, modify, merge, publish,
-distribute, sublicense, and/or sell copies of the Work, and to
-permit persons to whom the Work is furnished to do so, subject to
-the following conditions:
-
-The above copyright notice and this permission notice shall be
-included in all copies or substantial portions of the Work.
-
-3. Grant of Patent License. Subject to the terms and conditions of
-this License, each Contributor hereby grants to You a perpetual,
-worldwide, non-exclusive, no-charge, royalty-free, irrevocable
-(except as stated in this section) patent license to make, have made,
-use, offer to sell, sell, import, and otherwise transfer the Work,
-where such license applies only to those patent claims licensable
-by such Contributor that are necessarily infringed by their
-Contribution(s) alone or by combination of their Contribution(s)
-with the Work to which such Contribution(s) was submitted. If You
-institute patent litigation against any entity (including a
-cross-claim or counterclaim in a lawsuit) alleging that the Work
-or a Contribution incorporated within the Work constitutes direct
-or contributory patent infringement, then any patent licenses
-granted to You under this License for that Work shall terminate
-as of the date such litigation is filed.
-
-4. Redistribution. You may reproduce and distribute copies of the
-Work or Derivative Works thereof in any medium, with or without
-modifications, and in Source or Object form, provided that You
-meet the following conditions:
-
-(a) You must give any other recipients of the Work or
-Derivative Works a copy of this License; and
-
-(b) You must cause any modified files to carry prominent notices
-stating that You changed the files; and
-
-(c) You must retain, in the Source form of any Derivative Works
-that You distribute, all copyright, trademark, patent, and
-attribution notices from the Source form of the Work,
-excluding those notices that do not pertain to any part of
-the Derivative Works; and
-
-(d) If the Work includes a "NOTICE" text file as part of its
-distribution, then any Derivative Works that You distribute must
-include a readable copy of the attribution notices contained
-within such NOTICE file, excluding those notices that do not
-pertain to any part of the Derivative Works, in at least one
-of the following places: within a NOTICE text file distributed
-as part of the Derivative Works; within the Source form or
-documentation, if provided along with the Derivative Works; or,
-within a display generated by the Derivative Works, if and
-wherever such third-party notices normally appear. The contents
-of the NOTICE file are for informational purposes only and
-do not modify the License. You may add Your own attribution
-notices within Derivative Works that You distribute, alongside
-or as an addendum to the NOTICE text from the Work, provided
-that such additional attribution notices cannot be construed
-as modifying the License.
-
-You may add Your own copyright notice to Your modifications and
-may provide additional or different license terms and conditions
-for use, reproduction, or distribution of Your modifications, or
-for any such Derivative Works as a whole, provided Your use,
-reproduction, and distribution of the Work otherwise complies with
-the conditions stated in this License.
-
-5. Submission of Contributions. Unless You explicitly state otherwise,
-any Contribution intentionally submitted for inclusion in the Work
-by You to the Licensor shall be under the terms and conditions of
-this License, without any additional terms or conditions.
-Notwithstanding the above, nothing herein shall supersede or modify
-the terms of any separate license agreement you may have executed
-with Licensor regarding such Contributions.
-
-6. Trademarks. This License does not grant permission to use the trade
-names, trademarks, service marks, or product names of the Licensor,
-except as required for reasonable and customary use in describing the
-origin of the Work and reproducing the content of the NOTICE file.
-
-7. Disclaimer of Warranty. Unless required by applicable law or
-agreed to in writing, Licensor provides the Work (and each
-Contributor provides its Contributions) on an "AS IS" BASIS,
-WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
-implied, including, without limitation, any warranties or conditions
-of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
-PARTICULAR PURPOSE. You are solely responsible for determining the
-appropriateness of using or redistributing the Work and assume any
-risks associated with Your exercise of permissions under this License.
-
-8. Limitation of Liability. In no event and under no legal theory,
-whether in tort (including negligence), contract, or otherwise,
-unless required by applicable law (such as deliberate and grossly
-negligent acts) or agreed to in writing, shall any Contributor be
-liable to You for damages, including any direct, indirect, special,
-incidental, or consequential damages of any character arising as a
-result of this License or out of the use or inability to use the
-Work (including but not limited to damages for loss of goodwill,
-work stoppage, computer failure or malfunction, or any and all
-other commercial damages or losses), even if such Contributor
-has been advised of the possibility of such damages.
-
-9. Accepting Warranty or Support. You may choose to offer,
-and charge a fee for, warranty, support, indemnity or other
-liability obligations and/or rights consistent with this License.
-However, in accepting such obligations, You may act only on Your
-own behalf and on Your sole responsibility, not on behalf of any
-other Contributor, and only if You agree to indemnify, defend,
-and hold each Contributor harmless for any liability incurred by,
-or claims asserted against, such Contributor by reason of your
-accepting any such warranty or support.
-
-END OF TERMS AND CONDITIONS
-
-Copyright 2024 Athena Team
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-See the License for the specific language governing permissions and
-limitations under the License.