code-loader 1.0.196.dev0__tar.gz → 1.0.197__tar.gz

{code_loader-1.0.196.dev0 → code_loader-1.0.197}/PKG-INFO

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

{code_loader-1.0.196.dev0 → code_loader-1.0.197}/code_loader/contract/datasetclasses.py

@@ -59,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)
@@ -68,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), \
@@ -81,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.")
 

{code_loader-1.0.196.dev0 → code_loader-1.0.197}/code_loader/contract/enums.py

@@ -28,6 +28,7 @@ class LeapDataType(Enum):
     ImageWithBBox = 'ImageWithBBox'
     ImageWithHeatmap = 'ImageWithHeatmap'
     Video = 'Video'
+    Audio = 'Audio'
 
 
 class MetricDirection(Enum):

{code_loader-1.0.196.dev0 → code_loader-1.0.197}/code_loader/contract/visualizer_classes.py

@@ -343,6 +343,47 @@ class LeapImageWithHeatmap:
                 'Number of heatmaps and labels must be equal')
 
 
+@dataclass
+class LeapAudio:
+    """
+    Visualizer representing a playable audio sample for Tensorleap.
+
+    The user plays the sample in the right panel while a playhead tick runs
+    through `visual` — either a spectrogram (a LeapImage) or a time-series
+    waveform (a LeapGraph). Build one with `default_audio_visualizer`.
+
+    Attributes:
+    audio (npt.NDArray): Mono waveform samples, shaped [N]. Multi-channel input
+        should be mixed down before constructing LeapAudio.
+    sample_rate (int): Samples per second — needed to encode the clip and to map
+        the playhead position to time.
+    visual (Union[LeapImage, LeapGraph]): The thing you look at.
+    x_range (Optional[Tuple[float, float]]): Time extent (in seconds) that the
+        visual's horizontal axis spans, used to align the tick. Defaults to
+        (0, N / sample_rate).
+    type (LeapDataType): The data type, default is LeapDataType.Audio.
+    """
+    audio: npt.NDArray
+    sample_rate: int
+    visual: Union[LeapImage, LeapGraph]
+    x_range: Optional[Tuple[float, float]] = None
+    type: LeapDataType = LeapDataType.Audio
+
+    def __post_init__(self) -> None:
+        validate_type(self.type, LeapDataType.Audio)
+        validate_type(type(self.audio), np.ndarray)
+        validate_type(len(self.audio.shape), 1, 'Audio must be mono, of shape 1 [N]')
+        validate_type(type(self.visual), [LeapImage, LeapGraph],
+                      'Audio visual must be a LeapImage (spectrogram) or LeapGraph (waveform)')
+        if not isinstance(self.sample_rate, int) or self.sample_rate <= 0:
+            raise LeapValidationError(f'sample_rate must be a positive int, got {self.sample_rate}')
+        if self.x_range is not None:
+            if len(self.x_range) != 2:
+                raise LeapValidationError('x_range must be a tuple of length 2')
+            validate_type(type(self.x_range[0]), [float, int], 'x_range must be a tuple of floats or integers')
+            validate_type(type(self.x_range[1]), [float, int], 'x_range must be a tuple of floats or integers')
+
+
 map_leap_data_type_to_visualizer_class = {
     LeapDataType.Image.value: LeapImage,
     LeapDataType.Graph.value: LeapGraph,
@@ -352,5 +393,6 @@ map_leap_data_type_to_visualizer_class = {
     LeapDataType.ImageMask.value: LeapImageMask,
     LeapDataType.TextMask.value: LeapTextMask,
     LeapDataType.ImageWithBBox.value: LeapImageWithBBox,
-    LeapDataType.ImageWithHeatmap.value: LeapImageWithHeatmap
+    LeapDataType.ImageWithHeatmap.value: LeapImageWithHeatmap,
+    LeapDataType.Audio.value: LeapAudio
 }

{code_loader-1.0.196.dev0 → code_loader-1.0.197}/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
@@ -16,7 +19,8 @@ from code_loader.contract.datasetclasses import SectionCallableInterface, InputH
     SimulationHandler, _simulation_context, AutoregressiveStepHandler, AutoregressiveStepCallableInterface, \
     AUTOREGRESSIVE_LATENT_SPACE_AGGREGATIONS, AUTOREGRESSIVE_IMPLICIT_ARG_NAMES, \
     AutoregressiveMetricHandler, AutoregressiveLossHandler, AutoregressiveVisualizerHandler
-from code_loader.contract.enums import LeapDataType, DataStateEnum, DataStateType, MetricDirection, DatasetMetadataType
+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
@@ -33,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
@@ -744,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):
@@ -792,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:
@@ -813,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:
@@ -848,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:

