fixed various typing and logic errors

Author: fabnemEPFL

src/mmirage/config/loading.py

@@ -47,7 +47,7 @@ def is_unresolved_env_var(s: str) -> bool:
                 if self.num_shards < 1:
                     raise ValueError()
             except (ValueError, TypeError):
-                if is_unresolved_env_var(self.num_shards):
+                if isinstance(self.num_shards, str) and is_unresolved_env_var(self.num_shards):
                     self.num_shards = 1
                 else:
                     raise ValueError(f"Invalid value for num_shards: {self.num_shards!r}")
@@ -56,7 +56,7 @@ def is_unresolved_env_var(s: str) -> bool:
             try:
                 self.shard_id = int(self.shard_id)
             except (ValueError, TypeError):
-                if is_unresolved_env_var(self.shard_id):
+                if isinstance(self.shard_id, str) and is_unresolved_env_var(self.shard_id):
                     self.shard_id = 0
                 else:
                     raise ValueError(f"Invalid value for shard_id: {self.shard_id!r}")

src/mmirage/core/process/base.py

@@ -72,7 +72,7 @@ def batch_process_sample(
         """
         raise NotImplementedError()
 
-    @abstract
+    @abc.abstractmethod
     def get_token_counts(self) -> TokenCounts:
         """Get cumulative token counts from this processor.
 
@@ -84,6 +84,18 @@ def get_token_counts(self) -> TokenCounts:
         """
         raise NotImplementedError()
 
+    @abc.abstractmethod
+    def get_load_time(self) -> float:
+        """Get the time taken to load any necessary resources (e.g., models).
+
+        Returns:
+            Time in seconds taken to load resources.
+
+        Raises:
+            NotImplementedError: If not implemented by subclass.
+        """
+        raise NotImplementedError()
+
 
 class ProcessorRegistry:
     """Registry for managing and accessing available processors.

src/mmirage/core/process/mapper.py

@@ -75,7 +75,7 @@ def validate_vars(self) -> bool:
     def rewrite_batch(
         self,
         batch: Dict[str, List[Any]],
-        image_base_path: str = None,
+        image_base_path: Optional[str] = None,
     ) -> List[VariableEnvironment]:
         """Transform a batch of samples by computing output variables.
 

src/mmirage/merge_shards.py

@@ -49,7 +49,7 @@ def _merge_datasetdict(shard_dsets: List[DatasetDict]) -> DatasetDict:
         merged[str(split)] = concatenate_datasets(split_dsets)
     if not merged:
         raise RuntimeError("All splits were empty after merging.")
-    return DatasetDict(merged)
+    return DatasetDict(**merged)
 
 
 def _merge_shards(shard_dsets: List[DatasetLike]) -> DatasetLike:

src/mmirage/shard_process.py

@@ -37,7 +37,7 @@ def rewrite_batch(
     batch: Dict[str, List[Any]],
     mapper: MMIRAGEMapper,
     renderer: TemplateRenderer,
-    image_base_path: str = None,
+    image_base_path: Optional[str] = None,
 ) -> Dict[str, List[Any]]:
     """Rewrite a batch of samples by applying transformations.
     Args:
@@ -91,6 +91,8 @@ def main():
 
     state_dir = shard_state_dir(shard_id, loading_params.get_state_root())
 
+    gpu_poller: Optional[GpuUtilizationPoller] = None
+
     collect_stats = os.environ.get("MMIRAGE_COLLECT_STATS", "") == "1"
     if collect_stats:
         # Determine which physical GPU indices SGLang will use so the poller
@@ -112,9 +114,11 @@ def main():
             gpu_indices_for_polling: List[str] = all_visible[:tp_size] if all_visible else [str(i) for i in range(tp_size)]
         else:
             gpu_indices_for_polling = [str(i) for i in range(tp_size)]
-        gpu_poller: GpuUtilizationPoller = GpuUtilizationPoller(
+
+        gpu_poller = GpuUtilizationPoller(
             interval_seconds=5.0, gpu_indices=gpu_indices_for_polling
         )
+
     try:
         retry_count = _mark_running(state_dir, shard_id, datasets_config)
         logger.info(f"Starting shard {shard_id}/{last_shard_id} (attempt #{retry_count})")
@@ -144,7 +148,7 @@ def main():
 
         # Start GPU polling after model loading so utilisation samples reflect
         # inference only, not weight transfers during sgl.Engine() init.
-        if collect_stats:
+        if collect_stats and gpu_poller is not None:
             gpu_poller.start()
 
         ds_processed_all: List[DatasetLike] = []
@@ -180,7 +184,7 @@ def main():
             _save_dataset_atomic(ds_processed, out_dir)
             logger.info(f"✅ Saved dataset {ds_idx} shard in: {out_dir}")
 
-        gpu_info = gpu_poller.stop() if collect_stats else {"mean": None, "min": None, "max": None, "samples": 0}
+        gpu_info = gpu_poller.stop() if collect_stats and gpu_poller is not None else {"mean": None, "min": None, "max": None, "samples": 0}
 
         # Collect token counts accumulated by LLM processor(s).
         token_counts = mapper.get_token_counts()
@@ -214,7 +218,7 @@ def main():
         error_msg = f"{type(e).__name__}: {str(e)}"
         logger.error(f"❌ Shard {shard_id} failed: {error_msg}")
         logger.error(traceback.format_exc())
-        if collect_stats:
+        if collect_stats and gpu_poller is not None:
             gpu_poller.stop()
         _mark_failure(state_dir, error_msg)
         sys.exit(1)