scripts/train_all_tasks.py

import argparse
import copy
import json
import os
import re
import sys
from pathlib import Path
from typing import Any

from omegaconf import OmegaConf

if "src" not in sys.path:
    sys.path.insert(0, "src")

from mixhub.data.data import DATA_CATALOG
from run_model import main as run_model_main


ROOT_DIR = Path(__file__).resolve().parents[1]
CONFIG_DIR = ROOT_DIR / "config"
OUTPUTS_DIR = ROOT_DIR / "outputs"
DEFAULT_BASE_CONFIG = CONFIG_DIR / "example.yaml"

SMILES_HASH_DATASETS = {
    "calisol23",
    "mixturesoldb",
    "electrolytomics-conductivity",
    "designsolvents-density",
    "designsolvents-viscosity",
    "designsolvents-melting",
}


def normalize_token(value: str) -> str:
    normalized = re.sub(r"[^a-z0-9]+", "_", value.lower()).strip("_")
    return normalized or "task"


def build_task_key(dataset_name: str, property_name: str) -> str:
    return f"{normalize_token(dataset_name)}__{normalize_token(property_name)}"


def choose_device(device_arg: str) -> str:
    if device_arg != "auto":
        return device_arg

    import torch

    return "cuda" if torch.cuda.is_available() else "cpu"


def discover_tasks() -> tuple[list[dict[str, str]], list[dict[str, str]]]:
    tasks: list[dict[str, str]] = []
    skipped: list[dict[str, str]] = []

    for dataset_key, dataset_cls in DATA_CATALOG.items():
        try:
            dataset = dataset_cls()
        except FileNotFoundError as exc:
            skipped.append(
                {
                    "dataset": dataset_key,
                    "reason": str(exc),
                }
            )
            continue

        for property_name in dataset.properties:
            tasks.append(
                {
                    "dataset": dataset_key,
                    "property": property_name,
                    "task_key": build_task_key(dataset_key, property_name),
                }
            )

    tasks.sort(key=lambda item: item["task_key"])
    skipped.sort(key=lambda item: item["dataset"])
    return tasks, skipped


def build_task_config(
    base_config: Any,
    dataset_name: str,
    property_name: str,
    device: str,
):
    config = copy.deepcopy(base_config)
    config.dataset.name = dataset_name
    config.dataset.property = property_name
    if dataset_name in SMILES_HASH_DATASETS:
        config.dataset.featurization = "smiles_hash_features"
    config.device = device
    config.root_dir = str(OUTPUTS_DIR)
    return config


def write_task_config(task: dict[str, str], config) -> Path:
    CONFIG_DIR.mkdir(parents=True, exist_ok=True)
    config_path = CONFIG_DIR / f"{task['task_key']}.yaml"
    OmegaConf.save(config, config_path)
    return config_path


def write_manifest(tasks: list[dict[str, str]], skipped: list[dict[str, str]]) -> Path:
    OUTPUTS_DIR.mkdir(parents=True, exist_ok=True)
    manifest_path = OUTPUTS_DIR / "task_manifest.json"
    manifest = {
        "tasks": tasks,
        "skipped_datasets": skipped,
    }
    manifest_path.write_text(json.dumps(manifest, indent=2, ensure_ascii=False), encoding="utf-8")
    return manifest_path


def should_select_task(
    task: dict[str, str],
    dataset_filters: set[str],
    task_filters: set[str],
) -> bool:
    if dataset_filters and task["dataset"] not in dataset_filters:
        return False
    if task_filters and task["task_key"] not in task_filters:
        return False
    return True


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Generate configs and train one prediction checkpoint for every available CheMixHub task."
    )
    parser.add_argument(
        "--base-config",
        type=Path,
        default=DEFAULT_BASE_CONFIG,
        help="Base YAML config used as the template for every task.",
    )
    parser.add_argument(
        "--device",
        type=str,
        default="auto",
        help='Training device. Use "auto" to prefer CUDA when available, otherwise CPU.',
    )
    parser.add_argument(
        "--dataset",
        action="append",
        default=[],
        help="Restrict training to one or more dataset keys, e.g. --dataset miscible-solvent.",
    )
    parser.add_argument(
        "--task",
        action="append",
        default=[],
        help="Restrict training to one or more generated task keys.",
    )
    parser.add_argument(
        "--list",
        action="store_true",
        help="List discovered tasks and exit.",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Generate configs and manifest without launching training.",
    )
    parser.add_argument(
        "--skip-existing",
        action="store_true",
        help="Skip tasks whose checkpoint already exists in outputs/.",
    )
    parser.add_argument(
        "--limit",
        type=int,
        default=None,
        help="Only process the first N selected tasks.",
    )

    args = parser.parse_args()

    if not args.base_config.exists():
        raise FileNotFoundError(f"Base config does not exist: {args.base_config}")

    base_config = OmegaConf.load(args.base_config)
    device = choose_device(args.device)

    tasks, skipped = discover_tasks()
    dataset_filters = set(args.dataset)
    task_filters = set(args.task)
    selected_tasks = [task for task in tasks if should_select_task(task, dataset_filters, task_filters)]

    if args.limit is not None:
        selected_tasks = selected_tasks[: args.limit]

    if args.list:
        print("Discovered tasks:")
        for task in selected_tasks:
            print(f"- {task['task_key']}: dataset={task['dataset']} property={task['property']}")
        if skipped:
            print("\nSkipped datasets:")
            for item in skipped:
                print(f"- {item['dataset']}: {item['reason']}")
        return

    if not selected_tasks:
        raise ValueError("No tasks selected for training.")

    manifest_path = write_manifest(tasks=selected_tasks, skipped=skipped)
    print(f"Wrote task manifest to {manifest_path}")

    failures: list[dict[str, str]] = []
    completed: list[str] = []

    for index, task in enumerate(selected_tasks, start=1):
        checkpoint_path = OUTPUTS_DIR / f"best_model_dict_{task['task_key']}.pt"
        if args.skip_existing and checkpoint_path.exists():
            print(f"[{index}/{len(selected_tasks)}] Skipping {task['task_key']} because {checkpoint_path.name} already exists")
            completed.append(task["task_key"])
            continue

        config = build_task_config(
            base_config=base_config,
            dataset_name=task["dataset"],
            property_name=task["property"],
            device=device,
        )
        config_path = write_task_config(task, config)
        print(
            f"[{index}/{len(selected_tasks)}] "
            f"Prepared {task['task_key']} "
            f"(dataset={task['dataset']}, property={task['property']}, device={device})"
        )
        print(f"Config: {config_path}")

        if args.dry_run:
            completed.append(task["task_key"])
            continue

        try:
            run_model_main(
                config=config,
                experiment_name=task["task_key"],
                wandb_logger=None,
            )
            completed.append(task["task_key"])
        except Exception as exc:
            failures.append(
                {
                    "task_key": task["task_key"],
                    "dataset": task["dataset"],
                    "property": task["property"],
                    "error": str(exc),
                }
            )
            print(f"Training failed for {task['task_key']}: {exc}")

    summary = {
        "completed": completed,
        "failed": failures,
        "device": device,
        "base_config": str(args.base_config),
    }
    summary_path = OUTPUTS_DIR / "train_all_tasks_summary.json"
    summary_path.write_text(json.dumps(summary, indent=2, ensure_ascii=False), encoding="utf-8")
    print(f"Wrote training summary to {summary_path}")

    if failures:
        raise RuntimeError(f"{len(failures)} task(s) failed. See {summary_path} for details.")


if __name__ == "__main__":
    main()