Repository: open-compass/opencompass
Issue: [RFC]Support multi dataset tasks running as a merged dataset task
Root cause: Unimplemented Feature Gap


## Motivation.

The current task partitioning strategy of the OpenCompass command line uses the `NumWorkerPartitioner`. This strategy merges all datasets specified in the command into a single dataset collection, which is then split into `max-num-workers` sub-dataset tasks. Each sub-dataset task further contains multiple atomic dataset tasks. Each partitioned sub-dataset task is paired with a model task and executed as one `task`. The logic of the task's `run` method during execution is as follows:

```python
def run(self, cur_model=None, cur_model_abbr=None):
    self.logger.info(f'Task {task_abbr_from_cfg(self.cfg)}')
    for model_cfg, dataset_cfgs in zip(self.model_cfgs, self.dataset_cfgs):
        # ......
        for dataset_cfg in dataset_cfgs: # Each atomic task within the sub-dataset task
            # ......
            self._inference() # By default, initializes the retriever and calls `BaseInferencer.inference`
```

Each atomic task within a sub-dataset task initializes the retriever independently and calls `BaseInferencer.inference` once. This means that within a partitioned sub-dataset task, each atomic dataset task serially loops through the "load -> execute" process. When there are many atomic datasets but each is small in scale (e.g., the C-Eval dataset has 52 atomic datasets, but each has only about 20 cases on average), the overall execution efficiency is very low. If the `inferencer` configurations for all dataset tasks are identical, it is theoretically possible to merge these dataset tasks for joint execution.

## Proposed Change.

We aim to implement a method to merge multiple atomic dataset tasks within a partitioned sub-dataset task (only those with identical inferencer configurations will be merged) for execution by the `BaseInferencer`, without altering the structure of any inference result files generated within the `predictions/` directory.

### Key Implementation Class Diagram

<img width="1775" height="975" alt="Image" src="https://github.com/user-attachments/assets/fbc22719-f5da-4bf2-acb9-3978cb54204d" />

- The implementation of the `OpenICLInferTask.run_datasets` method reuses the logic from `OpenICLInferTask.run` that iterates over `dataset_cfgs`:

    ```python
    def run(self, cur_model=None, cur_model_abbr=None):
        self.logger.info(f'Task {task_abbr_from_cfg(self.cfg)}')
        for model_cfg, dataset_cfgs in zip(self.model_cfgs, self.dataset_cfgs):
            # ......
            # The following code logic is extracted into the OpenICLInferTask.run_datasets method
            for dataset_cfg in dataset_cfgs:
                self.model_cfg = model_cfg
                self.dataset_cfg = dataset_cfg
                self.infer_cfg = self.dataset_cfg['infer_cfg']
                self.dataset = build_dataset_from_cfg(self.dataset_cfg)
                self.sub_cfg = {
                    'models': [self.model_cfg],
                    'datasets': [[self.dataset_cfg]],
                }
                out_path = get_infer_output_path(
                    self.model_cfg, self.dataset_cfg,
                    osp.join(self.work_dir, 'predictions'))
                if osp.exists(out_path):
                    continue
                self._inference()
            # The above code logic is extracted into the OpenICLInferTask.run_datasets method
    ```

- The `OpenICLInferTask.run_merged_datasets` method serves as the entry point for merged dataset inference.
- Define a `BaseInferencer.inference_multi_datasets` method in the base class, which raises an exception indicating that the current type of icl_inferencer does not support merged dataset inference (this can first be implemented in a few commonly used icl_inferencers, such as `GenInferencer`).

### How to Trigger the Merged Dataset Logic

<img width="1737" height="987" alt="Image" src="https://github.com/user-attachments/assets/c0c8c4a5-d185-406b-ac83-bbadd9c1896f" />

- The decision to use a command-line argument to trigger this feature is primarily to allow it to act directly within the task, bypassing runner logic, thus avoiding modifications to various runners.
- A feasible implementation for passing the command-line argument into the cfg is as follows:

    ```python
    # opencompass/utils/run.py:get_config_from_arg
    def get_config_from_arg(args):
        # ......
        return Config(dict(models=models, datasets=datasets, summarizer=summarizer, cli_args=vars(args)), format_python_code=False)
    ```

### `inference_multi_datasets` Implementation Logic (Pseudocode, exemplified with `GenInferencer` method)

