code-loader 1.0.195.dev1__py3-none-any.whl → 1.0.196__py3-none-any.whl

code_loader/contract/datasetclasses.py

@@ -1,6 +1,6 @@
 import warnings
 from dataclasses import dataclass, field
-from typing import Any, Callable, List, Optional, Dict, Union, Type, Literal, cast
+from typing import Any, Callable, List, Optional, Dict, Union, Type, Literal, Tuple, cast
 import re
 import numpy as np
 import numpy.typing as npt
@@ -48,6 +48,8 @@ class PreprocessResponse:
     instance_to_sample_ids_mappings: Optional[Dict[str, str]] = None # in use only for element instance
     tl_generated: bool = False
     _grouped: bool = field(default=False, init=False, repr=False, compare=False)
+    _group_pos_cache: Optional[Dict[SampleId, Tuple[List[SampleId], int]]] = field(
+        default=None, init=False, repr=False, compare=False)
 
     def __post_init__(self) -> None:
         assert self.sample_ids_to_instance_mappings is None, f"Keep sample_ids_to_instance_mappings None when initializing PreprocessResponse"
@@ -57,7 +59,8 @@ class PreprocessResponse:
             self.sample_ids = cast(List[SampleId], [i for i in range(self.length)])
             self.sample_id_type = int
             self._grouped = False
-        elif self.length is None and self.sample_ids is not None:
+        elif self.sample_ids is not None:
+            given_length = self.length
             self._grouped = len(self.sample_ids) > 0 and isinstance(self.sample_ids[0], list)
             if self._grouped:
                 groups = cast(List[List[SampleId]], self.sample_ids)
@@ -66,6 +69,8 @@ class PreprocessResponse:
                     assert isinstance(group, list) and len(group) > 0, \
                         "Each group in a grouped PreprocessResponse must be a non-empty list."
                 self.sample_id_type = type(groups[0][0])
+                assert issubclass(self.sample_id_type, (str, int)), \
+                    f"Sample id should be of type str or int. Got: {self.sample_id_type.__name__}"
                 for group in groups:
                     for sample_id in group:
                         assert isinstance(sample_id, self.sample_id_type), \
@@ -79,6 +84,11 @@ class PreprocessResponse:
                 if self.sample_id_type == str:
                     for sample_id in flat_ids:
                         assert isinstance(sample_id, str), f"Sample id should be of type str. Got: {type(sample_id)}"
+            # dataclasses.replace() re-invokes __init__ with both fields already populated
+            # (length previously derived from sample_ids); only reject a genuine mismatch.
+            if given_length is not None:
+                assert given_length == self.length, \
+                    f"Inconsistent PreprocessResponse: length={given_length} but sample_ids implies {self.length}."
         else:
             raise Exception("length is deprecated, please use sample_ids instead.")
 
