Function check_taskset_path

benchmark/bench.py:81–175

Ensures the taskset path exists for the given dataset; generates it if necessary. This function checks whether `taskset_path` exists. If not, it uses a corresponding data generation script (e.g., gen_countdown_data.py) to create the dataset at the default or provided location. The gener

(dataset_name: str, taskset_path: str)

def check_taskset_path(dataset_name: str, taskset_path: str) -> str:
    """Ensures the taskset path exists for the given dataset; generates it if necessary.

    This function checks whether `taskset_path` exists. If not,
    it uses a corresponding data generation script (e.g., gen_countdown_data.py) to create
    the dataset at the default or provided location. The generator scripts are expected
    to be located in the 'scripts/' subdirectory relative to this file.

    Args:
        dataset_name: Name of the dataset (e.g., "countdown", "guru").
            Must be one of the supported datasets defined in `dataset_script_map`.
        taskset_path: Path to the dataset.

    Returns:
        str: The resolved path to the dataset.

    Raises:
        ValueError: If the `dataset_name` is not supported.
        FileNotFoundError: If the corresponding generator script does not exist.
        ImportError: If the generator module fails to load.
        AttributeError: If the loaded module does not define 'DEFAULT_DATA_PATH'.
        subprocess.CalledProcessError: If the generation script fails (due to check=True).

    Side Effects:
        - May create directories and files on disk via the external generation script.
        - Executes a subprocess to run the dataset generation script.

    Examples:
        For dataset_name='guru_math' and taskset_path=None, this function will runs the
        following command and generate the guru_math dataset to default location
        (DEFAULT_DATA_PATH in scripts/gen_guru_math_data.py):

        ```bash
        python scripts/gen_guru_math_data.py --local_dir DEFAULT_DATA_PATH
        ```
    """
    if taskset_path:
        if os.path.exists(taskset_path):
            return taskset_path
        if dataset_name == "gsm8k" and taskset_path == "openai/gsm8k":
            return taskset_path

    base_dir = os.path.dirname(__file__)
    frozenlake_data_script_path = os.path.abspath(
        os.path.join(
            base_dir,
            "..",
            "examples",
            "grpo_frozen_lake",
            "get_frozen_lake_data.py",
        )
    )
    dataset_script_map = {
        "countdown": "gen_countdown_data.py",
        "guru_math": "gen_guru_math_data.py",
        "alfworld": "get_alfworld_full_data.py",
        "frozenlake": frozenlake_data_script_path,
    }