suggestion by Copilot, mostly changing the CLI arguement to accept list of files

Author: legstar67

src/mmirage/core/process/batch/collector.py

@@ -13,19 +13,26 @@
 from mmirage.core.process.batch.registry import BatchAdapterFactory
 
 
-def _read_metadata_records(metadata_output_path: str) -> List[Dict[str, Any]]:
+def _normalize_metadata_paths(metadata_paths: str | Sequence[str]) -> List[str]:
+    if isinstance(metadata_paths, str):
+        return [metadata_paths]
+    return [str(path) for path in metadata_paths]
+
+
+def _read_metadata_records(metadata_output_paths: str | Sequence[str]) -> List[Dict[str, Any]]:
     records: List[Dict[str, Any]] = []
-    with open(metadata_output_path, "r", encoding="utf-8") as f:
-        for line in f:
-            raw = line.strip()
-            if not raw:
-                continue
-            try:
-                parsed = json.loads(raw)
-            except json.JSONDecodeError:
-                continue
-            if isinstance(parsed, dict):
-                records.append(parsed)
+    for metadata_output_path in _normalize_metadata_paths(metadata_output_paths):
+        with open(metadata_output_path, "r", encoding="utf-8") as f:
+            for line in f:
+                raw = line.strip()
+                if not raw:
+                    continue
+                try:
+                    parsed = json.loads(raw)
+                except json.JSONDecodeError:
+                    continue
+                if isinstance(parsed, dict):
+                    records.append(parsed)
     return records
 
 
