code-loader 1.0.184.dev2__py3-none-any.whl → 1.0.184.dev3__py3-none-any.whl

code_loader/leaploader.py

@@ -45,6 +45,7 @@ class LeapLoader(LeapLoaderBase):
 
         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
@@ -188,11 +189,30 @@ 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 _resolve_synthetic(self, sample_id: Union[int, str]) -> Optional[Tuple[PreprocessResponse, Any]]:
+        if not isinstance(sample_id, str):
+            return None
+        if sample_id in self._synthetic_lookup:
+            return self._synthetic_lookup[sample_id]
+        if self._synthetic_populator is not None:
+            self._synthetic_populator(sample_id)
+            return self._synthetic_lookup.get(sample_id)
+        return None
+
     def get_sample(self, state: DataStateEnum, sample_id: Union[int, str], instance_id: int = None) -> DatasetSample:
         self.exec_script()
 
-        if isinstance(sample_id, str) and sample_id in self._synthetic_lookup:
-            sim_preprocess, original_local_id = self._synthetic_lookup[sample_id]
+        resolved = self._resolve_synthetic(sample_id)
+        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
             )
@@ -674,9 +694,9 @@ class LeapLoader(LeapLoaderBase):
                               state: DataStateEnum, sample_id: Union[int, str]) -> Dict[str, npt.NDArray[np.float32]]:
         result_agg = {}
         preprocess_result = self._preprocess_result()
-        if (state == DataStateEnum.additional and isinstance(sample_id, str)
-                and sample_id in self._synthetic_lookup):
-            sim_preprocess, original_local_id = self._synthetic_lookup[sample_id]
+        resolved = self._resolve_synthetic(sample_id) if state == DataStateEnum.additional else None
+        if resolved is not None:
+            sim_preprocess, original_local_id = resolved
             preprocess_state = sim_preprocess
             sample_id = original_local_id
         else:
@@ -694,9 +714,9 @@ class LeapLoader(LeapLoaderBase):
         if instance_id is None:
             return None
         preprocess_result = self._preprocess_result()
-        if (state == DataStateEnum.additional and isinstance(sample_id, str)
-                and sample_id in self._synthetic_lookup):
-            sim_preprocess, original_local_id = self._synthetic_lookup[sample_id]
+        resolved = self._resolve_synthetic(sample_id) if state == DataStateEnum.additional else None
+        if resolved is not None:
+            sim_preprocess, original_local_id = resolved
             preprocess_state = sim_preprocess
             sample_id = original_local_id
         else:
@@ -760,9 +780,9 @@ class LeapLoader(LeapLoaderBase):
         result_agg = {}
         is_none = {}
         preprocess_result = self._preprocess_result()
-        if (state == DataStateEnum.additional and isinstance(sample_id, str)
-                and sample_id in self._synthetic_lookup):
-            sim_preprocess, original_local_id = self._synthetic_lookup[sample_id]
+        resolved = self._resolve_synthetic(sample_id) if state == DataStateEnum.additional else None
+        if resolved is not None:
+            sim_preprocess, original_local_id = resolved
             preprocess_state = sim_preprocess
             sample_id = original_local_id
         else:

{code_loader-1.0.184.dev2.dist-info → code_loader-1.0.184.dev3.dist-info}/METADATA

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

{code_loader-1.0.184.dev2.dist-info → code_loader-1.0.184.dev3.dist-info}/RECORD

@@ -23,7 +23,7 @@ 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=70enWfzlCfp4OkWQYMIXJ__V7aoJQfgiKSHrikRnpTc,41600
+code_loader/leaploader.py,sha256=OuN53z2KCO-I0FRxwjB_mfyzNZizH5Ud4vrZyFNn1wc,42595
 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
@@ -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.184.dev2.dist-info/LICENSE,sha256=qIwWjdspQeSMTtnFZBC8MuT-95L02FPvzRUdWFxrwJY,1067
-code_loader-1.0.184.dev2.dist-info/METADATA,sha256=ye_o6x8GQGs0cQsHnOIh27A8TRnCRq64p_6XwslFsv4,1095
-code_loader-1.0.184.dev2.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-code_loader-1.0.184.dev2.dist-info/RECORD,,
+code_loader-1.0.184.dev3.dist-info/LICENSE,sha256=qIwWjdspQeSMTtnFZBC8MuT-95L02FPvzRUdWFxrwJY,1067
+code_loader-1.0.184.dev3.dist-info/METADATA,sha256=lH0fmNWqwyUFhJXXu8ATIV8vhwOhU7y6u6WfWS3ckfw,1095
+code_loader-1.0.184.dev3.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+code_loader-1.0.184.dev3.dist-info/RECORD,,