code-loader 1.0.184.dev2__tar.gz → 1.0.184.dev4__tar.gz

{code_loader-1.0.184.dev2 → code_loader-1.0.184.dev4}/PKG-INFO

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

{code_loader-1.0.184.dev2 → code_loader-1.0.184.dev4}/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
@@ -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,58 @@ 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()
 
-        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, 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
             )
@@ -397,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, sample_ids=None):
-        # type: (str, Optional[Dict[str, Any]], int, int, Optional[List[str]]) -> Dict[str, Any]
+    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),
@@ -436,7 +489,8 @@ class LeapLoader(LeapLoaderBase):
                     )
             for synth_id, original_local_id in zip(sample_ids, original_sample_ids):
                 self._synthetic_lookup[synth_id] = (sim_preprocess, original_local_id)
-            self._extend_additional_preprocess(list(sample_ids))
+            if extend_preprocess:
+                self._extend_additional_preprocess(list(sample_ids))
             returned_sample_ids = list(sample_ids)
         else:
             returned_sample_ids = original_sample_ids
@@ -674,9 +728,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, state=state)
+        if resolved is not None:
+            sim_preprocess, original_local_id = resolved
             preprocess_state = sim_preprocess
             sample_id = original_local_id
         else:
@@ -694,9 +748,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, state=state)
+        if resolved is not None:
+            sim_preprocess, original_local_id = resolved
             preprocess_state = sim_preprocess
             sample_id = original_local_id
         else:
@@ -760,9 +814,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, state=state)
+        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 → code_loader-1.0.184.dev4}/pyproject.toml

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