code-loader 1.0.94.dev11__py3-none-any.whl → 1.0.153.dev3__py3-none-any.whl

code_loader/leaploader.py

@@ -2,6 +2,7 @@
 import importlib.util
 import inspect
 import io
+import os
 import sys
 from contextlib import redirect_stdout
 from functools import lru_cache
@@ -15,7 +16,7 @@ from code_loader.contract.datasetclasses import DatasetSample, DatasetBaseHandle
     PreprocessResponse, VisualizerHandler, LeapData, \
     PredictionTypeHandler, MetadataHandler, CustomLayerHandler, MetricHandler, VisualizerHandlerData, MetricHandlerData, \
     MetricCallableReturnType, CustomLossHandlerData, CustomLossHandler, RawInputsForHeatmap, SamplePreprocessResponse, \
-    ElementInstance
+    ElementInstance, custom_latent_space_attribute
 from code_loader.contract.enums import DataStateEnum, TestingSectionEnum, DataStateType, DatasetMetadataType
 from code_loader.contract.exceptions import DatasetScriptException
 from code_loader.contract.responsedataclasses import DatasetIntegParseResult, DatasetTestResultPayload, \
@@ -23,6 +24,7 @@ from code_loader.contract.responsedataclasses import DatasetIntegParseResult, Da
     VisualizerInstance, PredictionTypeInstance, ModelSetup, CustomLayerInstance, MetricInstance, CustomLossInstance, \
     EngineFileContract
 from code_loader.inner_leap_binder import global_leap_binder
+from code_loader.inner_leap_binder.leapbinder import mapping_runtime_mode_env_var_mame
 from code_loader.leaploaderbase import LeapLoaderBase
 from code_loader.utils import get_root_exception_file_and_line_number
 
@@ -33,10 +35,23 @@ class LeapLoader(LeapLoaderBase):
 
         self._preprocess_result_cached = None
 
+        try:
+            from code_loader.mixpanel_tracker import track_code_loader_loaded
+            track_code_loader_loaded({
+                'event_type': 'leap_loader_instantiated',
+                'code_path': code_path,
+                'code_entry_name': code_entry_name
+            })
+        except Exception:
+            pass
+
     @lru_cache()
     def exec_script(self) -> None:
         try:
+            os.environ[mapping_runtime_mode_env_var_mame] = 'TRUE'
             self.evaluate_module()
+            if global_leap_binder.integration_test_func is not None:
+                global_leap_binder.integration_test_func(None, PreprocessResponse(state=DataStateType.training, length=0))
         except TypeError as e:
             import traceback
             if "leap_binder.set_metadata(" in traceback.format_exc(5):
@@ -45,6 +60,10 @@ class LeapLoader(LeapLoaderBase):
             raise DatasetScriptException(getattr(e, 'message', repr(e))) from e
         except Exception as e:
             raise DatasetScriptException(getattr(e, 'message', repr(e))) from e
+        finally:
+            # ensure that the environment variable is removed after the script execution
+            if mapping_runtime_mode_env_var_mame in os.environ:
+                del os.environ[mapping_runtime_mode_env_var_mame]
 
     def evaluate_module(self) -> None:
         def append_path_recursively(full_path: str) -> None:
@@ -58,13 +77,15 @@ class LeapLoader(LeapLoaderBase):
         file_path = Path(self.code_path, self.code_entry_name)
         append_path_recursively(str(file_path))
 
+        importlib.invalidate_caches()
+
         spec = importlib.util.spec_from_file_location(self.code_path, file_path)
         if spec is None or spec.loader is None:
-            raise DatasetScriptException(f'Something is went wrong with spec file from: {file_path}')
+            raise DatasetScriptException(f'Something went wrong with spec file from: {file_path}')
 
         file = importlib.util.module_from_spec(spec)
         if file is None:
-            raise DatasetScriptException(f'Something is went wrong with import module from: {file_path}')
+            raise DatasetScriptException(f'Something went wrong with import module from: {file_path}')
 
         spec.loader.exec_module(file)
 
@@ -136,35 +157,28 @@ class LeapLoader(LeapLoaderBase):
             for prediction_type in setup.prediction_types
         }
 
