code-loader 1.0.193.dev0__tar.gz → 1.0.193.dev1__tar.gz

{code_loader-1.0.193.dev0 → code_loader-1.0.193.dev1}/PKG-INFO

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

{code_loader-1.0.193.dev0 → code_loader-1.0.193.dev1}/code_loader/inner_leap_binder/leapbinder_decorators.py

@@ -2050,7 +2050,7 @@ def tensorleap_autoregressive_step(latent_space_aggregation: str = 'last_step'):
                 for key, value in first_inputs.items():
                     if not isinstance(value, np.ndarray):
                         continue
-                    assert np.array_equal(value, second_inputs[key]), error_message
+                    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):
@@ -2087,19 +2087,24 @@ def tensorleap_autoregressive_step(latent_space_aggregation: str = 'last_step'):
             sample_id, prev_inputs, prev_outputs, state, preprocess_response = \
                 _normalize_args(args, kwargs)
 
-            is_top_level_in_test = (_called_from_inside_tl_decorator == 0
+            is_top_level = _called_from_inside_tl_decorator == 0
+            is_top_level_in_test = (is_top_level
                                     and _called_from_inside_tl_integration_test_decorator)
             loop_context = _active_model_loop
-            if is_top_level_in_test:
-                if loop_context is None:
-                    raise LeapValidationError(
-                        f'{user_function.__name__}() validation failed: inside the integration '
-                        f'test the autoregressive step may only be called from within the '
-                        f'tensorleap_model_loop function — including the initial call '
-                        f'(prev_inputs=None). The platform drives the chain through the model '
-                        f'loop only, so a hook call outside it would not run on the platform.')
+            if is_top_level_in_test and loop_context is None:
+                raise LeapValidationError(
+                    f'{user_function.__name__}() validation failed: inside the integration '
+                    f'test the autoregressive step may only be called from within the '
+                    f'tensorleap_model_loop function — including the initial call '
+                    f'(prev_inputs=None). The platform drives the chain through the model '
+                    f'loop only, so a hook call outside it would not run on the platform.')
+            # Outside the integration test a model loop may still be running (local debugging);
+            # the loop's phase machine must advance there too, or its model proxy always raises.
+            notify_loop = is_top_level and loop_context is not None
+            if notify_loop:
                 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)
 
@@ -2123,10 +2128,10 @@ def tensorleap_autoregressive_step(latent_space_aggregation: str = 'last_step'):
             _validate_consistency(next_inputs)
             _validate_state_types(new_state)
 
-            if is_top_level_in_test:
-                if next_inputs is not None:
-                    next_inputs = {key: np.expand_dims(value, axis=0)
-                                   for key, value in next_inputs.items()}
+            if is_top_level_in_test and next_inputs is not None:
+                next_inputs = {key: np.expand_dims(value, axis=0)
+                               for key, value in next_inputs.items()}
+            if notify_loop:
                 loop_context.on_hook_return(next_inputs, new_state)
 
             if not _call_from_tl_platform:
@@ -2256,9 +2261,10 @@ def tensorleap_model_loop():
 
         def _run_with_context(model, sample_id, preprocess_response, is_mapping):
             global _active_model_loop, _model_loop_run_in_current_test
+            in_test = bool(_called_from_inside_tl_integration_test_decorator)
             if _active_model_loop is not None:
                 _fail('model loops cannot nest or run concurrently.')
-            if _model_loop_run_in_current_test:
+            if in_test and _model_loop_run_in_current_test:
                 _fail('only one model loop run is allowed per integration test.')
             if model is None or model is not _last_loaded_model:
                 _fail('the model argument must be the object returned by the '
@@ -2278,7 +2284,8 @@ def tensorleap_model_loop():
                                        preprocess_response)
             finally:
                 _active_model_loop = None
-            _model_loop_run_in_current_test = True
+            if in_test:
+                _model_loop_run_in_current_test = True
             if context.phase == 'awaiting_first_hook':
                 _fail('the loop never called the autoregressive step — it must start with the '
                       'initial hook call (prev_inputs=None, prev_outputs=None, state=None).')

{code_loader-1.0.193.dev0 → code_loader-1.0.193.dev1}/code_loader/leaploader.py

@@ -7,7 +7,7 @@ import sys
 from contextlib import redirect_stdout
 from functools import lru_cache
 from pathlib import Path
-from typing import Dict, List, Iterable, Set, Union, Any, Type, Optional, Callable, Tuple
+from typing import Dict, List, Iterable, Set, FrozenSet, Union, Any, Type, Optional, Callable, Tuple
 
 import numpy as np
 import numpy.typing as npt
