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

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

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

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

@@ -36,6 +36,10 @@ from code_loader.mixpanel_tracker import clear_integration_events, AnalyticsEven
 
 _called_from_inside_tl_decorator = 0
 _called_from_inside_tl_integration_test_decorator = False
+# Whether any integration test wrapper was entered in this process. The crash reporter uses it
+# to tell a raise in the test body (report as code-flow failure) from a module-level crash
+# before any test ran (report as a plain script crash).
+_integration_test_started = False
 _call_from_tl_platform = os.environ.get('IS_TENSORLEAP_PLATFORM') == 'true'
 
 # ---- warnings store (module-level) ----
@@ -299,7 +303,8 @@ def tensorleap_integration_test():
                                     func_name='integration_test', expected_names=["idx", "preprocess"], **kwargs)
             _validate_input_args(*args, **kwargs)
 
-            global _called_from_inside_tl_integration_test_decorator
+            global _called_from_inside_tl_integration_test_decorator, _integration_test_started
+            _integration_test_started = True
             # Clear integration test events for new test
             try:
                 clear_integration_events()
@@ -2050,7 +2055,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 +2092,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 +2133,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 +2266,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 +2289,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).')
@@ -3395,7 +3407,11 @@ def tensorleap_status_table():
             return
 
         if _crashed["value"]:
-            print(f"\nScript crashed before completing all steps. crashed at function '{_current_func['name']}'.")
+            crashed_at = _current_func['name']
+            if crashed_at:
+                print(f"\nScript crashed before completing all steps. crashed at function '{crashed_at}'.")
+            else:
+                print("\nScript crashed before completing all steps.")
             return
 
         print(ready_mess) if ready else print(
@@ -3470,9 +3486,11 @@ def tensorleap_status_table():
             row = _find_row(crashed_name)
             if row:
                 row["Added to integration"] = CROSS
-        elif not crashed_name:
-            # Crash with no decorator currently running = the integration-test body /
-            # code flow itself failed; report that rather than blaming a decorator.
+        elif not crashed_name and _integration_test_started:
+            # Crash with no decorator currently running while an integration test has been
+            # entered = the test body / code flow itself failed; report that rather than
+            # blaming a decorator. Without a started test (a module-level crash between
+            # decorator registrations) the plain script-crash message is the truthful one.
             code_mapping_failure[0] = 1
 
         traceback.print_exception(exc_type, exc_value, exc_traceback)

{code_loader-1.0.193.dev0 → code_loader-1.0.193.dev2}/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 '
@@ -534,21 +545,39 @@ class LeapLoader(LeapLoaderBase):
             test_result.is_passed = False
         return test_result
 
+    @staticmethod
+    def _validate_autoregressive_step_inputs(step_inputs: Dict[str, Any],
+                                             sample_id: Union[int, str]) -> None:
+        # The parse-time check only covers the first sample of each state; the engine runtime
+        # paths must enforce the same key/value contract for every sample they touch.
+        for key, value in step_inputs.items():
+            if key.startswith('_'):
+                raise Exception(f'The autoregressive step hook returned model input key "{key}" '
+                                f'for sample {sample_id} — underscore passthrough keys were '
+                                'replaced by the explicit state channel; move passthrough values '
+                                'into the returned state.')
+            if not isinstance(value, np.ndarray):
+                raise Exception('The autoregressive step hook returned a non-numpy value for '
+                                f'model input "{key}" ({type(value).__name__}) for sample '
+                                f'{sample_id}.')
+
     def _ensure_autoregressive_input_shapes(self) -> Dict[str, List[int]]:
         handler = global_leap_binder.setup_container.autoregressive_step
         assert handler is not None
         if handler.input_shapes is None:
             preprocess_result = self._preprocess_result()
             first_state_response = next(iter(preprocess_result.values()))
-            first_result = handler.function(first_state_response.sample_ids[0], None, None, None,
+            first_sample_id = first_state_response.sample_ids[0]
+            first_result = handler.function(first_sample_id, None, None, None,
                                             first_state_response)
             if not isinstance(first_result, tuple) or len(first_result) != 2 or \
                     not isinstance(first_result[0], dict):
                 raise Exception('The autoregressive step hook must return a (next_inputs, state) '
                                 'tuple with a dict of initial model inputs on its first call, '
                                 f'got {type(first_result)}.')
-            handler.input_shapes = {key: list(value.shape) for key, value in first_result[0].items()
-                                    if isinstance(value, np.ndarray)}
+            self._validate_autoregressive_step_inputs(first_result[0], first_sample_id)
+            handler.input_shapes = {key: list(value.shape)
+                                    for key, value in first_result[0].items()}
         return handler.input_shapes
 
     def run_simulation(self, sim_name, params=None, n_samples=1, seed=0,
@@ -886,6 +915,7 @@ class LeapLoader(LeapLoaderBase):
                     'The autoregressive step hook must return a (next_inputs, state) tuple with a '
                     f'dict of initial model inputs on its first call, got {type(step_zero_result)} '
                     f'for sample {sample_id}.')
+            self._validate_autoregressive_step_inputs(step_zero_result[0], sample_id)
             inputs.update(step_zero_result[0])
         return inputs
 
@@ -1051,6 +1081,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 +1100,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.dev2}/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.dev2}/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.dev2}/pyproject.toml

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