```python
def inference_multi_datasets(self,
                  retrievers: list[BaseRetriever],
                  ice_templates: list[Optional[PromptTemplate]] = [],
                  prompt_templates: list[Optional[PromptTemplate]] = [],
                  output_json_filepaths: list[Optional[str]] = [],
                  output_json_filenames: list[Optional[str]] = []) -> List:

        # Use os.path.join(output_json_filepaths[i], output_json_filenames[i]) as the unique identifier for the dataset
        output_handlers_map = {} # Maps dataset identifiers to their output handlers. Key: os.path.join(output_json_filepaths[i], output_json_filenames[i]), Value: GenInferencerOutputHandler
        merged_prompt_list = [] # All merged prompts to be actually executed. Value: prompt
        prompt_owner_list = [] # The atomic dataset identifier corresponding to each prompt in merged_prompt_list. Value: os.path.join(output_json_filepaths[i], output_json_filenames[i])

        for i, retriever in enumerate(retrievers):
            output_json_filepath = output_json_filepaths[i]
            output_json_filename = output_json_filenames[i]
            ice_template = ice_templates[i]
            prompt_template = prompt_templates[i]
            output_handler = GenInferencerOutputHandler()

            # Link output_handler in the map
            output_handlers_map[os.path.join(output_json_filepath, output_json_filename)] = output_handler

            # 2. Get results of retrieval process
            # ...... (same as in the `inference` function)

            # 3. Generate prompts for testing input
            # ...... (same as in the `inference` function)

            # 3.1 Fetch and zip prompt & gold answer if output column exists
            # ......

            # Create tmp json file for saving intermediate results and future
            # resuming
            # ...... (same as in the `inference` function)

            prompt_list # list of prompts
            index # index of the first prompt of this retriever (likely the starting index in merged_prompt_list)

            # Add unprocessed prompt_list to merged_prompt_list
            merged_prompt_list.extend(prompt_list[index:])
            # Add corresponding owners to prompt_owner_list
            prompt_owner_list.extend([os.path.join(output_json_filepath, output_json_filename)] * len(prompt_list[index:]))

        # 4. Wrap prompts with Dataloader
        dataloader = self.get_dataloader(merged_prompt_list, self.batch_size)

        # 5. Inference for prompts in each batch
        logger.info('Starting inference process...')

        start_time_stamp = time.time()
        num_sample = 0
        # Note: The original code uses 'index' which might be reused; renaming to 'prompt_idx_in_batch' for clarity within the loop.
        # However, based on the original logic, 'index' here likely refers to the position within the current batch's processed prompts.
        # A more accurate variable might be 'global_prompt_index' to track position in merged_prompt_list.
        # For strict adherence to original pseudocode structure, we'll keep 'index' but note the potential ambiguity.
        # It seems 'index' here is intended to track the progress through the merged_prompt_list/gold_list.
        # Let's assume 'index' starts at 0 for the entire batch processing loop.
        index = 0 
        for batch_index, datum in enumerate(tqdm(dataloader, disable=not self.is_main_process)):
            # Assuming 'datum' is a list of (prompt_entry, gold_answer) tuples
            entry, golds = list(zip(*datum)) 
            
            # 5-1. Inference with local model
            # (same as in the `inference` function)
            # ... (code to generate 'generated' text from 'entry')
            # Assume 'parsed_entries' corresponds to the processed 'entry' ready for output handling
            # Assume 'generated' is the list of model outputs for this batch
            parsed_entries = entry 
            
            # 5-3. Save current output
            # Note: The original code uses 'batched(generated, num_return_sequences)', assuming 'generated' is flat.
            # This part assumes 'num_return_sequences' is handled correctly, possibly within the inference step above.
            for prompt, prediction, gold in zip(
                    parsed_entries, generated, # Adjusted: 'batched' removed as per loop context, assuming 'generated' is structured correctly
                    golds):
                if num_return_sequences == 1:
                    prediction = prediction[0]
                
                # Use the prompt's linked output_handler to save results
                # Find the owner (dataset ID) of the current prompt
                current_prompt_owner_id = prompt_owner_list[index]
                output_handler_for_prompt = output_handlers_map.get(current_prompt_owner_id)
                
                output_handler_for_prompt.save_results(
                    prompt,
                    prediction,
                    # The 'index' here is likely the position within the original dataset's prompt list,
                    # not the global index in merged_prompt_list. This might need a separate tracker.
                    # For now, assuming the original logic intends to use the global 'index' for intermediate saving.
                    index, 
                    gold=gold
                )
                index += 1

            # 5-4. Save intermediate results
            for output_handler in output_handlers_map.values():
                if (self.save_every is not None and len(output_handler.result_dict) % self.save_every == 0
                    and self.is_main_process):
                    # Note: Original uses 'output_json_filepath' from loop variable, but here we use the handler's knowledge or stored path.
                    # Assuming the handler has this information or it's passed correctly.
                    # For pseudocode clarity, we'll use placeholders.
                    handler_filepath = output_handler.output_path # This is an assumption; adjust based on actual implementation.
                    handler_filename = output_handler.output_filename # This is an assumption; adjust based on actual implementation.
                    output_handler.write_to_json(handler_filepath, f'tmp_{handler_filename}')

            num_sample += len(datum)

        end_time_stamp = time.time()

        # 6. Output
        if self.is_main_process:
            # Note: Original loop uses 'output_json_filepaths', but we need to iterate through handlers.
            for output_handler_id, output_handler in output_handlers_map.items():
                # ...... (same as in the `inference` function for final output)
                # Again, assuming handlers have path info.
                handler_filepath = output_handler.output_path
                handler_filename = output_handler.output_filename
                output_handler.write_to_json(
                    handler_filepath,
                    handler_filename
                )
                # ...... (same as in the `inference` function)
        
        if self.dump_timer and self.is_main_process:
            for output_handler_id, output_handler in output_handlers_map.items():
                pass
                # ...... (same as in the `inference` function)

        # Compile and return all results
        results = []
        for output_handler in output_handlers_map.values():
            # Assuming 'results_dict' values contain the 'prediction' key
            results.extend([sample['prediction']
            for sample in output_handler.result_dict.values()])

        return results
```

- To ensure that the JSON files in the `predictions` directory remain unchanged, each atomic dataset task is managed by an independent `output_handler`. Therefore, the main modification in `inference_multi_datasets` is to ensure that when all sub-datasets are merged into a single `datum` for inference, the inference results are still managed by the `output_handler` corresponding to each original atomic dataset task.
- Additionally, the reading and generation of temporary files related to checkpoint resumption are also handled by the `output_handler` corresponding to each atomic dataset task.

## Feedback Period.

1~2 weeks

## CC List.

## Any Other Things.

### Will you implement it?

- [ ] I would like to implement this feature and create a PR!