@@ -480,6 +480,8 @@ class LeapLoader(LeapLoaderBase):
         try:
             preprocess_result = self._preprocess_result()
             input_shapes: Dict[str, List[int]] = {}
+            expected_keys: Optional[FrozenSet[str]] = None
+            expected_keys_state = ''
             for state, preprocess_response in preprocess_result.items():
                 if preprocess_response.sample_ids_to_instance_mappings:
                     raise Exception('Element instances are not supported together with '
@@ -506,6 +508,15 @@ class LeapLoader(LeapLoaderBase):
                                     f'(state: {state.name}). Seed any stochasticity from '
                                     'sample_id: the engine replays steps after crash recovery '
                                     'and relies on identical results.')
+                if expected_keys is None:
+                    expected_keys = frozenset(first_inputs)
+                    expected_keys_state = state.name
+                elif frozenset(first_inputs) != expected_keys:
+                    raise Exception(f'The autoregressive step hook returned different model input '
+                                    f'keys across states: {sorted(expected_keys)} '
+                                    f'({expected_keys_state}) vs {sorted(first_inputs)} '
+                                    f'({state.name}). The key set must be identical for every '
+                                    f'state.')
                 for key, value in first_inputs.items():
                     if key.startswith('_'):
                         raise Exception(f'The autoregressive step hook returned model input key '
@@ -520,7 +531,7 @@ class LeapLoader(LeapLoaderBase):
                                         f'model input "{key}" across states: {input_shapes[key]} vs '
                                         f'{list(value.shape)}. Shapes must be fixed.')
                     input_shapes[key] = list(value.shape)