-    def get_sample(self, state: DataStateEnum, sample_id: Union[int, str]) -> DatasetSample:
+    def get_sample(self, state: DataStateEnum, sample_id: Union[int, str], instance_id: int = None) -> DatasetSample:
         self.exec_script()
         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)
 
-        metadata, metadata_is_none = self._get_metadata(state, sample_id)
-        sample = DatasetSample(inputs=self._get_inputs(state, sample_id),
-                               gt=None if state == DataStateEnum.unlabeled else self._get_gt(state, sample_id),
-                               metadata=metadata,
-                               metadata_is_none=metadata_is_none,
-                               index=sample_id,
-                               state=state)
-        return sample
-
-    def get_sample_with_masks(self, state: DataStateEnum, sample_id: Union[int, str]) -> DatasetSample:
-        self.exec_script()
-        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)
+        metadata, metadata_is_none = self.get_metadata(state, sample_id)
 
-        metadata, metadata_is_none = self._get_metadata(state, sample_id)
+        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(sample_id,
+                                                                                                  preprocess_result[
+                                                                                                      state])
+        instance_mask = self._get_instances_masks(state, sample_id, instance_id)
         sample = DatasetSample(inputs=self._get_inputs(state, sample_id),
                                gt=None if state == DataStateEnum.unlabeled else self._get_gt(state, sample_id),
                                metadata=metadata,
                                metadata_is_none=metadata_is_none,
                                index=sample_id,
                                state=state,
-                               instance_masks=self._get_instances_masks(state, sample_id))
+                               custom_latent_space=custom_latent_space,
+                               instance_masks=instance_mask)
         return sample
 
     def check_dataset(self) -> DatasetIntegParseResult:
@@ -175,6 +189,7 @@ class LeapLoader(LeapLoaderBase):
         with redirect_stdout(stdout_steam):
             try:
                 self.exec_script()
+
                 preprocess_test_payload = self._check_preprocess()
                 test_payloads.append(preprocess_test_payload)
                 handlers_test_payloads = self._check_handlers()
@@ -454,12 +469,14 @@ class LeapLoader(LeapLoaderBase):
     def _get_inputs(self, state: DataStateEnum, sample_id: Union[int, str]) -> Dict[str, npt.NDArray[np.float32]]:
         return self._get_dataset_handlers(global_leap_binder.setup_container.inputs, state, sample_id)
 
-    def _get_instances_masks(self, state: DataStateEnum, sample_id: Union[int, str]) -> Dict[str, List[ElementInstance]]:
+    def _get_instances_masks(self, state: DataStateEnum, sample_id: Union[int, str], instance_id: int) -> Optional[Dict[str, ElementInstance]]:
+        if instance_id is None:
+            return None
         preprocess_result = self._preprocess_result()
         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)
+            handler_result = handler.function(sample_id, preprocess_state, instance_id)
             handler_name = handler.name
             result_agg[handler_name] = handler_result
         return result_agg
@@ -504,20 +521,38 @@ class LeapLoader(LeapLoaderBase):
 
         return converted_value, is_none
 
-    def _get_metadata(self, state: DataStateEnum, sample_id: Union[int, str]) -> Tuple[Dict[str, Union[str, int, bool, float]], Dict[str, bool]]:
+    def get_metadata(self, state: DataStateEnum, sample_id: Union[int, str], requested_metadata_names: Optional[List[str]] = None) -> Tuple[
+        Dict[str, Union[str, int, bool, float]], Dict[str, bool]]:
+
+        def is_metadata_name_starts_with_handler_name(_handler):
+            for metadata_name in requested_metadata_names:
+                if metadata_name.startswith(_handler.name + '_') or metadata_name == _handler.name:
+                    return True
+            return False
+
         result_agg = {}
         is_none = {}
         preprocess_result = self._preprocess_result()
         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):
+                    continue
+
             handler_result = handler.function(sample_id, preprocess_state)
             if isinstance(handler_result, dict):
                 for single_metadata_name, single_metadata_result in handler_result.items():
                     handler_name = f'{handler.name}_{single_metadata_name}'
+                    if requested_metadata_names:
+                        if handler_name not in requested_metadata_names:
+                            continue
                     result_agg[handler_name], is_none[handler_name] = self._convert_metadata_to_correct_type(
                         handler_name, single_metadata_result)
             else:
                 handler_name = handler.name
