code-loader 1.0.194__tar.gz → 1.0.195.dev0__tar.gz

{code_loader-1.0.194 → code_loader-1.0.195.dev0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: code-loader
-Version: 1.0.194
+Version: 1.0.195.dev0
 Summary: 
 Home-page: https://github.com/tensorleap/code-loader
 License: MIT

{code_loader-1.0.194 → code_loader-1.0.195.dev0}/code_loader/inner_leap_binder/leapbinder_decorators.py

@@ -1933,9 +1933,10 @@ def tensorleap_autoregressive_step(latent_space_aggregation: str = 'last_step'):
     state returned by the previous call, and return the next step's full input dict together with
     the updated state; returning (None, state) ends the chain — the terminating call still
     returns state, so the final step participates in state aggregation. State is never fed to
-    the model; it may be any nest of dicts/lists holding numpy arrays, numbers, strings or bools,
-    and it rides the platform queue on every step — accumulate reductions, not raw per-step
-    tensors. The hook must be stateless and deterministically seeded (seed stochasticity from
+    the model; any picklable nest of builtin/numpy values works (dicts, lists, tuples, sets,
+    numbers, strings, arrays — not user-defined classes, lambdas or open resources, since the
+    platform deserializes state without the integration code), and it rides the platform queue
+    on every step — accumulate reductions, not raw per-step tensors. The hook must be stateless and deterministically seeded (seed stochasticity from
     sample_id) — the engine replays steps after crash recovery and relies on identical results.
 
     latent_space_aggregation ('last_step' | 'mean') declares how the chain's latent-space

{code_loader-1.0.194 → code_loader-1.0.195.dev0}/code_loader/utils.py

@@ -1,4 +1,6 @@
+import io
 import math
+import pickle
 import sys
 from pathlib import Path
 from types import TracebackType
@@ -104,27 +106,31 @@ def get_metadata_type_from_variable(variable: Union[Optional[str], int, bool, Op
 
 
 
-AUTOREGRESSIVE_STATE_LEAF_TYPES = (np.ndarray, np.generic, int, float, str, bool, type(None))
+# The platform queue transports state as pickled payloads which the engine deserializes
+# WITHOUT the integration code loaded, so a class defined in the user's modules round-trips
+# locally but crashes on the engine. Only modules importable on the engine may resolve.
+_STATE_SAFE_PICKLE_MODULES = ('builtins', 'copyreg', 'collections', 'numpy', 'datetime')
+
+
+class _EngineSafeUnpickler(pickle.Unpickler):
+    def find_class(self, module: str, name: str) -> Any:
+        if module.split('.')[0] in _STATE_SAFE_PICKLE_MODULES:
+            return super().find_class(module, name)
+        raise pickle.UnpicklingError(
+            f'{module}.{name} is defined in integration code, which is not importable on the '
+            f'platform engine')
 
 
 def validate_autoregressive_state_types(value: Any, path: str = 'state') -> None:
-    if isinstance(value, dict):
-        for key, item in value.items():
-            if not isinstance(key, str):
-                raise AssertionError(
-                    f'autoregressive state validation failed: dict key {key!r} at {path} must be '
-                    f'a string. Got {type(key).__name__}.')
-            validate_autoregressive_state_types(item, f'{path}[{key!r}]')
-        return
-    if isinstance(value, (list, tuple)):
-        for i, item in enumerate(value):
-            validate_autoregressive_state_types(item, f'{path}[{i}]')
-        return
-    if not isinstance(value, AUTOREGRESSIVE_STATE_LEAF_TYPES):
+    try:
+        _EngineSafeUnpickler(io.BytesIO(pickle.dumps(value))).load()
+    except Exception as e:
         raise AssertionError(
-            f'autoregressive state validation failed: {path} holds a {type(value).__name__}. '
-            f'State must be built only from numpy arrays, numbers, strings, bools and nests of '
-            f'dicts/lists of those — it is serialized to the platform queue on every chain step.')
+            f'autoregressive state validation failed: {path} is not serializable ({e}). State '
+            'is serialized to the platform queue on every chain step and deserialized without '
+            'the integration code — any picklable nest of builtin/numpy values works (dicts, '
+            'lists, tuples, sets, numbers, strings, arrays), but user-defined classes, lambdas '
+            'and open resources do not.') from e
 
 
 def autoregressive_nests_equal(a: Any, b: Any) -> bool:

{code_loader-1.0.194 → code_loader-1.0.195.dev0}/pyproject.toml

@@ -1,7 +1,7 @@
 [tool.poetry]
 name = "code-loader"
 
-version = "1.0.194"
+version = "1.0.195.dev0"
 description = ""
 authors = ["dorhar <doron.harnoy@tensorleap.ai>"]
 license = "MIT"