-                    if not np.array_equal(value, second_result[0][key]):
+                    if not autoregressive_nests_equal(value, second_result[0][key]):
                         raise Exception('The autoregressive step hook is not deterministic — two calls '
                                         'with identical arguments returned different outputs for model '
                                         f'input "{key}". Seed any stochasticity from sample_id: the '
@@ -1051,6 +1062,11 @@ class LeapLoader(LeapLoaderBase):
         next_inputs, new_chain_state = result
         validate_autoregressive_state_types(new_chain_state)
         if next_inputs is None:
+            if prev_inputs is None:
+                raise Exception('The autoregressive step hook returned None for next_inputs on '
+                                f'the first call of the chain (sample {sample_id}) — the first '
+                                'call must return the initial model inputs; None would produce '
+                                'an empty chain with no forward pass.')
             return None, new_chain_state
         if not isinstance(next_inputs, dict):
             raise Exception('The autoregressive step hook must return a dict of model inputs or '
@@ -1065,7 +1081,12 @@ class LeapLoader(LeapLoaderBase):
                             f'on every call. Expected {sorted(expected_shapes)}, got '
                             f'{sorted(next_inputs)} for sample {sample_id}.')
         for key, value in next_inputs.items():
-            if isinstance(value, np.ndarray) and list(value.shape) != expected_shapes[key]:
+            if not isinstance(value, np.ndarray):
+                raise Exception('The autoregressive step hook must return numpy arrays for every '
+                                f'model input on every call — input "{key}" is a '
+                                f'{type(value).__name__} for sample {sample_id} (forgot .numpy() '
+                                'on a fed-back model output?).')
+            if list(value.shape) != expected_shapes[key]:
                 raise Exception('The autoregressive step hook must return fixed tensor shapes on '
                                 f'every call (use fixed-length padding + mask for growing '
                                 f'sequences). Input "{key}" was {expected_shapes[key]} on the '

{code_loader-1.0.193.dev0 → code_loader-1.0.193.dev1}/code_loader/leaploaderbase.py

@@ -149,10 +149,74 @@ class LeapLoaderBase:
     def has_custom_latent_space_decorator(self) -> bool:
         pass
 
+    # The autoregressive entry points raise instead of `pass`-ing like the abstract methods
+    # above: this class has no ABCMeta, so @abstractmethod is not enforced at runtime and an
+    # un-overridden `pass` body would return None — has_autoregressive_step would silently
+    # report False for an AR integration and the engine would skip the chain rollout with no
+    # error. A stale loader subclass keeps working for non-AR flows and fails loudly here the
+    # moment an AR flow reaches it.
+    @abstractmethod
     def has_autoregressive_step(self) -> bool:
-        # Default False so existing LeapLoaderBase implementations stay valid; the concrete
-        # loaders that can answer (LeapLoader in-process, LeapLoaderWithRedis via RPC) override.
-        return False
+        raise NotImplementedError(f'{type(self).__name__} does not implement '
+                                  'has_autoregressive_step — it must be overridden '
+                                  '(LeapLoader in-process, LeapLoaderWithRedis via RPC).')
+
+    @abstractmethod
+    def get_gt(self, state: DataStateEnum, sample_id: Union[int, str]
+               ) -> Dict[str, npt.NDArray[np.float32]]:
+        raise NotImplementedError(f'{type(self).__name__} does not implement get_gt.')
+
+    @abstractmethod
+    def get_prediction_names_in_order(self) -> List[str]:
+        raise NotImplementedError(f'{type(self).__name__} does not implement '
+                                  'get_prediction_names_in_order.')
+
+    @abstractmethod
+    def get_autoregressive_latent_space_aggregation(self) -> str:
+        raise NotImplementedError(f'{type(self).__name__} does not implement '
+                                  'get_autoregressive_latent_space_aggregation.')
+
+    @abstractmethod
+    def get_autoregressive_decorator_names(self) -> Dict[str, List[str]]:
+        raise NotImplementedError(f'{type(self).__name__} does not implement '
+                                  'get_autoregressive_decorator_names.')
+
+    @abstractmethod
+    def run_autoregressive_step(self, sample_id: Union[int, str],
+                                prev_inputs: Optional[Dict[str, npt.NDArray[np.float32]]],
+                                prev_outputs: Optional[Dict[str, npt.NDArray[np.float32]]],
+                                chain_state: Any,
+                                state: DataStateEnum
+                                ) -> Tuple[Optional[Dict[str, npt.NDArray[np.float32]]], Any]:
+        raise NotImplementedError(f'{type(self).__name__} does not implement '
+                                  'run_autoregressive_step.')
+
+    @abstractmethod
+    def run_autoregressive_metric(self, name: str, inputs: Dict[str, npt.NDArray[np.float32]],
+                                  outputs: Dict[str, npt.NDArray[np.float32]], chain_state: Any,
+                                  wired_args: Dict[str, npt.NDArray[np.float32]]) -> Any:
+        raise NotImplementedError(f'{type(self).__name__} does not implement '
+                                  'run_autoregressive_metric.')
+
+    @abstractmethod
+    def run_autoregressive_loss(self, name: str, inputs: Dict[str, npt.NDArray[np.float32]],
+                                outputs: Dict[str, npt.NDArray[np.float32]], chain_state: Any,
+                                wired_args: Dict[str, npt.NDArray[np.float32]]) -> Any:
+        raise NotImplementedError(f'{type(self).__name__} does not implement '
+                                  'run_autoregressive_loss.')
+
+    @abstractmethod
+    def run_autoregressive_visualizer(self, name: str, inputs: Dict[str, npt.NDArray[np.float32]],
+                                      outputs: Dict[str, npt.NDArray[np.float32]],
+                                      chain_state: Any,
+                                      wired_args: Dict[str, npt.NDArray[np.float32]]) -> Any:
+        raise NotImplementedError(f'{type(self).__name__} does not implement '
+                                  'run_autoregressive_visualizer.')
+
+    @abstractmethod
+    def autoregressive_visualizer_by_name(self) -> Dict[str, VisualizerHandlerData]:
+        raise NotImplementedError(f'{type(self).__name__} does not implement '
+                                  'autoregressive_visualizer_by_name.')
 
     @abstractmethod
     def get_heatmap_visualizer_raw_vis_input_arg_name(self, visualizer_name: str) -> Optional[str]:

{code_loader-1.0.193.dev0 → code_loader-1.0.193.dev1}/code_loader/utils.py

@@ -1,3 +1,4 @@
+import math
 import sys
 from pathlib import Path
 from types import TracebackType
@@ -134,6 +135,14 @@ def autoregressive_nests_equal(a: Any, b: Any) -> bool:
         return isinstance(a, (list, tuple)) and isinstance(b, (list, tuple)) and \
             len(a) == len(b) and all(autoregressive_nests_equal(x, y) for x, y in zip(a, b))
     if isinstance(a, np.ndarray) or isinstance(b, np.ndarray):
-        return isinstance(a, np.ndarray) and isinstance(b, np.ndarray) and \
-            a.shape == b.shape and bool(np.array_equal(a, b))
+        if not (isinstance(a, np.ndarray) and isinstance(b, np.ndarray) and a.shape == b.shape):
+            return False
+        # NaN leaves must equal themselves: this comparison backs the determinism checks, where
+        # NaN padding / running stats in a perfectly deterministic hook are legitimate.
+        if a.dtype.kind in 'fc' and b.dtype.kind in 'fc':
+            return bool(np.array_equal(a, b, equal_nan=True))
+        return bool(np.array_equal(a, b))
+    if isinstance(a, (float, np.floating)) and isinstance(b, (float, np.floating)) and \
+            math.isnan(a) and math.isnan(b):
+        return True
     return bool(a == b)

{code_loader-1.0.193.dev0 → code_loader-1.0.193.dev1}/pyproject.toml

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