{code_loader-1.0.196.dev0 → code_loader-1.0.197}/code_loader/inner_leap_binder/leapbinder_decorators.py

@@ -29,7 +29,7 @@ from code_loader.contract.enums import MetricDirection, LeapDataType, DatasetMet
 from code_loader import leap_binder, LeapLoader
 from code_loader.contract.mapping import NodeMapping, NodeMappingType, NodeConnection
 from code_loader.contract.visualizer_classes import LeapImage, LeapImageMask, LeapTextMask, LeapText, LeapGraph, \
-    LeapHorizontalBar, LeapImageWithBBox, LeapImageWithHeatmap, LeapVideo, LeapValidationError, \
+    LeapHorizontalBar, LeapImageWithBBox, LeapImageWithHeatmap, LeapVideo, LeapAudio, LeapValidationError, \
     map_leap_data_type_to_visualizer_class
 from code_loader.inner_leap_binder.leapbinder import mapping_runtime_mode_env_var_mame
 from code_loader.mixpanel_tracker import clear_integration_events, AnalyticsEvent, emit_integration_event_once
@@ -686,10 +686,29 @@ def tensorleap_load_model(prediction_types: Optional[List[PredictionTypeHandler]
 
             class ModelPlaceholder:
 
+                @staticmethod
+                def _reject_reused_input_source(elem, seen):
+                    # A model input slot is wired by mutating elem.node_mapping.type, so the same
+                    # source object at two positions silently overwrites the first slot. Each input
+                    # encoder mints its own NodeMapping, so a collision means one source fed two
+                    # slots — e.g. indexing a single grouped input, model([imgs[0], imgs[1]]), where
+                    # imgs[i] returns the same node (a group is the batch dim of ONE input, not many
+                    # inputs), or the flat model([img, img]). Reject it instead of mis-wiring.
+                    if id(elem.node_mapping) in seen:
+                        raise Exception(
+                            "The same input source was passed to more than one model input slot. Each "
+                            "model input must come from a distinct input encoder. Indexing a single "
+                            "grouped input (e.g. model([imgs[0], imgs[1]])) feeds several samples of ONE "
+                            "input into separate slots — a group is the batch dimension of one input, "
+                            "not multiple inputs.")
+                    seen.add(id(elem.node_mapping))
+
                 # keras interface
                 def __call__(self, arg):
                     if isinstance(arg, list):
+                        seen: set = set()
                         for i, elem in enumerate(arg):
+                            self._reject_reused_input_source(elem, seen)
                             elem.node_mapping.type = _safe_get_item(i)
                     else:
                         arg.node_mapping.type = NodeMappingType.Input0
@@ -703,7 +722,9 @@ def tensorleap_load_model(prediction_types: Optional[List[PredictionTypeHandler]
                                                   "use model.run(None, inputs) to return all outputs.")
                     assert isinstance(input_dict, dict), \
                         f'Expected input_dict to be a dict, got {type(input_dict)} instead.'
+                    seen: set = set()
                     for i, (input_key, elem) in enumerate(input_dict.items()):
+                        self._reject_reused_input_source(elem, seen)
                         if isinstance(input_key, NodeMappingType):
                             elem.node_mapping.type = input_key
                         else:
@@ -1522,7 +1543,8 @@ def tensorleap_custom_visualizer(name: str, visualizer_type: LeapDataType,
                 LeapDataType.HorizontalBar: LeapHorizontalBar,
                 LeapDataType.ImageWithBBox: LeapImageWithBBox,
                 LeapDataType.ImageWithHeatmap: LeapImageWithHeatmap,
-                LeapDataType.Video: LeapVideo
+                LeapDataType.Video: LeapVideo,
+                LeapDataType.Audio: LeapAudio
             }
             validate_output_structure(result, func_name=user_function.__name__,
                                       expected_type_name=result_type_map[visualizer_type])
@@ -2010,9 +2032,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
@@ -2844,6 +2867,13 @@ def tensorleap_element_instance_preprocess(
             result = user_function()
             found_instance_metadata = False
             for preprocess_response in result:
+                if preprocess_response.is_grouped:
+                    raise Exception(
+                        "tensorleap_element_instance_preprocess validation failed: a grouped "
+                        "(nested sample_ids) PreprocessResponse cannot also declare element "
+                        "instances -- instance ids are derived per scalar sample id, and the "
+                        "engine's grouped/batch fetch path has no per-instance id to mask by. "
+                        "Remove the instance encoders or don't group this dataset.")
                 sample_ids_to_instance_mappings = {}
                 instance_to_sample_ids_mappings = {}
                 all_sample_ids = preprocess_response.sample_ids.copy()
@@ -3114,10 +3144,12 @@ def tensorleap_input_encoder(name: str, channel_dim=_UNSET, model_input_index=No
             if _called_from_inside_tl_decorator == 0 and _called_from_inside_tl_integration_test_decorator:
                 if grouped:
                     # A grouped result is a list of per-sample arrays (never stacked, so a
-                    # B=1-only model can be fed one sample at a time). Add a batch dim to each
-                    # sample, mirroring the flat path's single-sample expand_dims.
-                    if isinstance(result, list):
-                        result = [np.expand_dims(r, axis=0) for r in result]
+                    # B=1-only model can be fed one sample at a time). A user may also return a
+                    # supported (B, *) array; split it back into per-sample rows first (mirrors
+                    # LeapLoader._to_grouped_list) so both forms match the runtime path. Then add
+                    # a batch dim to each sample, mirroring the flat path's single-sample expand_dims.
+                    rows = result if isinstance(result, list) else [row for row in np.asarray(result)]
+                    result = [np.expand_dims(r, axis=0) for r in rows]
                 else:
                     batch_warning(result, user_function.__name__)
                     result = np.expand_dims(result, axis=0)
@@ -3216,7 +3248,7 @@ def tensorleap_gt_encoder(name: str):
                 set_current("tensorleap_gt_encoder")
             validate_args_structure(*args, types_order=[Union[int, str, list], PreprocessResponse],
                                     func_name=user_function.__name__, expected_names=["idx", "preprocess"], **kwargs)
-            sample_id, preprocess_response = args
+            sample_id, preprocess_response = args if len(args) != 0 else kwargs.values()
             _validate_input_args(sample_id, preprocess_response)
 
             grouped = isinstance(sample_id, list)
@@ -3228,10 +3260,12 @@ def tensorleap_gt_encoder(name: str):
             if _called_from_inside_tl_decorator == 0 and _called_from_inside_tl_integration_test_decorator:
                 if grouped:
                     # A grouped result is a list of per-sample arrays (never stacked, so a
-                    # B=1-only model can be fed one sample at a time). Add a batch dim to each
-                    # sample, mirroring the flat path's single-sample expand_dims.
-                    if isinstance(result, list):
-                        result = [np.expand_dims(r, axis=0) for r in result]
+                    # B=1-only model can be fed one sample at a time). A user may also return a
+                    # supported (B, *) array; split it back into per-sample rows first (mirrors
+                    # LeapLoader._to_grouped_list) so both forms match the runtime path. Then add
+                    # a batch dim to each sample, mirroring the flat path's single-sample expand_dims.
+                    rows = result if isinstance(result, list) else [row for row in np.asarray(result)]
+                    result = [np.expand_dims(r, axis=0) for r in rows]
                 else:
                     batch_warning(result, user_function.__name__)
                     result = np.expand_dims(result, axis=0)

{code_loader-1.0.196.dev0 → code_loader-1.0.197}/code_loader/leaploader.py

@@ -41,6 +41,11 @@ def _serialize_sim_bounds(bounds) -> dict:
 
 
 class LeapLoader(LeapLoaderBase):
+    # Opt-in ceiling (env var) on samples materialized per get_samples call — the engine sets it to
+    # turn a mistaken whole-group ask into a clear error instead of a silent worker OOM. Unset => no
+    # ceiling. See get_samples and grouped-fetch-oom-bug.md.
+    _grouped_fetch_max_samples_env = "LEAP_GROUPED_FETCH_MAX_SAMPLES_PER_CALL"
+
     def __init__(self, code_path: str, code_entry_name: str):
         super().__init__(code_path, code_entry_name)
 
@@ -77,8 +82,6 @@ class LeapLoader(LeapLoaderBase):
                     preprocess_result = global_leap_binder.get_preprocess_result()
                     self._preprocess_result_cached = preprocess_result
                     is_grouped = any(r.is_grouped for r in preprocess_result.values())
-                except Exception:
-                    is_grouped = False
                 finally:
                     os.environ[mapping_runtime_mode_env_var_mame] = 'TRUE'
 
@@ -285,8 +288,10 @@ class LeapLoader(LeapLoaderBase):
             latent_fn = global_leap_binder.setup_container.custom_latent_space.function
             preprocess_state = preprocess_result[state]
             if preprocess_state.is_grouped:
-                group_ids, pos = self._locate_group(preprocess_state, sample_id)
-                custom_latent_space = self._to_grouped_list(latent_fn(group_ids, preprocess_state))[pos]
+                # Single-sample fetch: encode this sample only, not the whole group (same
+                # memory rationale as _get_dataset_handlers; see grouped-fetch-oom-bug.md).
+                self._locate_group(preprocess_state, sample_id)  # validates group membership
+                custom_latent_space = self._to_grouped_list(latent_fn([sample_id], preprocess_state))[0]
             else:
                 custom_latent_space = latent_fn(sample_id, preprocess_state)
         instance_mask = self._get_instances_masks(state, sample_id, instance_id)
@@ -416,7 +421,7 @@ class LeapLoader(LeapLoaderBase):
         test_result = DatasetTestResultPayload('preprocess')
         try:
             preprocess_result = self._preprocess_result()
-            if self.get_sample_id_type() is str:
+            if issubclass(self.get_sample_id_type(), str):
                 max_allowed_item_size = np.dtype('<U256').itemsize
                 for state, preprocess_response in preprocess_result.items():
                     sample_ids_array = np.array(preprocess_response.flat_sample_ids)
@@ -480,6 +485,10 @@ class LeapLoader(LeapLoaderBase):
                 preprocess_response.tl_generated = True
                 if preprocess_response.length < 1:
                     raise ValueError("Simulation returned PreprocessResponse with length < 1")
+                if preprocess_response.is_grouped:
+                    raise ValueError(
+                        "Simulation returned a grouped PreprocessResponse; simulations must "
+                        "return a flat (non-grouped) PreprocessResponse.")
                 preprocess_response.sample_ids = [0]
                 for handler in global_leap_binder.setup_container.inputs:
                     out1 = handler.function(preprocess_response.sample_ids[0], preprocess_response)
@@ -1001,11 +1010,17 @@ class LeapLoader(LeapLoaderBase):
         else:
             preprocess_state = preprocess_result[state]
         if preprocess_state.is_grouped:
-            # Grouped dataset: call the encoder once with the whole group, index out the sample.
-            group_ids, pos = self._locate_group(preprocess_state, sample_id)
+            # Grouped dataset, SINGLE-sample fetch: encode only THIS sample, never the whole group.
+            # Grouping is an I/O read-unit, not an encode-unit — the integration's per-file cache
+            # makes the group's file decode once across consecutive per-sample calls, so a singleton
+            # ([sample_id]) is both a valid single-group subset of the get_samples contract and
+            # correct (samples in a group are independent; see _to_grouped_list). Handing the encoder
+            # the whole group here just to index one sample out materialized B encoded tensors per
+            # sample — a driver of the grouped-fetch worker OOM (see grouped-fetch-oom-bug.md).
+            self._locate_group(preprocess_state, sample_id)  # validates the id belongs to a group
             for handler in handlers:
-                grouped = self._to_grouped_list(handler.function(group_ids, preprocess_state))
-                result_agg[handler.name] = grouped[pos]
+                result_agg[handler.name] = self._to_grouped_list(
+                    handler.function([sample_id], preprocess_state))[0]
             return result_agg
         for handler in handlers:
             handler_result = handler.function(sample_id, preprocess_state)
@@ -1038,6 +1053,23 @@ class LeapLoader(LeapLoaderBase):
             f"get_samples received sample ids spanning {len(distinct_groups)} groups; a request "
             f"must be confined to a single group. The engine partitions cross-group requests and "
             f"issues one get_samples call per group.")
+        # Opt-in safety net for the grouped-fetch OOM (grouped-fetch-oom-bug.md). This call
+        # materializes len(group_ids) encoded samples at once; on a large group/file that peak is
+        # what kernel-OOM-kills the worker. A group's file is decoded once by the integration's
+        # per-file cache regardless of how the encode is chunked, so the engine should request
+        # memory-bounded sub-chunks of a group (each a valid single-group subset, e.g.
+        # get_samples(state, group_ids[i:j])) rather than the whole group. When the engine sets the
+        # ceiling below, an over-large ask fails here with a message that names the lever instead of
+        # a silent GENERIC_MEMORY_LIMIT kill. Unset => no ceiling (preserves current whole-group callers).
+        max_per_call = os.environ.get(self._grouped_fetch_max_samples_env)
+        if max_per_call is not None and len(group_ids) > int(max_per_call):
+            raise ValueError(
+                f"get_samples asked to materialize {len(group_ids)} samples in a single call, above "
+                f"the configured ceiling {max_per_call} ({self._grouped_fetch_max_samples_env}). The "
+                f"group's file is decoded once via the integration's per-file cache regardless of "
+                f"chunking, so request memory-bounded sub-chunks of the group "
+                f"(get_samples(state, group_ids[i:j])) instead of the whole group. "
+                f"See grouped-fetch-oom-bug.md.")
         inputs = {handler.name: self._to_grouped_list(handler.function(group_ids, preprocess_state))
                   for handler in global_leap_binder.setup_container.inputs}
         autoregressive_handler = global_leap_binder.setup_container.autoregressive_step
@@ -1059,8 +1091,9 @@ class LeapLoader(LeapLoaderBase):
                         f'{type(step_zero_result)} for sample {sid}.')
                 self._validate_autoregressive_step_inputs(step_zero_result[0], sid)
                 step_zero_per_sample.append(step_zero_result[0])
-            for key in (step_zero_per_sample[0].keys() if step_zero_per_sample else []):
-                inputs[key] = [row[key] for row in step_zero_per_sample]
+            if step_zero_per_sample:
+                for key in self._grouped_dict_keys(step_zero_per_sample, autoregressive_handler.name):
+                    inputs[key] = [row[key] for row in step_zero_per_sample]
         gt = None
         if state != DataStateEnum.unlabeled:
             gt = {handler.name: self._to_grouped_list(handler.function(group_ids, preprocess_state))
@@ -1069,8 +1102,8 @@ class LeapLoader(LeapLoaderBase):
         metadata: Dict[str, Any] = {}
         metadata_is_none: Dict[str, Any] = {}
         for handler in global_leap_binder.setup_container.metadata:
-            rows = handler.function(group_ids, preprocess_state)  # list of B scalars/dicts
-            if rows and isinstance(rows[0], dict):
+            rows = handler.function(group_ids, preprocess_state)  # length-B sequence of scalars/dicts
+            if len(rows) > 0 and isinstance(rows[0], dict):
                 for key in self._grouped_dict_keys(rows, handler.name):
                     name = "{}_{}".format(handler.name, key)
                     converted = [self._convert_metadata_to_correct_type(name, row[key]) for row in rows]
@@ -1197,16 +1230,23 @@ class LeapLoader(LeapLoaderBase):
             sample_id = original_local_id
         else:
             preprocess_state = preprocess_result[state]
-        group_ids, pos = (None, None)
         if preprocess_state.is_grouped:
-            group_ids, pos = self._locate_group(preprocess_state, sample_id)
+            self._locate_group(preprocess_state, sample_id)  # validates the id belongs to a group
         for handler in global_leap_binder.setup_container.metadata:
             if requested_metadata_names:
                 if not is_metadata_name_starts_with_handler_name(handler):
                     continue
 
             if preprocess_state.is_grouped:
-                handler_result = handler.function(group_ids, preprocess_state)[pos]
+                # Single-sample metadata fetch: encode only THIS sample, never the whole group —
+                # same memory rationale as _get_dataset_handlers / the get_sample latent branch
+                # (grouped-fetch-oom-bug.md). Metadata values are not guaranteed small scalars: a
+                # metadata encoder may return an np.asarray or read the sample's heavy data to
+                # compute a statistic, so materializing all B rows just to index one out is the same
+                # OOM path this fix closes. The batched get_metadata_multiple_samples correctly keeps
+                # the whole-group encode (amortized across many ids via its group_cache); a single
+                # id does not need it. A singleton is a valid single-group subset of the contract.
+                handler_result = handler.function([sample_id], preprocess_state)[0]
             else:
                 handler_result = handler.function(sample_id, preprocess_state)
             if isinstance(handler_result, dict):
@@ -1245,7 +1285,7 @@ class LeapLoader(LeapLoaderBase):
             if requested_metadata_names and not is_wanted(handler):
                 continue
             rows = handler.function(group_ids, preprocess_state)
-            if rows and isinstance(rows[0], dict):
+            if len(rows) > 0 and isinstance(rows[0], dict):
                 for key in self._grouped_dict_keys(rows, handler.name):
                     name = f'{handler.name}_{key}'
                     if requested_metadata_names and name not in requested_metadata_names:

{code_loader-1.0.196.dev0 → code_loader-1.0.197}/code_loader/utils.py

@@ -1,4 +1,6 @@
+import io
 import math
+import pickle
 import sys
 from pathlib import Path
 from types import TracebackType
@@ -15,11 +17,16 @@ from code_loader.contract.enums import DatasetMetadataType
 def to_numpy_return_wrapper(encoder_function: SectionCallableInterface) -> SectionCallableInterface:
     def numpy_encoder_function(idx: Union[int, str], samples: PreprocessResponse) -> npt.NDArray[np.float32]:
         result = encoder_function(idx, samples)
-        if isinstance(result, list) and len(result) > 0 and all(isinstance(sample, np.ndarray) for sample in result):
-            # A list of per-sample arrays only comes from a grouped encoder (a flat encoder returns
-            # a single array). Never stack it into a (B, *) batch — grouping is only a storage
-            # read-unit; batching the group to the model's B is the engine's single responsibility.
-            return result
+        if isinstance(idx, list):
+            # Grouped call (idx is a list of sample ids): never stack the per-sample results into a
+            # (B, *) batch — grouping is only a storage read-unit; batching the group to the model's
+            # B is the engine's single responsibility. Gate on the call shape (idx is a list), not
+            # the result shape, so a flat encoder returning a list of arrays still stacks below.
+            # Normalize each sample independently (a user-stacked (B, *) array is split back into its
+            # rows) so a ragged group — variable per-sample shapes, e.g. token lists of differing
+            # length — is preserved instead of crashing np.array() on an inhomogeneous shape.
+            rows = result if isinstance(result, list) else np.asarray(result)
+            return [np.asarray(sample) for sample in rows]
         numpy_result: npt.NDArray[np.float32] = np.array(result)
         return numpy_result
 
@@ -109,27 +116,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.197/code_loader/visualizers/default_visualizers.py

@@ -0,0 +1,149 @@
+from enum import Enum
+
+import numpy as np
+import numpy.typing as npt
+
+from code_loader.contract.visualizer_classes import LeapImage, LeapGraph, LeapHorizontalBar, LeapText, \
+    LeapImageMask, LeapTextMask, LeapVideo, LeapAudio
+from code_loader.utils import rescale_min_max
+
+
+class DefaultVisualizer(Enum):
+    Image = 'Image'
+    Video = 'Video'
+    Audio = 'Audio'
+    Graph = 'Graph'
+    HorizontalBar = 'HorizontalBar'
+    Text = 'Text'
+    ImageMask = 'ImageMask'
+    TextMask = 'TextMask'
+    RawData = 'RawData'
+
+
+def default_image_visualizer(data: npt.NDArray[np.float32]) -> LeapImage:
+    rescaled_data = rescale_min_max(data[0])
+    return LeapImage(rescaled_data)
+
+
+def default_video_visualizer(data: npt.NDArray[np.float32]) -> LeapVideo:
+    return LeapVideo(data[0])
+
+
+_MAX_WAVEFORM_POINTS = 2000
+_VISUAL_KINDS = ('waveform', 'spectrogram', 'image')
+
+
+def _to_mono(audio: npt.NDArray) -> npt.NDArray:
+    a = np.squeeze(np.asarray(audio))
+    if a.ndim == 2:  # mix down: channels are the shorter axis
+        a = a.mean(axis=int(np.argmin(a.shape)))
+    if a.ndim != 1:
+        raise ValueError(f'audio must be reducible to a mono 1-D signal, got shape {np.asarray(audio).shape}')
+    return a.astype(np.float32)
+
+
+def _waveform_graph(series: npt.NDArray, duration: float) -> LeapGraph:
+    s = np.squeeze(np.asarray(series)).astype(np.float32)
+    # ponytail: recharts chokes on 100k+ points; stride down for the default view. x_range keeps full time span.
+    step = max(1, s.shape[0] // _MAX_WAVEFORM_POINTS)
+    s = s[::step].reshape(-1, 1)
+    return LeapGraph(s, x_label='Time [s]', y_label='Amplitude', x_range=(0.0, duration))
+
+
+def default_audio_visualizer(audio: npt.NDArray,
+                             sample_rate: int,
+                             visual: npt.NDArray,
+                             visual_kind: str,
+                             time_axis_index: int = -1) -> LeapAudio:
+    """Package a playable audio sample plus the visual a playhead runs through.
+
+    audio: waveform samples (mono [N], or multi-channel that is mixed down).
+    sample_rate: samples per second.
+    visual: the array to display. For 'waveform' a 1-D series; for 'spectrogram'
+        a raw 2-D [freq, time] (or [time, freq]) array to be rendered grayscale;
+        for 'image' an already-rendered [H, W, 1] or [H, W, 3] image (e.g. a
+        spectrogram the caller already ran through their own colormap).
+    visual_kind: 'waveform', 'spectrogram', or 'image'. The caller states this
+        explicitly; the shape of `visual` is never inspected to guess it.
+    time_axis_index: for 'spectrogram', which axis of `visual` is time (moved to columns).
+        Unused for 'waveform'/'image'.
+    """
+    if visual_kind not in _VISUAL_KINDS:
+        raise ValueError(f"visual_kind must be one of {_VISUAL_KINDS}, got {visual_kind!r}")
+
+    mono = _to_mono(audio)
+    duration = mono.shape[0] / float(sample_rate)
+    x_range = (0.0, duration)
+    v = np.asarray(visual)
+
+    if visual_kind == 'waveform':
+        return LeapAudio(mono, sample_rate, _waveform_graph(v, duration), x_range)
+
+    if visual_kind == 'image':
+        return LeapAudio(mono, sample_rate, LeapImage(v), x_range)
+
+    # spectrogram
+    spec = np.squeeze(v)
+    if spec.ndim != 2:
+        raise ValueError(f'spectrogram visual must be 2-D (freq x time), got shape {v.shape}')
+    if time_axis_index % spec.ndim == 0:  # move time onto the horizontal axis (columns)
+        spec = spec.T
+    img = rescale_min_max(spec.astype(np.float32))[..., None]  # [freq, time, 1]
+    return LeapAudio(mono, sample_rate, LeapImage(img), x_range)
+
+
+def default_graph_visualizer(data: npt.NDArray[np.float32]) -> LeapGraph:
+    return LeapGraph(data[0])
+
+
+def default_horizontal_bar_visualizer(data: npt.NDArray[np.float32]) -> LeapHorizontalBar:
+    labels = [str(index) for index in range(data.shape[-1])]
+    return LeapHorizontalBar(data[0], labels)
+
+
+def default_word_visualizer(data: npt.NDArray[np.float32]) -> LeapText:
+    if len(data.shape) == 2 and data.shape[0] == 1:
+        data = data[0]
+    if hasattr(data, 'tolist'):
+        data = data.tolist()
+    words = [str(index[0]) if type(index) is list else str(index) for index in data]
+    return LeapText(words)
+
+
+def default_raw_data_visualizer(data: npt.NDArray[np.float32]) -> LeapText:
+    return LeapText([str(data)])
+
+
+def default_image_mask_visualizer(mask: npt.NDArray[np.float32], image: npt.NDArray[np.float32]) -> LeapImageMask:
+    if image.shape[0] == 1:
+        image = image[0]
+    if mask.shape[0] == 1:
+        mask = mask[0]
+
+    n_different_labels = mask.shape[-1]
+    labels = [str(i) for i in range(n_different_labels)]
+
+    if len(mask.shape) > 2:
+        if mask.shape[-1] == 1:
+            mask = np.squeeze(mask, axis=-1)
+        else:
+            mask = np.argmax(mask, axis=-1)
+
+    return LeapImageMask(mask.astype(np.uint8), image.astype(np.float32), labels)
+
+
+def default_text_mask_visualizer(mask: npt.NDArray[np.float32], text_data: npt.NDArray[np.float32]) -> LeapTextMask:
+    mask = mask[0]
+    text_data = text_data[0]
+
+    words = default_word_visualizer(text_data).data
+    n_different_labels = mask.shape[-1]
+    labels = [str(i) for i in range(n_different_labels)]
+
+    if len(mask.shape) > 1:
+        if mask.shape[-1] == 1:
+            mask = np.squeeze(mask, axis=-1)
+        else:
+            mask = np.argmax(mask, axis=-1)
+
+    return LeapTextMask(mask.astype(np.uint8), words, labels)

{code_loader-1.0.196.dev0 → code_loader-1.0.197}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "code-loader"
-version = "1.0.196.dev"
+version = "1.0.197"
 description = ""
 authors = ["dorhar <doron.harnoy@tensorleap.ai>"]
 license = "MIT"

code_loader-1.0.196.dev0/code_loader/visualizers/default_visualizers.py

@@ -1,85 +0,0 @@
-from enum import Enum
-
-import numpy as np
-import numpy.typing as npt
-
-from code_loader.contract.visualizer_classes import LeapImage, LeapGraph, LeapHorizontalBar, LeapText, \
-    LeapImageMask, LeapTextMask, LeapVideo
-from code_loader.utils import rescale_min_max
-
-
-class DefaultVisualizer(Enum):
-    Image = 'Image'
-    Video = 'Video'
-    Graph = 'Graph'
-    HorizontalBar = 'HorizontalBar'
-    Text = 'Text'
-    ImageMask = 'ImageMask'
-    TextMask = 'TextMask'
-    RawData = 'RawData'
-    
-
-def default_image_visualizer(data: npt.NDArray[np.float32]) -> LeapImage:
-    rescaled_data = rescale_min_max(data[0])
-    return LeapImage(rescaled_data)
-
-
-def default_video_visualizer(data: npt.NDArray[np.float32]) -> LeapVideo:
-    return LeapVideo(data[0])
-
-
-def default_graph_visualizer(data: npt.NDArray[np.float32]) -> LeapGraph:
-    return LeapGraph(data[0])
-
-
-def default_horizontal_bar_visualizer(data: npt.NDArray[np.float32]) -> LeapHorizontalBar:
-    labels = [str(index) for index in range(data.shape[-1])]
-    return LeapHorizontalBar(data[0], labels)
-
-
-def default_word_visualizer(data: npt.NDArray[np.float32]) -> LeapText:
-    if len(data.shape) == 2 and data.shape[0] == 1:
-        data = data[0]
-    if hasattr(data, 'tolist'):
-        data = data.tolist()
-    words = [str(index[0]) if type(index) is list else str(index) for index in data]
-    return LeapText(words)
-
-
-def default_raw_data_visualizer(data: npt.NDArray[np.float32]) -> LeapText:
-    return LeapText([str(data)])
-
-
-def default_image_mask_visualizer(mask: npt.NDArray[np.float32], image: npt.NDArray[np.float32]) -> LeapImageMask:
-    if image.shape[0] == 1:
-        image = image[0]
-    if mask.shape[0] == 1:
-        mask = mask[0]
-
-    n_different_labels = mask.shape[-1]
-    labels = [str(i) for i in range(n_different_labels)]
-
-    if len(mask.shape) > 2:
-        if mask.shape[-1] == 1:
-            mask = np.squeeze(mask, axis=-1)
-        else:
-            mask = np.argmax(mask, axis=-1)
-
-    return LeapImageMask(mask.astype(np.uint8), image.astype(np.float32), labels)
-
-
-def default_text_mask_visualizer(mask: npt.NDArray[np.float32], text_data: npt.NDArray[np.float32]) -> LeapTextMask:
-    mask = mask[0]
-    text_data = text_data[0]
-
-    words = default_word_visualizer(text_data).data
-    n_different_labels = mask.shape[-1]
-    labels = [str(i) for i in range(n_different_labels)]
-
-    if len(mask.shape) > 1:
-        if mask.shape[-1] == 1:
-            mask = np.squeeze(mask, axis=-1)
-        else:
-            mask = np.argmax(mask, axis=-1)
-
-    return LeapTextMask(mask.astype(np.uint8), words, labels)