code-loader 1.0.193.dev4__py3-none-any.whl → 1.0.195.dev0__py3-none-any.whl

code_loader/inner_leap_binder/leapbinder_decorators.py

@@ -1730,6 +1730,8 @@ def _ndarray_leaves(value):
 
 
 def _squeeze_leading_batch_axis(tensors_dict):
+    if tensors_dict is None:
+        return None
     squeezed = {}
     for key, value in tensors_dict.items():
         if isinstance(value, np.ndarray) and value.ndim > 0 and value.shape[0] == 1:
@@ -1931,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
@@ -2077,21 +2080,6 @@ def tensorleap_autoregressive_step(latent_space_aggregation: str = 'last_step'):
                     assert _nest_equal(value, second_inputs[key]), error_message
             assert _nest_equal(first_state, second_state), error_message
 
-        def _squeeze_test_batch_dim(tensors_dict):
-            # Inside the integration test the model works on batch-1 tensors (this wrapper adds the
-            # batch axis to returned inputs, and model outputs come back batched). The user's hook
-            # body must see the same unbatched per-sample tensors it will see at engine runtime, so
-            # strip the leading batch-1 axis before calling it.
-            if tensors_dict is None:
-                return None
-            squeezed = {}
-            for key, value in tensors_dict.items():
-                if isinstance(value, np.ndarray) and value.ndim > 0 and value.shape[0] == 1:
-                    squeezed[key] = value[0]
-                else:
-                    squeezed[key] = value
-            return squeezed
-
         def inner_without_validate(sample_id, prev_inputs, prev_outputs, state, preprocess_response):
             global _called_from_inside_tl_decorator
             _called_from_inside_tl_decorator += 1
@@ -2129,8 +2117,11 @@ def tensorleap_autoregressive_step(latent_space_aggregation: str = 'last_step'):
                 loop_context.on_hook_call(sample_id, prev_inputs, prev_outputs, state,
                                           preprocess_response)
             if is_top_level_in_test:
-                prev_inputs = _squeeze_test_batch_dim(prev_inputs)
-                prev_outputs = _squeeze_test_batch_dim(prev_outputs)
+                # The wrapper adds the batch axis to returned inputs and model outputs come back
+                # batched; the hook body must see the same unbatched per-sample tensors it will
+                # see at engine runtime.
+                prev_inputs = _squeeze_leading_batch_axis(prev_inputs)
+                prev_outputs = _squeeze_leading_batch_axis(prev_outputs)
 
             _validate_input_args(sample_id, prev_inputs, prev_outputs, state, preprocess_response)
 

code_loader/leaploader.py

@@ -581,8 +581,8 @@ class LeapLoader(LeapLoaderBase):
         return handler.input_shapes
 
     def run_simulation(self, sim_name, params=None, n_samples=1, seed=0,
-                       sample_ids=None, extend_preprocess=True):
-        # type: (str, Optional[Dict[str, Any]], int, int, Optional[List[str]], bool) -> Dict[str, Any]
+                       sample_ids=None, extend_preprocess=True, allow_partial=False):
+        # type: (str, Optional[Dict[str, Any]], int, int, Optional[List[str]], bool, bool) -> Dict[str, Any]
         # extend_preprocess=True (default): also extend preprocess_result[additional] with the
         # new synthetic sample_ids (producer-side behavior — synthetic worker needs this for
         # enumeration / data_length tracking). Pass False from consumer-side populators that
@@ -609,19 +609,30 @@ class LeapLoader(LeapLoaderBase):
         for sample_id in original_sample_ids:
             for handler in global_leap_binder.setup_container.inputs:
                 per_encoder[handler.name].append(handler.function(sample_id, sim_preprocess))
-        encoded = {name: np.stack(arrays) for name, arrays in per_encoder.items()}
+        encoded = {
+            name: np.stack(arrays) if arrays else np.empty((0,), dtype=np.float32)
+            for name, arrays in per_encoder.items()
+        }
         if sample_ids is not None:
-            if len(sample_ids) != len(original_sample_ids):
-                raise ValueError(
-                    "sample_ids length ({}) does not match simulation output length ({})".format(
-                        len(sample_ids), len(original_sample_ids)
-                    )
-                )
             for sid in sample_ids:
                 if not isinstance(sid, str):
                     raise TypeError(
                         "All sample_ids must be of type str. Got: {}".format(type(sid))
                     )
+            if len(sample_ids) != len(original_sample_ids):
+                if not allow_partial:
+                    raise ValueError(
+                        "sample_ids length ({}) does not match simulation output length ({})".format(
+                            len(sample_ids), len(original_sample_ids)
+                        )
+                    )
+                n_bound = min(len(sample_ids), len(original_sample_ids))
+                sample_ids = list(sample_ids)[:n_bound]
+                original_sample_ids = original_sample_ids[:n_bound]
+                # Over-generation trims sample_ids/original_sample_ids above but encoded was
+                # stacked from the full simulation output; slice it too so the encoded batch dim
+                # stays 1:1 with the returned sample_ids.
+                encoded = {name: arr[:n_bound] for name, arr in encoded.items()}
             for synth_id, original_local_id in zip(sample_ids, original_sample_ids):
                 self._synthetic_lookup[synth_id] = (sim_preprocess, original_local_id)
             if extend_preprocess:
@@ -631,6 +642,19 @@ class LeapLoader(LeapLoaderBase):
             returned_sample_ids = original_sample_ids
         return {"encoded": encoded, "sample_ids": returned_sample_ids}
 