@@ -138,15 +148,19 @@ SectionCallableInterface = Callable[[Union[int, str], PreprocessResponse], npt.N
 InstanceCallableInterface = Callable[[Union[int, str], PreprocessResponse, int], Optional[ElementInstance]]
 InstanceLengthCallableInterface = Callable[[Union[int, str], PreprocessResponse], int]
 
-# (sample_id, prev_inputs, prev_outputs, preprocess) -> next model inputs, or None to end the chain.
-# First call per chain receives prev_inputs=None, prev_outputs=None and returns the initial inputs.
-# Keys starting with '_' are passthrough state: never fed to the model, handed back verbatim in
-# prev_inputs on the next call. The hook must be stateless and deterministically seeded (the engine
-# may replay a step after crash recovery and relies on identical results).
+# (sample_id, prev_inputs, prev_outputs, state, preprocess) -> (next model inputs | None, state).
+# First call per chain receives prev_inputs=None, prev_outputs=None, state=None and returns the
+# initial model inputs plus the initial state. Returning next_inputs=None ends the chain — the
+# terminating call still returns state, so the final step participates in state aggregation.
+# State is an arbitrary nest of dicts/lists holding numpy arrays, numbers, strings or bools; it is
+# never fed to the model and rides the engine's redis queue on every step, so accumulate
+# reductions, not raw per-step tensors. The hook must be stateless and deterministically seeded
+# (the engine may replay a step after crash recovery and relies on identical results).
+AutoregressiveChainState = Any
 AutoregressiveStepCallableInterface = Callable[
     [Union[int, str], Optional[Dict[str, npt.NDArray[np.float32]]], Optional[Dict[str, npt.NDArray[np.float32]]],
-     PreprocessResponse],
-    Optional[Dict[str, npt.NDArray[np.float32]]]]
+     AutoregressiveChainState, PreprocessResponse],
+    Tuple[Optional[Dict[str, npt.NDArray[np.float32]]], AutoregressiveChainState]]
 
 MetadataSectionCallableInterface = Union[
     Callable[[Union[int, str], PreprocessResponse], int],
@@ -302,16 +316,42 @@ class CustomLatentSpaceHandler:
 # dominates the model inputs). 'mean': every latent space is the elementwise mean over all steps.
 AUTOREGRESSIVE_LATENT_SPACE_AGGREGATIONS = ('last_step', 'mean')
 
+# Reserved argument names of autoregressive metrics/losses/visualizers, fed implicitly by the
+# platform from the finished chain (per-chain dicts, no batch axis). Any other argument is wired
+# to a ground-truth encoder through the integration test.
+AUTOREGRESSIVE_IMPLICIT_ARG_NAMES = ('inputs', 'outputs', 'state')
+
 
 @dataclass
 class AutoregressiveStepHandler:
     function: AutoregressiveStepCallableInterface
     name: str = 'autoregressive_step'
-    # Per-model-input shapes (unbatched, underscore passthrough keys excluded), discovered by running
-    # the hook's first call at parse time. Fills the role InputHandler.shape plays for input encoders.
+    # Per-model-input shapes (unbatched), discovered by running the hook's first call at parse
+    # time. Fills the role InputHandler.shape plays for input encoders.
     input_shapes: Optional[Dict[str, List[int]]] = None
     latent_space_aggregation: str = 'last_step'
 
+
+# Per-chain, unbatched callables: called once per finished chain with the final step's tensors —
+# inputs/outputs/state are fed implicitly by the platform (dicts, no batch axis; outputs keyed by
+# prediction-type names), remaining args are wired ground-truth encoder results.
+@dataclass
+class AutoregressiveMetricHandler:
+    metric_handler_data: MetricHandlerData
+    function: CustomCallableInterfaceMultiArgs
+
+
+@dataclass
+class AutoregressiveLossHandler:
+    custom_loss_handler_data: CustomLossHandlerData
+    function: CustomCallableInterface
+
+
+@dataclass
+class AutoregressiveVisualizerHandler:
+    visualizer_handler_data: VisualizerHandlerData
+    function: VisualizerCallableInterface
+
 @dataclass
 class PredictionTypeHandler:
     name: str
@@ -353,6 +393,9 @@ class DatasetIntegrationSetup:
     custom_latent_space: Optional[CustomLatentSpaceHandler] = None
     simulations: List[SimulationHandler] = field(default_factory=list)
     autoregressive_step: Optional[AutoregressiveStepHandler] = None
+    autoregressive_metrics: List[AutoregressiveMetricHandler] = field(default_factory=list)
+    autoregressive_losses: List[AutoregressiveLossHandler] = field(default_factory=list)
+    autoregressive_visualizers: List[AutoregressiveVisualizerHandler] = field(default_factory=list)
 
 
 @dataclass

code_loader/inner_leap_binder/leapbinder.py

@@ -1,6 +1,9 @@
 
+import builtins
 import inspect
-from typing import Callable, List, Optional, Dict, Any, Type, Union, get_args, cast
+import os
+from contextlib import contextmanager
+from typing import Callable, List, Optional, Dict, Any, Type, Union, get_args, cast, Iterator, Set
 
 import numpy as np
 import numpy.typing as npt
@@ -14,8 +17,10 @@ from code_loader.contract.datasetclasses import SectionCallableInterface, InputH
     RawInputsForHeatmap, VisualizerHandlerData, MetricHandlerData, CustomLossHandlerData, SamplePreprocessResponse, \
     ElementInstanceMasksHandler, InstanceCallableInterface, CustomLatentSpaceHandler, InstanceMetricHandler, \
     SimulationHandler, _simulation_context, AutoregressiveStepHandler, AutoregressiveStepCallableInterface, \
-    AUTOREGRESSIVE_LATENT_SPACE_AGGREGATIONS
-from code_loader.contract.enums import LeapDataType, DataStateEnum, DataStateType, MetricDirection, DatasetMetadataType
+    AUTOREGRESSIVE_LATENT_SPACE_AGGREGATIONS, AUTOREGRESSIVE_IMPLICIT_ARG_NAMES, \
+    AutoregressiveMetricHandler, AutoregressiveLossHandler, AutoregressiveVisualizerHandler
+from code_loader.contract.enums import LeapDataType, DataStateEnum, DataStateType, MetricDirection, DatasetMetadataType, \
+    TestingSectionEnum
 from code_loader.contract.mapping import NodeConnection, NodeMapping, NodeMappingType
 from code_loader.contract.responsedataclasses import DatasetTestResultPayload, LeapAnalysisConfiguration
 from code_loader.contract.visualizer_classes import map_leap_data_type_to_visualizer_class
@@ -32,6 +37,25 @@ from code_loader.visualizers.default_visualizers import DefaultVisualizer, \
 mapping_runtime_mode_env_var_mame = '__MAPPING_RUNTIME_MODE__'
 
 
+@contextmanager
+def _track_opened_files() -> Iterator[Set[str]]:
+    opened: Set[str] = set()
+    original_open = builtins.open
+
+    def tracking_open(file: Any, *args: Any, **kwargs: Any) -> Any:
+        mode = args[0] if args else kwargs.get('mode', 'r')
+        is_read_mode = isinstance(mode, str) and not any(c in mode for c in ('a', 'w', 'x'))
+        if is_read_mode and isinstance(file, (str, os.PathLike)):
+            opened.add(os.fspath(file))
+        return original_open(file, *args, **kwargs)
+
+    builtins.open = tracking_open
+    try:
+        yield opened
+    finally:
+        builtins.open = original_open
+
+
 def _stringized_annotation_type_name(annotation: Any) -> Optional[str]:
     """Bare type name if ``annotation`` is a stringized annotation (from ``from __future__
     import annotations`` or a quoted hint), else ``None``. code_loader inspects raw
@@ -548,11 +572,90 @@ class LeapBinder:
         # Builtin chain metadata, declared at parse time so it survives the reporter's
         # metadata type mapping; the placeholder values are overwritten by the engine when a
         # chain finalizes (realized length, truncated-by-safety-cap flag).
-        def builtin_chain_metadata(idx: Any, preprocess: PreprocessResponse) -> Dict[str, Any]:
+        def builtin_chain_metadata(idx: Any, preprocess: PreprocessResponse) -> Any:
+            if isinstance(idx, list):
+                return [{'length': -1, 'truncated': False} for _ in idx]
             return {'length': -1, 'truncated': False}
 
         self.set_metadata(builtin_chain_metadata, 'builtin_chain')
 
+    def _autoregressive_arg_names(self, function: Callable[..., Any], decorator_name: str,
+                                  name: str) -> List[str]:
+        """Split an autoregressive callable's signature into implicit chain args (validated
+        present) and wired arg names (returned) — the wired args connect to ground-truth
+        encoders through the integration test."""
+        argspec = inspect.getfullargspec(function)
+        for arg_name, arg_type in argspec.annotations.items():
+            _reject_stringized_sample_preprocess_response(function, arg_name, arg_type)
+            if arg_type == SamplePreprocessResponse:
+                raise Exception(
+                    f'{decorator_name} "{name}": SamplePreprocessResponse arguments are not '
+                    f'supported for autoregressive decorators yet. Derive what you need inside '
+                    f'the autoregressive step hook and carry it in the state.')
+        implicit = [arg_name for arg_name in argspec[0]
+                    if arg_name in AUTOREGRESSIVE_IMPLICIT_ARG_NAMES]
+        if not implicit:
+            raise Exception(
+                f'{decorator_name} "{name}": the function must declare at least one of the '
+                f'implicit chain arguments {list(AUTOREGRESSIVE_IMPLICIT_ARG_NAMES)} — they are '
+                f'fed by the platform from the finished chain.')
+        return [arg_name for arg_name in argspec[0]
+                if arg_name not in AUTOREGRESSIVE_IMPLICIT_ARG_NAMES]
+
+    def _assert_unique_metric_name(self, name: str) -> None:
+        existing = [handler.metric_handler_data.name for handler in self.setup_container.metrics] \
+            + [handler.metric_handler_data.name
+               for handler in self.setup_container.autoregressive_metrics]
+        if name in existing:
+            raise Exception(f'Metric with name {name} already exists. Please choose another')
+
+    def _assert_unique_loss_name(self, name: str) -> None:
+        existing = [handler.custom_loss_handler_data.name
+                    for handler in self.setup_container.custom_loss_handlers] \
+            + [handler.custom_loss_handler_data.name
+               for handler in self.setup_container.autoregressive_losses]
+        if name in existing:
+            raise Exception(f'Custom loss with name {name} already exists. Please choose another')
+
+    def _assert_unique_visualizer_name(self, name: str) -> None:
+        existing = [handler.visualizer_handler_data.name
+                    for handler in self.setup_container.visualizers] \
+            + [handler.visualizer_handler_data.name
+               for handler in self.setup_container.autoregressive_visualizers]
+        if name in existing:
+            raise Exception(f'Visualizer with name {name} already exists. Please choose another')
+
+    def add_autoregressive_metric(self, function: CustomCallableInterfaceMultiArgs, name: str,
+                                  direction: Optional[Union[MetricDirection, Dict[str, MetricDirection]]]
+                                  = MetricDirection.Downward,
+                                  compute_insights: Optional[Union[bool, Dict[str, bool]]] = None) -> None:
+        self._assert_unique_metric_name(name)
+        arg_names = self._autoregressive_arg_names(function, 'tensorleap_autoregressive_metric',
+                                                   name)
+        metric_handler_data = MetricHandlerData(name, arg_names, direction, compute_insights)
+        self.setup_container.autoregressive_metrics.append(
+            AutoregressiveMetricHandler(metric_handler_data, function))
+
+    def add_autoregressive_loss(self, function: CustomCallableInterface, name: str) -> None:
+        self._assert_unique_loss_name(name)
+        arg_names = self._autoregressive_arg_names(function, 'tensorleap_autoregressive_loss',
+                                                   name)
+        self.setup_container.autoregressive_losses.append(
+            AutoregressiveLossHandler(CustomLossHandlerData(name, arg_names), function))
+
+    def add_autoregressive_visualizer(self, function: VisualizerCallableInterface, name: str,
+                                      visualizer_type: LeapDataType) -> None:
+        self._assert_unique_visualizer_name(name)
+        if visualizer_type.value not in map_leap_data_type_to_visualizer_class:
+            raise Exception(
+                f'The visualizer_type is invalid. current visualizer_type: {visualizer_type}, '
+                f'should be one of : {", ".join([arg.__name__ for arg in get_args(LeapData)])}')
+        arg_names = self._autoregressive_arg_names(
+            function, 'tensorleap_autoregressive_visualizer', name)
+        self.setup_container.autoregressive_visualizers.append(
+            AutoregressiveVisualizerHandler(VisualizerHandlerData(name, visualizer_type,
+                                                                  arg_names), function))
+
     def set_custom_layer(self, custom_layer: Type[Any], name: str, inspect_layer: bool = False,
                          kernel_index: Optional[int] = None, use_custom_latent_space: bool = False) -> None:
         """
@@ -664,15 +767,34 @@ class LeapBinder:
             preprocess_response: PreprocessResponse, test_result: List[DatasetTestResultPayload],
             dataset_base_handler: Union[DatasetBaseHandler, MetadataHandler], state: DataStateEnum) -> List[DatasetTestResultPayload]:
         assert preprocess_response.sample_ids is not None
+        opened_files: Optional[Set[str]] = None
         if preprocess_response.is_grouped:
             # Grouped: probe with the first group, then reduce to a single sample's result so the
             # recorded shape/type is per-sample (matching the flat case), not the (B, *) batch.
             # (casts: sample_ids[0] is a group list here, and the grouped result is a per-sample list.)
             group = preprocess_response.sample_ids[0]
-            raw_result = cast(Any, dataset_base_handler.function(cast(Any, group), preprocess_response))[0]
+            with _track_opened_files() as opened_files:
+                raw_result = cast(Any, dataset_base_handler.function(cast(Any, group), preprocess_response))[0]
         else:
             raw_result = dataset_base_handler.function(
                 cast(Union[int, str], preprocess_response.sample_ids[0]), preprocess_response)
+
+        def _warn_if_group_spans_multiple_files(payloads: List[DatasetTestResultPayload]) -> None:
+            if opened_files is None or len(opened_files) <= 1:
+                return
+            shown = sorted(opened_files)[:5]
+            suffix = f" (+{len(opened_files) - 5} more)" if len(opened_files) > 5 else ""
+            warning = (
+                f"Encoding declared group 0 (size {len(cast(List[Any], group))}) for '{dataset_base_handler.name}' opened "
+                f"{len(opened_files)} distinct files: {shown}{suffix}. If this group is meant to be one "
+                f"physical source (e.g. one Parquet file), it may be grouped incorrectly - check your "
+                f"preprocess grouping logic. (Only files opened via Python's open() are tracked, so this "
+                f"can miss other I/O paths or be a false positive if grouping isn't meant to reflect file "
+                f"locality.)")
+            for payload in payloads:
+                existing = payload.display.get(TestingSectionEnum.Errors.name, '')
+                payload.display[TestingSectionEnum.Errors.name] = (existing + '\n' if existing else '') + warning
+
         handler_type = 'metadata' if isinstance(dataset_base_handler, MetadataHandler) else None
         if isinstance(dataset_base_handler, MetadataHandler):
             if isinstance(raw_result, dict):