@@ -56,7 +63,7 @@ def _aggregate_batch_mappings(
 
 
 def collect_and_merge(
-    metadata_output_path: str,
+    metadata_output_path: str | Sequence[str],
     provider_configs: Mapping[str, BatchProviderConfig],
     output_path: str,
 ) -> List[Dict[str, Any]]:
@@ -161,8 +168,9 @@ def _build_arg_parser() -> argparse.ArgumentParser:
     )
     parser.add_argument(
         "--metadata-path",
+        nargs="+",
         required=True,
-        help="Path to metadata JSONL receipt file.",
+        help="Path(s) to metadata JSONL receipt file(s). Supports multiple files.",
     )
     parser.add_argument(
         "--output-path",

src/mmirage/core/process/batch/status_checker.py

@@ -14,42 +14,55 @@
 from mmirage.core.process.batch.registry import BatchAdapterFactory
 
 
-def extract_unique_provider_batches(metadata_output_path: str) -> List[Tuple[str, str]]:
+def _normalize_metadata_paths(metadata_paths: str | Sequence[str]) -> List[str]:
+    if isinstance(metadata_paths, str):
+        return [metadata_paths]
+    return [str(path) for path in metadata_paths]
+
+
+def _read_metadata_records(metadata_output_paths: str | Sequence[str]) -> List[Dict[str, str]]:
+    records: List[Dict[str, str]] = []
+    for metadata_output_path in _normalize_metadata_paths(metadata_output_paths):
+        with open(metadata_output_path, "r", encoding="utf-8") as f:
+            for line in f:
+                raw = line.strip()
+                if not raw:
+                    continue
+                try:
+                    record = json.loads(raw)
+                except json.JSONDecodeError:
+                    continue
+                if isinstance(record, dict):
+                    records.append(record)
+    return records
+
+
+def extract_unique_provider_batches(metadata_output_path: str | Sequence[str]) -> List[Tuple[str, str]]:
     """Parse metadata JSONL and return unique ``(provider, provider_batch_id)`` pairs.
 
     Malformed lines and records missing required keys are skipped safely.
     """
     unique_pairs: List[Tuple[str, str]] = []
     seen = set()
 
-    with open(metadata_output_path, "r", encoding="utf-8") as f:
-        for line in f:
-            raw = line.strip()
-            if not raw:
-                continue
+    for record in _read_metadata_records(metadata_output_path):
+        provider = str(record.get("provider", "")).strip().lower()
+        provider_batch_id = str(record.get("provider_batch_id", "")).strip()
 
-            try:
-                record = json.loads(raw)
-            except json.JSONDecodeError:
-                continue
-
-            provider = str(record.get("provider", "")).strip().lower()
-            provider_batch_id = str(record.get("provider_batch_id", "")).strip()
-
-            if not provider or not provider_batch_id:
-                continue
+        if not provider or not provider_batch_id:
+            continue
 
-            pair = (provider, provider_batch_id)
-            if pair in seen:
-                continue
-            seen.add(pair)
-            unique_pairs.append(pair)
+        pair = (provider, provider_batch_id)
+        if pair in seen:
+            continue
+        seen.add(pair)
+        unique_pairs.append(pair)
 
     return unique_pairs
 
 
 def run_status_checker(
-    metadata_output_path: str,
+    metadata_output_path: str | Sequence[str],
     provider_configs: Mapping[str, BatchProviderConfig],
     output: TextIO = sys.stdout,
 ) -> List[BatchSubmissionResult]:
@@ -84,7 +97,7 @@ def run_status_checker(
 
 
 def _build_provider_configs_from_metadata(
-    metadata_output_path: str,
+    metadata_output_path: str | Sequence[str],
 ) -> Dict[str, BatchProviderConfig]:
     provider_names = {provider for provider, _ in extract_unique_provider_batches(metadata_output_path)}
     configs: Dict[str, BatchProviderConfig] = {}
@@ -104,8 +117,9 @@ def _build_arg_parser() -> argparse.ArgumentParser:
     parser = argparse.ArgumentParser(description="Check provider batch statuses from metadata receipts.")
     parser.add_argument(
         "--metadata-path",
+        nargs="+",
         required=True,
-        help="Path to metadata JSONL receipt file.",
+        help="Path(s) to metadata JSONL receipt file(s). Supports multiple files.",
     )
     return parser
 

src/mmirage/core/process/processors/llm/config.py

@@ -18,20 +18,19 @@
 
 def _parse_tp_size_from_env() -> int:
     """Parse tensor parallelism size from SLURM_GPUS_ON_NODE environment variable.
-    
+
     Defensively parses the environment variable, handling invalid values:
     - Returns 1 if the variable is None or empty
     - Strips whitespace before parsing
     - Returns 1 for non-integer values
     - Returns 1 for values <= 0
-    
+
     Returns:
         Tensor parallelism size (>= 1), defaults to 1 on any parsing error.
     """
     env_value = os.environ.get("SLURM_GPUS_ON_NODE")
     if not env_value:
         return 1
-    
     try:
         tp_size = int(env_value.strip())
         # Ensure tp_size is positive (must be >= 1)

src/mmirage/core/process/processors/llm/llm_processor.py

@@ -63,7 +63,7 @@ def __init__(self, engine_args: SGLangLLMConfig, **kwargs) -> None:
         """
         super().__init__(engine_args, **kwargs)
         provider_cfg_raw = dict(getattr(engine_args, "batch_provider", {}) or {})
-        batch_mode_requested = bool(provider_cfg_raw.get("enabled", False))
+        batch_mode_requested = bool(provider_cfg_raw.get("enabled", True))
 
         # In provider-batch mode we only build payloads/metadata and should not
         # initialize GPU-backed SGLang runtime.
@@ -82,14 +82,15 @@ def __init__(self, engine_args: SGLangLLMConfig, **kwargs) -> None:
         self._text_orchestrator: Optional[BatchSubmissionOrchestrator] = None
         self._multimodal_orchestrator: Optional[BatchSubmissionOrchestrator] = None
         self._batch_request_counter = 0
+        self._global_row_offset = 0
         self._setup_batch_runtime()
 
     def _setup_batch_runtime(self) -> None:
         provider_cfg_raw = dict(getattr(self.config, "batch_provider", {}) or {})
         if not provider_cfg_raw:
             return
 
-        if not provider_cfg_raw.get("enabled", False):
+        if not provider_cfg_raw.get("enabled", True):
             return
 
         provider = str(provider_cfg_raw.get("provider", "openai")).strip().lower()
@@ -390,7 +391,7 @@ def _batch_process_sample(
                     config=self._batch_provider_config,
                 )
                 requests.append(dict(request))
-                source_indices.append(self._batch_request_counter)
+                source_indices.append(self._global_row_offset + global_i)
 
             self._text_orchestrator.add_requests(
                 requests=requests,
@@ -440,7 +441,7 @@ def _batch_process_sample(
                     config=self._batch_provider_config,
                 )
                 requests.append(dict(request))
-                source_indices.append(self._batch_request_counter)
+                source_indices.append(self._global_row_offset + global_i)
 
             self._multimodal_orchestrator.add_requests(
                 requests=requests,
@@ -458,6 +459,8 @@ def _batch_process_sample(
             placeholder = f"__BATCH_SUBMITTED__:{unique_id}"
             placeholders.append(batch[i].with_variable(output_var.name, placeholder))
 
+        self._global_row_offset += nb_samples
+
         return placeholders
 
     def finalize(self) -> None: