code-loader 1.0.182.dev0__py3-none-any.whl → 1.0.184__py3-none-any.whl

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, Union, Any, Type, Optional, Callable, Tuple
+from typing import Dict, List, Iterable, Set, Union, Any, Type, Optional, Callable, Tuple
 
 import numpy as np
 import numpy.typing as npt
@@ -44,6 +44,8 @@ class LeapLoader(LeapLoaderBase):
         super().__init__(code_path, code_entry_name)
 
         self._preprocess_result_cached = None
+        self._synthetic_lookup: Dict[str, Tuple[PreprocessResponse, Any]] = {}
+        self._synthetic_populator: Optional[Callable[[str], None]] = None
 
         try:
             from code_loader.mixpanel_tracker import track_code_loader_loaded
@@ -187,8 +189,62 @@ class LeapLoader(LeapLoaderBase):
             for prediction_type in setup.prediction_types
         }
 
+    def set_synthetic_populator(self, populator: Optional[Callable[[str], None]]) -> None:
+        # Hook called by `_resolve_synthetic` when a synthetic sample_id is requested
+        # but not yet in `_synthetic_lookup`. The populator should call `run_simulation`
+        # for the recipe that produced sample_id, or raise if the recipe can't be
+        # resolved. After the call, the entire batch the recipe produced is in
+        # `_synthetic_lookup`, so subsequent lookups for sibling sample_ids skip the hook.
+        self._synthetic_populator = populator
+
+    def _user_additional_ids_set(self) -> Set[str]:
+        # Membership-test helper for distinguishing user-defined additional samples
+        # from synthetic ones. In consumer pods, the populator passes
+        # extend_preprocess=False to run_simulation, so additional.sample_ids stays
+        # uncontaminated and direct membership check is reliable.
+        additional = self._preprocess_result().get(DataStateEnum.additional)
+        if additional is None:
+            return set()
+        return set(additional.sample_ids)
+
+    def _resolve_synthetic(self, sample_id: Union[int, str],
+                           state: Optional[DataStateEnum] = None
+                           ) -> Optional[Tuple[PreprocessResponse, Any]]:
+        # Resolution flow for state==additional:
+        #   1. already-replayed synthetic ID (member of _synthetic_lookup) → cached synthetic flow
+        #   2. user-defined ID (member of preprocess[additional].sample_ids) → return None
+        #      so caller falls through to normal preprocess flow (no recipe lookup built)
+        #   3. populator: builds recipe index lazily on first call, replays the missing
+        #      recipe, populates _synthetic_lookup. Loud-raises on unresolvable IDs.
+        #
+        # Order rationale: _synthetic_lookup check sits before the user-response check
+        # to be robust against `_extend_additional_preprocess` mutation (producer side
+        # calls run_simulation with extend_preprocess=True, which appends synthetic IDs
+        # into preprocess[additional]). Once an ID is in _synthetic_lookup, that wins
+        # regardless of contamination in the preprocess response.
+        if not isinstance(sample_id, str):
+            return None
+        if state != DataStateEnum.additional:
+            return None
+        if sample_id in self._synthetic_lookup:
+            return self._synthetic_lookup[sample_id]
+        if sample_id in self._user_additional_ids_set():
+            return None
+        if self._synthetic_populator is None:
+            return None
+        self._synthetic_populator(sample_id)
+        return self._synthetic_lookup.get(sample_id)
+
     def get_sample(self, state: DataStateEnum, sample_id: Union[int, str], instance_id: int = None) -> DatasetSample:
         self.exec_script()
+
+        resolved = self._resolve_synthetic(sample_id, state=state)
+        if resolved is not None:
+            sim_preprocess, original_local_id = resolved
+            return self._get_sample_from_preprocess(
+                sim_preprocess, original_local_id, synthetic_index=sample_id
+            )
+
         preprocess_result = self._preprocess_result()
         if state == DataStateEnum.unlabeled and sample_id not in preprocess_result[state].sample_ids:
             self._preprocess_result(update_unlabeled_preprocess=True)
@@ -211,6 +267,50 @@ class LeapLoader(LeapLoaderBase):
                                instance_masks=instance_mask)
         return sample
 
+    def _get_sample_from_preprocess(
+        self,
+        preprocess: "PreprocessResponse",
+        original_sample_id: Union[int, str],
+        synthetic_index: str,
+    ) -> DatasetSample:
+        inputs = {}
+        for handler in global_leap_binder.setup_container.inputs:
+            inputs[handler.name] = handler.function(original_sample_id, preprocess)
+
+        gt = {}
+        for handler in global_leap_binder.setup_container.ground_truths:
+            gt[handler.name] = handler.function(original_sample_id, preprocess)
+
+        metadata = {}
+        metadata_is_none = {}
+        for handler in global_leap_binder.setup_container.metadata:
+            handler_result = handler.function(original_sample_id, preprocess)
+            if isinstance(handler_result, dict):
+                for k, v in handler_result.items():
+                    key = "{}_{}".format(handler.name, k)
+                    metadata[key], metadata_is_none[key] = self._convert_metadata_to_correct_type(key, v)
+            else:
+                metadata[handler.name], metadata_is_none[handler.name] = self._convert_metadata_to_correct_type(
+                    handler.name, handler_result
+                )
+
+        custom_latent_space = None
+        if global_leap_binder.setup_container.custom_latent_space is not None:
+            custom_latent_space = global_leap_binder.setup_container.custom_latent_space.function(
+                original_sample_id, preprocess
+            )
+
+        return DatasetSample(
+            inputs=inputs,
+            gt=gt if gt else None,
+            metadata=metadata,
+            metadata_is_none=metadata_is_none,
+            index=synthetic_index,
+            state=DataStateEnum.additional,
+            custom_latent_space=custom_latent_space,
+            instance_masks=None,
+        )
+
     def check_dataset(self) -> DatasetIntegParseResult:
         test_payloads: List[DatasetTestResultPayload] = []
         setup_response = None
@@ -345,8 +445,13 @@ class LeapLoader(LeapLoaderBase):
             result_payloads.append(test_result)
         return result_payloads
 
-    def run_simulation(self, sim_name, params=None, n_samples=1, seed=0):
-        # type: (str, Optional[Dict[str, Any]], int, int) -> Dict[str, npt.NDArray[np.float32]]
+    def run_simulation(self, sim_name, params=None, n_samples=1, seed=0,
+                       sample_ids=None, extend_preprocess=True):
+        # type: (str, Optional[Dict[str, Any]], int, int, Optional[List[str]], bool) -> Dict[str, Any]
+        # extend_preprocess=True (default): also extend preprocess_result[additional] with the
+        # new synthetic sample_ids (producer-side behavior — synthetic worker needs this for
+        # enumeration / data_length tracking). Pass False from consumer-side populators that
+        # only need _synthetic_lookup filled, so user's additional response stays uncontaminated.
         self.exec_script()
         sim = next(
             (s for s in global_leap_binder.setup_container.simulations if s.name == sim_name),
@@ -364,11 +469,51 @@ class LeapLoader(LeapLoaderBase):
             _simulation_context["active"] = False
         sim_preprocess.state = DataStateType.additional
         sim_preprocess.tl_generated = True
+        original_sample_ids = list(sim_preprocess.sample_ids)
         per_encoder = {handler.name: [] for handler in global_leap_binder.setup_container.inputs}
-        for sample_id in sim_preprocess.sample_ids:
+        for sample_id in original_sample_ids:
             for handler in global_leap_binder.setup_container.inputs:
                 per_encoder[handler.name].append(handler.function(sample_id, sim_preprocess))
-        return {name: np.stack(arrays) for name, arrays in per_encoder.items()}
+        encoded = {name: np.stack(arrays) for name, arrays in per_encoder.items()}
+        if sample_ids is not None:
+            if len(sample_ids) != len(original_sample_ids):
+                raise ValueError(
+                    "sample_ids length ({}) does not match simulation output length ({})".format(
+                        len(sample_ids), len(original_sample_ids)
+                    )
+                )
+            for sid in sample_ids:
+                if not isinstance(sid, str):
+                    raise TypeError(
+                        "All sample_ids must be of type str. Got: {}".format(type(sid))
+                    )
+            for synth_id, original_local_id in zip(sample_ids, original_sample_ids):
+                self._synthetic_lookup[synth_id] = (sim_preprocess, original_local_id)
+            if extend_preprocess:
+                self._extend_additional_preprocess(list(sample_ids))
+            returned_sample_ids = list(sample_ids)
+        else:
+            returned_sample_ids = original_sample_ids
+        return {"encoded": encoded, "sample_ids": returned_sample_ids}
+
+    def _extend_additional_preprocess(self, new_sample_ids: List[str]) -> None:
+        if self._preprocess_result_cached is None:
+            self._preprocess_result()
+        additional = self._preprocess_result_cached.get(DataStateEnum.additional)
+        if additional is None:
+            placeholder = PreprocessResponse(
+                sample_ids=list(new_sample_ids),
+                data={},
+                state=DataStateType.additional,
+                sample_id_type=str,
+                tl_generated=True,
+            )
+            self._preprocess_result_cached[DataStateEnum.additional] = placeholder
+            global_leap_binder.setup_container.preprocess.data_length[DataStateType.additional] = placeholder.length
+        else:
+            additional.sample_ids = list(additional.sample_ids) + list(new_sample_ids)
+            additional.length = len(additional.sample_ids)
+            global_leap_binder.setup_container.preprocess.data_length[DataStateType.additional] = additional.length
 
     @staticmethod
     def _get_all_dataset_base_handlers() -> List[Union[DatasetBaseHandler, MetadataHandler]]:
@@ -583,7 +728,13 @@ class LeapLoader(LeapLoaderBase):
                               state: DataStateEnum, sample_id: Union[int, str]) -> Dict[str, npt.NDArray[np.float32]]:
         result_agg = {}
         preprocess_result = self._preprocess_result()