+                if requested_metadata_names:
+                    if handler_name not in requested_metadata_names:
+                        continue
                 result_agg[handler_name], is_none[handler_name] = self._convert_metadata_to_correct_type(
                     handler_name, handler_result)
 
@@ -533,7 +568,12 @@ class LeapLoader(LeapLoaderBase):
 
         return id_type
 
-    def get_instances_data(self, state: DataStateEnum) -> Tuple[Dict[Union[int, str], List[Union[int, str]]], Dict[Union[int, str], Union[int, str]], List[Union[int, str]]]:
+    @lru_cache()
+    def has_custom_latent_space_decorator(self) -> bool:
+        self.exec_script()
+        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.
@@ -546,4 +586,4 @@ class LeapLoader(LeapLoaderBase):
         """
         preprocess_result = self._preprocess_result()
         preprocess_state = preprocess_result[state]
-        return preprocess_state.sample_ids_to_instance_mappings, preprocess_state.instance_to_sample_ids_mappings, preprocess_state.sample_ids
+        return preprocess_state.sample_ids_to_instance_mappings, preprocess_state.instance_to_sample_ids_mappings

code_loader/leaploaderbase.py

@@ -61,16 +61,30 @@ class LeapLoaderBase:
         pass
 
     @abstractmethod
-    def get_sample(self, state: DataStateEnum, sample_id: Union[int, str]) -> DatasetSample:
+    def get_sample(self, state: DataStateEnum, sample_id: Union[int, str], instance_id: int = None) -> DatasetSample:
         pass
 
     @abstractmethod
-    def get_sample_with_masks(self, state: DataStateEnum, sample_id: Union[int, str]) -> DatasetSample:
+    def get_instances_data(self, state: DataStateEnum) -> Tuple[Dict[str, List[str]], Dict[str, str]]:
         pass
 
-    @abstractmethod
-    def get_instances_data(self, state: DataStateEnum) -> Tuple[Dict[Union[int, str], List[Union[int, str]]], Dict[Union[int, str], Union[int, str]], List[Union[int, str]]]:
-        pass
+    def get_metadata_multiple_samples(self, state: DataStateEnum, sample_ids: Union[List[int], List[str]],
+                                      requested_metadata_names: Optional[List[str]] = None
+                                      ) -> Tuple[Dict[str, Union[List[str], List[int], List[bool],
+    List[float]]], Dict[str, List[bool]]]:
+        aggregated_results: Dict[str, List[Union[str, int, bool, float]]] = {}
+        aggregated_is_none: Dict[str, List[bool]] = {}
+        sample_id_type = self.get_sample_id_type()
+        for sample_id in sample_ids:
+            sample_id = sample_id_type(sample_id)
+            metadata_result, is_none_result = self.get_metadata(state, sample_id, requested_metadata_names)
+            for metadata_name, metadata_value in metadata_result.items():
+                if metadata_name not in aggregated_results:
+                    aggregated_results[metadata_name] = []
+                    aggregated_is_none[metadata_name] = []
+                aggregated_results[metadata_name].append(metadata_value)
+                aggregated_is_none[metadata_name].append(is_none_result[metadata_name])
+        return aggregated_results, aggregated_is_none
 
     @abstractmethod
     def check_dataset(self) -> DatasetIntegParseResult:
@@ -91,6 +105,13 @@ class LeapLoaderBase:
                         input_tensors_by_arg_name: Dict[str, npt.NDArray[np.float32]]):
         pass
 
+    @abstractmethod
+    def get_metadata(
+            self, state: DataStateEnum, sample_id: Union[int, str],
+            requested_metadata_names: Optional[List[str]] = None
+    ) -> Tuple[Dict[str, Union[str, int, bool, float]], Dict[str, bool]]:
+        pass
+
     @abstractmethod
     def run_heatmap_visualizer(self, visualizer_name: str, sample_ids: np.array, state: DataStateEnum,
                                input_tensors_by_arg_name: Dict[str, npt.NDArray[np.float32]]
@@ -114,6 +135,10 @@ class LeapLoaderBase:
     def get_sample_id_type(self) -> Type:
         pass
 
+    @abstractmethod
+    def has_custom_latent_space_decorator(self) -> bool:
+        pass
+
     @abstractmethod
     def get_heatmap_visualizer_raw_vis_input_arg_name(self, visualizer_name: str) -> Optional[str]:
         pass

code_loader/mixpanel_tracker.py

@@ -0,0 +1,230 @@
+"""
+Mixpanel tracking utilities for code-loader.
+"""
+import os
+import sys
+import getpass
+import uuid
+import logging
+from enum import Enum
+from typing import Optional, Dict, Any, Set, Union, TypedDict
+import mixpanel  # type: ignore[import]
+
+logger = logging.getLogger(__name__)
+
+TRACKING_VERSION = '1'
+
+
+class AnalyticsEvent(str, Enum):
+    """Enumeration of all tracked analytics events."""
+    CODE_LOADER_LOADED = "code_loader_loaded"
+    LOAD_MODEL_INTEGRATION_TEST = "load_model_integration_test"
+    PREPROCESS_INTEGRATION_TEST = "preprocess_integration_test"
+    INPUT_ENCODER_INTEGRATION_TEST = "input_encoder_integration_test"
+    GT_ENCODER_INTEGRATION_TEST = "gt_encoder_integration_test"
+
+
+class CodeLoaderLoadedProps(TypedDict, total=False):
+    """Properties for code_loader_loaded event."""
+    event_type: str
+    code_path: str
+    code_entry_name: str
+
+
+class LoadModelEventProps(TypedDict, total=False):
+    """Properties for load_model_integration_test event."""
+    prediction_types_count: int
+
+
+class PreprocessEventProps(TypedDict, total=False):
+    """Properties for preprocess_integration_test event."""
+    preprocess_responses_count: int
+
+
+class InputEncoderEventProps(TypedDict, total=False):
+    """Properties for input_encoder_integration_test event."""
+    encoder_name: str
+    channel_dim: int
+
+
+class GtEncoderEventProps(TypedDict, total=False):
+    """Properties for gt_encoder_integration_test event."""
+    encoder_name: str
+
+
+class MixpanelTracker:
+    """Handles Mixpanel event tracking for code-loader."""
+    
+    def __init__(self, token: str = "0c1710c9656bbfb1056bb46093e23ca1"):
+        self.token = token
+        self.mp = mixpanel.Mixpanel(token)
+        self._user_id: Optional[str] = None
+    
+    def _get_whoami(self) -> str:
+        """Get the current system username (whoami) for device identification.
+        
+        Returns:
+            str: The system username, with fallbacks to environment variables or 'unknown'
+        """
+        if self._user_id is None:
+            try:
+                self._user_id = getpass.getuser()
+            except Exception as e:
+                logger.debug(f"Failed to get username via getpass: {e}")
+                # Fallback to environment variables or default
+                self._user_id = os.environ.get('USER', os.environ.get('USERNAME', 'unknown'))
+        return self._user_id or 'unknown'
+    
+       
+    def _get_tensorleap_user_id(self) -> Optional[str]:
+        """Get the TensorLeap user ID from ~/.tensorleap/user_id if it exists."""
+        try:
+            user_id_path = os.path.expanduser("~/.tensorleap/user_id")
+            if os.path.exists(user_id_path):
+                with open(user_id_path, 'r') as f:
+                    user_id = f.read().strip()
+                    if user_id:
+                        return user_id
+        except Exception as e:
+            logger.debug(f"Failed to read TensorLeap user ID: {e}")
+        return None
+    
+    def _get_or_create_device_id(self) -> str:
+        """Get or create a device ID from ~/.tensorleap/device_id file.
+        
+        If the file doesn't exist, creates it with a new UUID.
+        
+        Returns:
+            str: The device ID (UUID string)
+        """
+        try:
+            device_id_path = os.path.expanduser("~/.tensorleap/device_id")
+            
+            # Create directory if it doesn't exist
+            os.makedirs(os.path.dirname(device_id_path), exist_ok=True)
+            
+            if os.path.exists(device_id_path):
+                with open(device_id_path, 'r') as f:
+                    device_id = f.read().strip()
+                    if device_id:
+                        return device_id
+            
+            # Generate new device ID and save it
+            device_id = str(uuid.uuid4())
+            with open(device_id_path, 'w') as f:
+                f.write(device_id)
+            
+            return device_id
+        except Exception as e:
+            logger.debug(f"Failed to read/write device ID file: {e}")
+            # Fallback to generating a new UUID if file operations fail
+            return str(uuid.uuid4())
+    
+    def _get_distinct_id(self) -> str:
+        """Get the distinct ID for Mixpanel tracking.
+        
+        Priority order:
+        1. TensorLeap user ID (from ~/.tensorleap/user_id)
+        2. Device ID (from ~/.tensorleap/device_id, generated if not exists)
+        """
+        tensorleap_user_id = self._get_tensorleap_user_id()
+        if tensorleap_user_id:
+            return tensorleap_user_id
+        
+        return self._get_or_create_device_id()
+    
+    def _track_event(self, event_name: Union[str, AnalyticsEvent], event_properties: Optional[Dict[str, Any]] = None) -> None:
+        """Internal method to track any event with device identification.
+        
+        Args:
+            event_name: The name of the event to track (string or AnalyticsEvent enum)
+            event_properties: Optional additional properties to include in the event
+        """
+        # Skip tracking if IS_TENSORLEAP_PLATFORM environment variable is set to 'true'
+        if os.environ.get('IS_TENSORLEAP_PLATFORM') == 'true':
+            return
+            
+        try:
+            distinct_id = self._get_distinct_id()
+            
+            tensorleap_user_id = self._get_tensorleap_user_id()
+            whoami = self._get_whoami()
+            device_id = self._get_or_create_device_id()
+            
+            properties = {
+                'tracking_version': TRACKING_VERSION,
+                'service': 'code-loader',
+                'whoami': whoami,
+                '$device_id': device_id,  # Always use device_id for $device_id
+                'python_version': f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}",
+                'platform': os.name,
+            }
+            
+            if tensorleap_user_id:
+                properties['user_id'] = tensorleap_user_id
+            
+            if event_properties:
+                properties.update(event_properties)
+            
+            self.mp.track(distinct_id, str(event_name), properties)
+        except Exception as e:
+            logger.debug(f"Failed to track event '{event_name}': {e}")
+
+    def track_code_loader_loaded(self, event_properties: Optional[Dict[str, Any]] = None) -> None:
+        """Track code loader loaded event with device identification.
+        
+        Args:
+            event_properties: Optional additional properties to include in the event
+        """
+        self._track_event(AnalyticsEvent.CODE_LOADER_LOADED, event_properties)
+
+    def track_integration_test_event(self, event_name: Union[str, AnalyticsEvent], event_properties: Optional[Dict[str, Any]] = None) -> None:
+        """Track an integration test event with device identification.
+        
+        Args:
+            event_name: The name of the event to track (string or AnalyticsEvent enum)
+            event_properties: Optional additional properties to include in the event
+        """
+        self._track_event(event_name, event_properties)
+
+
+# Global tracker instance
+_tracker = None
+
+
+def get_tracker() -> MixpanelTracker:
+    global _tracker
+    if _tracker is None:
+        _tracker = MixpanelTracker()
+    return _tracker
+
+
+def track_code_loader_loaded(event_properties: Optional[Dict[str, Any]] = None) -> None:
+    get_tracker().track_code_loader_loaded(event_properties)
+
+
+def track_integration_test_event(event_name: Union[str, AnalyticsEvent], event_properties: Optional[Dict[str, Any]] = None) -> None:
+    get_tracker().track_integration_test_event(event_name, event_properties)
+
+
+# Module-level set to track which integration test events have been emitted
+_integration_events_emitted: Set[str] = set()
+
+
+def emit_integration_event_once(event_name: Union[str, AnalyticsEvent], props: Dict[str, Any]) -> None:
+    """Emit an integration test event only once per test run."""
+    event_name_str = str(event_name)
+    if event_name_str in _integration_events_emitted:
+        return
+    
+    try:
+        track_integration_test_event(event_name, props)
+        _integration_events_emitted.add(event_name_str)
+    except Exception as e:
+        logger.debug(f"Failed to emit integration event once '{event_name}': {e}")
+
+
+def clear_integration_events() -> None:
+    """Clear the integration events set for a new test run."""
+    global _integration_events_emitted
+    _integration_events_emitted.clear()