@@ -712,6 +834,7 @@ class LeapBinder:
             else:
                 if raw_result is None:
                     if state != DataStateEnum.training:
+                        _warn_if_group_spans_multiple_files(test_result)
                         return test_result
 
                     if dataset_base_handler.metadata_type is None:
@@ -733,6 +856,7 @@ class LeapBinder:
             # setting shape in setup for all encoders
             if isinstance(dataset_base_handler, (InputHandler, GroundTruthHandler)):
                 dataset_base_handler.shape = result_shape
+        _warn_if_group_spans_multiple_files(test_result)
         return test_result
 
     def check_handlers(self, preprocess_result: Dict[DataStateEnum, PreprocessResponse]) -> None:
@@ -768,6 +892,10 @@ class LeapBinder:
             preprocess_response.tl_generated = True
             if not preprocess_response.length or preprocess_response.length < 1:
                 raise Exception("Simulation '{}' returned PreprocessResponse with length < 1".format(sim.name))
+            if preprocess_response.is_grouped:
+                raise Exception(
+                    "Simulation '{}' returned a grouped PreprocessResponse; simulations must "
+                    "return a flat (non-grouped) PreprocessResponse.".format(sim.name))
             preprocess_response.sample_ids = cast(List[Union[int, str]], [0])
             sim_sample_ids = cast(List[Union[int, str]], preprocess_response.sample_ids)
             for handler in self.setup_container.inputs:
@@ -796,6 +924,23 @@ class LeapBinder:
         appear in any order in the integration file.
         """
         if self.setup_container.autoregressive_step is None:
+            declared_autoregressive = [
+                ('tensorleap_autoregressive_metric',
+                 [handler.metric_handler_data.name
+                  for handler in self.setup_container.autoregressive_metrics]),
+                ('tensorleap_autoregressive_loss',
+                 [handler.custom_loss_handler_data.name
+                  for handler in self.setup_container.autoregressive_losses]),
+                ('tensorleap_autoregressive_visualizer',
+                 [handler.visualizer_handler_data.name
+                  for handler in self.setup_container.autoregressive_visualizers]),
+            ]
+            for decorator_name, names in declared_autoregressive:
+                if names:
+                    raise Exception(
+                        f'{decorator_name} {names} requires a tensorleap_autoregressive_step '
+                        f'hook — autoregressive decorators consume the finished chain the hook '
+                        f'drives.')
             return
         if self.setup_container.inputs:
             raise Exception(