-        preprocess_state = preprocess_result[state]
+        resolved = self._resolve_synthetic(sample_id, state=state)
+        if resolved is not None:
+            sim_preprocess, original_local_id = resolved
+            preprocess_state = sim_preprocess
+            sample_id = original_local_id
+        else:
+            preprocess_state = preprocess_result[state]
         for handler in handlers:
             handler_result = handler.function(sample_id, preprocess_state)
             handler_name = handler.name
@@ -597,7 +748,13 @@ class LeapLoader(LeapLoaderBase):
         if instance_id is None:
             return None
         preprocess_result = self._preprocess_result()
-        preprocess_state = preprocess_result[state]
+        resolved = self._resolve_synthetic(sample_id, state=state)
+        if resolved is not None:
+            sim_preprocess, original_local_id = resolved
+            preprocess_state = sim_preprocess
+            sample_id = original_local_id
+        else:
+            preprocess_state = preprocess_result[state]
         result_agg = {}
         for handler in global_leap_binder.setup_container.instance_masks:
             handler_result = handler.function(sample_id, preprocess_state, instance_id)
@@ -657,7 +814,13 @@ class LeapLoader(LeapLoaderBase):
         result_agg = {}
         is_none = {}
         preprocess_result = self._preprocess_result()
-        preprocess_state = preprocess_result[state]
+        resolved = self._resolve_synthetic(sample_id, state=state)
+        if resolved is not None:
+            sim_preprocess, original_local_id = resolved
+            preprocess_state = sim_preprocess
+            sample_id = original_local_id
+        else:
+            preprocess_state = preprocess_result[state]
         for handler in global_leap_binder.setup_container.metadata:
             if requested_metadata_names:
                 if not is_metadata_name_starts_with_handler_name(handler):
@@ -698,17 +861,9 @@ class LeapLoader(LeapLoaderBase):
         return global_leap_binder.setup_container.custom_latent_space is not None
 
     def get_instances_data(self, state: DataStateEnum) -> Tuple[Dict[str, List[str]], Dict[str, str]]:
