Merge pull request #9 from commit-0/hydra

OmegaConf integration

Author: wenting-zhao

commit0/harness/run_pytest_ids.py

@@ -1,4 +1,5 @@
 import argparse
+from datasets import load_dataset
 import docker
 from enum import StrEnum, auto
 import os
@@ -7,6 +8,9 @@
 from pathlib import Path
 import logging
 
+from omegaconf import DictConfig, OmegaConf
+import hydra
+
 from commit0.harness.constants import RUN_PYTEST_LOG_DIR
 from commit0.harness.docker_build import (
     close_logger,
@@ -32,7 +36,7 @@
 
 
 class ExecutionBackend(StrEnum):
-    DOCKER = auto()
+    LOCAL = auto()
     MODAL = auto()
 
 
@@ -194,72 +198,41 @@ def run_modal(
                 )
 
 
-def main(
-    repo: str,
-    test_ids_ls: list[str],
-    timeout: int,
-    branch_name: str,
-    backend: ExecutionBackend,
-) -> None:
-    with open("config.yml", "r") as file:
-        data = yaml.safe_load(file)
-    spec = make_spec(data["repos"][repo])
-    test_ids = " ".join(test_ids_ls)
-    hashed_test_ids = get_hash_string(test_ids)
+@hydra.main(version_base=None, config_path="configs", config_name="base")
+def main(config: DictConfig) -> None:
+    OmegaConf.to_yaml(config)
+    dataset = load_dataset(config.dataset_name, split="test")
+    for example in dataset:
+        if example["repo"].endswith(config.repo):
+            spec = make_spec(example)
+            break
 
+    hashed_test_ids = get_hash_string(config.test_ids)
     # set up logging
-    log_dir = RUN_PYTEST_LOG_DIR / repo / hashed_test_ids
+    log_dir = RUN_PYTEST_LOG_DIR / config.repo / hashed_test_ids
     log_dir.mkdir(parents=True, exist_ok=True)
     log_file = log_dir / "run_pytest.log"
-    logger = setup_logger(repo, log_file)
+    logger = setup_logger(config.repo, log_file)
 
     # make eval file
     eval_script = spec.eval_script.format(
-        local_repo=f"{data['base_repo_dir']}/{repo}",
-        branch_name=branch_name,
-        test_ids=test_ids,
-        ip=get_ip(data["backend"]),
+        local_repo=f"{config.base_dir}/{config.repo}",
+        branch_name=config.branch,
+        test_ids=config.test_ids,
+        ip=get_ip(config.backend),
         user=get_user(),
     )
     eval_file = Path(log_dir / "eval.sh")
     eval_file.write_text(eval_script)
 
-    if ExecutionBackend(backend) == ExecutionBackend.DOCKER:
-        run_docker(spec, logger, eval_file, timeout, log_dir)
-    elif ExecutionBackend(backend) == ExecutionBackend.MODAL:
-        run_modal(spec, logger, eval_file, timeout, log_dir)
-
-
-def add_init_args(parser: argparse.ArgumentParser) -> None:
-    parser.add_argument("--repo", type=str, help="which repo to run unit tests")
-    parser.add_argument(
-        "--test_ids", type=str, nargs="+", help="which test ids / files / directories"
-    )
-    parser.add_argument(
-        "--branch_name", type=str, help="which git branch to run unit tests"
-    )
-    parser.add_argument(
-        "--timeout",
-        type=int,
-        default=1_800,
-        help="Timeout (in seconds) for running tests for each instance",
-    )
-    parser.add_argument(
-        "--backend",
-        choices=[backend.value for backend in ExecutionBackend],
-        default=ExecutionBackend.DOCKER.value,
-        help="Execution backend [docker, modal]",
-    )
+    if ExecutionBackend(config.backend) == ExecutionBackend.LOCAL:
+        run_docker(spec, logger, eval_file, config.timeout, log_dir)
+    elif ExecutionBackend(config.backend) == ExecutionBackend.MODAL:
+        run_modal(spec, logger, eval_file, config.timeout, log_dir)
 
 
 def run(args: argparse.Namespace) -> None:
-    main(
-        repo=args.repo,
-        test_ids_ls=args.test_ids,
-        timeout=args.timeout,
-        branch_name=args.branch_name,
-        backend=args.backend,
-    )
+    main()
 
 
 __all__ = []

commit0/harness/setup.py

@@ -1,27 +1,27 @@
-import argparse
 import logging
 import os
 
 import docker
+import hydra
 import yaml
 from datasets import load_dataset
+from omegaconf import DictConfig, OmegaConf
+
 from typing import Iterator
 from commit0.harness.utils import clone_repo
 from commit0.harness.constants import REPO_IMAGE_BUILD_DIR, RepoInstance
 from commit0.harness.docker_build import build_repo_images
 from commit0.harness.spec import make_spec
 
+
 logging.basicConfig(
     level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
 )
 logger = logging.getLogger(__name__)
 
 
-def main(
-    hf_name: str,
-    base_dir: str,
-    config_file: str,
-) -> None:
+@hydra.main(version_base=None, config_path="configs", config_name="base")
+def main(config: DictConfig) -> None:
     dataset: Iterator[RepoInstance] = load_dataset(hf_name, split="test")  # type: ignore
     out = dict()
     specs = []
@@ -38,40 +38,15 @@ def main(
             clone_url, out[repo_name]["local_path"], example["base_commit"], logger
         )
 
-    config_file = os.path.abspath(config_file)
-    with open(config_file, "w") as f:
-        yaml.dump(out, f, default_flow_style=False)
-    logger.info(f"Config file has been written to {config_file}")
     logger.info("Start building docker images")
     logger.info(f"Please check {REPO_IMAGE_BUILD_DIR} for build details")
     client = docker.from_env()
     build_repo_images(client, specs)
     logger.info("Done building docker images")
 
 
-def add_init_args(parser: argparse.ArgumentParser) -> None:
-    parser.add_argument("--hf_name", type=str, help="HF dataset name")
-    parser.add_argument(
-        "--base_dir",
-        type=str,
-        default="repos/",
-        help="base directory to write repos to",
-    )
-    parser.add_argument(
-        "--config_file",
-        type=str,
-        default="config.yml",
-        help="where to write config file to",
-    )
-    parser.set_defaults(func=run)
-
-
-def run(args: argparse.Namespace) -> None:
-    main(
-        hf_name=args.hf_name,
-        base_dir=args.base_dir,
-        config_file=args.config_file,
-    )
+def run() -> None:
+    main()
 
 
 __all__ = []

pyproject.toml

@@ -23,6 +23,7 @@ dependencies = [
     "wget",
     "ruff>=0.6.4",
     "pre-commit>=3.8.0",
+    "hydra-core>=1.3.2",
     "modal>=0.64.95",
 ]
 scripts = { commit0 = "commit0.__main__:main" }