+    def prune_synthetic_lookup(self, keep_ids):
+        # type: (List[str]) -> None
+        # Keep only the given synthetic sample_ids in the in-memory lookup and
+        # drop the rest, so preprocess responses don't accumulate across many
+        # run_simulation calls (e.g. a synthetic calibration loop keeps only its
+        # running top-K trials). Keep-set semantics (rather than evict-list) so
+        # the caller can bound the lookup to a known-good set each iteration
+        # without tracking everything it ever generated.
+        keep = set(keep_ids)
+        self._synthetic_lookup = {
+            sid: entry for sid, entry in self._synthetic_lookup.items() if sid in keep
+        }
+
     def _extend_additional_preprocess(self, new_sample_ids: List[str]) -> None:
         if self._preprocess_result_cached is None:
             self._preprocess_result()

code_loader/leaploaderbase.py

@@ -223,8 +223,13 @@ class LeapLoaderBase:
         pass
 
     @abstractmethod
-    def run_simulation(self, sim_name, params=None, n_samples=1, seed=0, sample_ids=None):
-        # type: (str, Optional[Dict[str, Any]], int, int, Optional[List[str]]) -> Dict[str, Any]
+    def run_simulation(self, sim_name, params=None, n_samples=1, seed=0, sample_ids=None,
+                       extend_preprocess=True, allow_partial=False):
+        # type: (str, Optional[Dict[str, Any]], int, int, Optional[List[str]], bool, bool) -> Dict[str, Any]
+        pass
+
+    def prune_synthetic_lookup(self, keep_ids):
+        # type: (List[str]) -> None
         pass
 
     def is_custom_latent_space(self) -> bool:

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.193.dev4.dist-info → code_loader-1.0.195.dev0.dist-info}/METADATA

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

{code_loader-1.0.193.dev4.dist-info → code_loader-1.0.195.dev0.dist-info}/RECORD

@@ -22,17 +22,17 @@ code_loader/experiment_api/utils.py,sha256=XZHtxge12TS4H4-8PjV3sKuhp8Ud6ojAiIzTZ
 code_loader/experiment_api/workingspace_config_utils.py,sha256=DLzXQCg4dgTV_YgaSbeTVzq-2ja_SQw4zi7LXwKL9cY,990
 code_loader/inner_leap_binder/__init__.py,sha256=koOlJyMNYzGbEsoIbXathSmQ-L38N_pEXH_HvL7beXU,99
 code_loader/inner_leap_binder/leapbinder.py,sha256=4cbSEbN6vVwjFUMxDEksxk1VSCip2ntLB32RxA81_JE,53194
-code_loader/inner_leap_binder/leapbinder_decorators.py,sha256=ShgqKiyfcJKMIpoB_QdxHoQJ0S1gdQDEQgoPIv6sSIw,176492
-code_loader/leaploader.py,sha256=sx-BzeVG3qt54S2uF5VbBi2xL0x70AIFN1to6OtjBYU,66979
-code_loader/leaploaderbase.py,sha256=ncKEWsVv_1Vhtgl8Dqgt4ampxTzOj5nVitGJwQQtTc4,10561
+code_loader/inner_leap_binder/leapbinder_decorators.py,sha256=eHdSRCYtfkHtq7NSlS4bF92TU8RGRzZ1gbb3ULeEV2c,176145
+code_loader/leaploader.py,sha256=ME1Ypq9VV3o9XXOCV6voDWWVP0gB4Z0LoJh6u79qDYY,68370
+code_loader/leaploaderbase.py,sha256=PvBU_2OQqBfJLDMsbK1s372954eYr4Y3O4r18y5S7uY,10739
 code_loader/mixpanel_tracker.py,sha256=rNwRmFifNbdUoqLQvvhhgpKczWpWiEmd8MfyJe27sxw,9131
 code_loader/plot_functions/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 code_loader/plot_functions/plot_functions.py,sha256=2DC-zlVaN13P4VNx5d8csgs80C6SisaeP1-Kq2LW7iM,16075
 code_loader/plot_functions/visualize.py,sha256=gsBAYYkwMh7jIpJeDMPS8G4CW-pxwx6LznoQIvi4vpo,657
-code_loader/utils.py,sha256=1XLY5hNQaOsqXeUIooGs6KXV0E1edeiQyxAdnG3eiy0,6338
+code_loader/utils.py,sha256=gEJCKpDWf6p9_KoEz-0EGiJqxD6QfPgZN6zdKdtSAz8,6628
 code_loader/visualizers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 code_loader/visualizers/default_visualizers.py,sha256=onRnLE_TXfgLN4o52hQIOOhUcFexGlqJ3xSpQDVLuZM,2604
-code_loader-1.0.193.dev4.dist-info/LICENSE,sha256=qIwWjdspQeSMTtnFZBC8MuT-95L02FPvzRUdWFxrwJY,1067
-code_loader-1.0.193.dev4.dist-info/METADATA,sha256=0HFcGksXyXxl3m20DjXWkMzAxAXZ7Zt3O0Vl8a8h0K4,1095
-code_loader-1.0.193.dev4.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-code_loader-1.0.193.dev4.dist-info/RECORD,,
+code_loader-1.0.195.dev0.dist-info/LICENSE,sha256=qIwWjdspQeSMTtnFZBC8MuT-95L02FPvzRUdWFxrwJY,1067
+code_loader-1.0.195.dev0.dist-info/METADATA,sha256=q8UczKAxuYThEseqP4aXcS5ET0DfpjPaRmIy4aKDl1I,1095
+code_loader-1.0.195.dev0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+code_loader-1.0.195.dev0.dist-info/RECORD,,