-        """
-        This Method get the data state and returns two dictionaries that holds the mapping of the sample ids to their
-         instances and the other way around and the sample ids array.
-        Args:
-            state: DataStateEnum state
-        Returns:
-            sample_ids_to_instance_mappings: sample id to instance mappings
-            instance_to_sample_ids_mappings: instance to sample ids mappings
-            sample_ids: sample ids array
-        """
         preprocess_result = self._preprocess_result()
-        preprocess_state = preprocess_result[state]
+        preprocess_state = preprocess_result.get(state)
+        if preprocess_state is None:
+            return None, None
         return preprocess_state.sample_ids_to_instance_mappings, preprocess_state.instance_to_sample_ids_mappings
 

code_loader/leaploaderbase.py

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

{code_loader-1.0.182.dev0.dist-info → code_loader-1.0.184.dist-info}/METADATA

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

{code_loader-1.0.182.dev0.dist-info → code_loader-1.0.184.dist-info}/RECORD

@@ -23,8 +23,8 @@ code_loader/experiment_api/workingspace_config_utils.py,sha256=DLzXQCg4dgTV_YgaS
 code_loader/inner_leap_binder/__init__.py,sha256=koOlJyMNYzGbEsoIbXathSmQ-L38N_pEXH_HvL7beXU,99
 code_loader/inner_leap_binder/leapbinder.py,sha256=XLYYcV50qjMvoC1S6WW0tLBch_0g5gl1UyHiVSWYbvg,40491
 code_loader/inner_leap_binder/leapbinder_decorators.py,sha256=YlK51b5Ryo396f7tOA7Ole3vYCLs3f5ZLm2qDQ9K1NE,105781
-code_loader/leaploader.py,sha256=Md-CsdLiBVfeGIlVqpVjNVF_fEAFbg2ahmVCCPl3iOw,36758
-code_loader/leaploaderbase.py,sha256=HrOGX9H8eRbyrFOEdvC2OISJvtZWKMVDah0cry8PXlo,6492
+code_loader/leaploader.py,sha256=K9Q5kiCZ_A2GkS5qOEantAMvWGaz3bmO1X8buXDSpgg,44682
+code_loader/leaploaderbase.py,sha256=l36qDA00GhZEG5NLKpEtAXgWJA-UQQIhNFGxywK7mUA,6530
 code_loader/mixpanel_tracker.py,sha256=rNwRmFifNbdUoqLQvvhhgpKczWpWiEmd8MfyJe27sxw,9131
 code_loader/plot_functions/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 code_loader/plot_functions/plot_functions.py,sha256=6Q7VWGxetL2W0EK2QeCdObVATvBuHs3YBA09H4uoIk0,14996
@@ -32,7 +32,7 @@ code_loader/plot_functions/visualize.py,sha256=gsBAYYkwMh7jIpJeDMPS8G4CW-pxwx6Lz
 code_loader/utils.py,sha256=YecipkdTA-VcE9F0RQcY9cFnY8P3AksPnHM2Db7xUSk,3972
 code_loader/visualizers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 code_loader/visualizers/default_visualizers.py,sha256=onRnLE_TXfgLN4o52hQIOOhUcFexGlqJ3xSpQDVLuZM,2604
-code_loader-1.0.182.dev0.dist-info/LICENSE,sha256=qIwWjdspQeSMTtnFZBC8MuT-95L02FPvzRUdWFxrwJY,1067
-code_loader-1.0.182.dev0.dist-info/METADATA,sha256=qmRxIYq5blbCPlBU8gavN-RNdC8fK3pxS-ckMVME2fI,1095
-code_loader-1.0.182.dev0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-code_loader-1.0.182.dev0.dist-info/RECORD,,
+code_loader-1.0.184.dist-info/LICENSE,sha256=qIwWjdspQeSMTtnFZBC8MuT-95L02FPvzRUdWFxrwJY,1067
+code_loader-1.0.184.dist-info/METADATA,sha256=iOj4cFbRfozrJ0PW7vnYNeTHxqDkT5F0IJHW8kbkrX4,1090
+code_loader-1.0.184.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+code_loader-1.0.184.dist-info/RECORD,,