pre-commit

Author: wenting-zhao

commit0/__main__.py

@@ -31,7 +31,14 @@ def main() -> None:
     # after hydra gets all configs, put command-line arguments back
     sys.argv = sys_argv
     # repo_split: split from command line has a higher priority than split in hydra
-    if command in ["clone", "build", "setup-git-user", "evaluate", "evaluate-reference", "save"]:
+    if command in [
+        "clone",
+        "build",
+        "setup-git-user",
+        "evaluate",
+        "evaluate-reference",
+        "save",
+    ]:
         if len(sys.argv) >= 3:
             if sys.argv[2] not in SPLIT:
                 raise ValueError(

commit0/harness/build.py

@@ -7,7 +7,7 @@
 from tqdm import tqdm
 from typing import Iterator
 
-from commit0.harness.constants import RepoInstance, SPLIT
+from commit0.harness.constants import EVAL_BACKENDS, RepoInstance, SPLIT
 from commit0.harness.docker_build import build_repo_images
 from commit0.harness.execution_context import (
     ExecutionBackend,
@@ -23,7 +23,12 @@
 
 
 def main(
-    dataset_name: str, dataset_split: str, repo_split: str, num_workers: int, backend: str, key_path: str
+    dataset_name: str,
+    dataset_split: str,
+    repo_split: str,
+    num_workers: int,
+    backend: str,
+    key_path: str,
 ) -> None:
     dataset: Iterator[RepoInstance] = load_dataset(dataset_name, split=dataset_split)  # type: ignore
     specs = []
@@ -40,6 +45,10 @@ def main(
         client = docker.from_env()
         build_repo_images(client, specs, num_workers)
         execution_context = Docker
+    else:
+        raise ValueError(
+            f"Evaluation must be from {', '.join(EVAL_BACKENDS)}, but {backend} is provided."
+        )
 
     # get ssh key from each docker image
     img2key = dict()
@@ -49,12 +58,9 @@ def main(
                 key = context.get_ssh_pubkey_from_remote(user="root")
                 img2key[spec.repo_image_key] = key
         except Exception as e:
-            error_msg = (
-                f"General error: {e}\n"
-                f"{traceback.format_exc()}\n"
-            )
+            error_msg = f"General error: {e}\n" f"{traceback.format_exc()}\n"
             raise RuntimeError(error_msg)
-    with open(key_path, 'w') as json_file:
+    with open(key_path, "w") as json_file:
         json.dump(img2key, json_file, indent=4)
 
 

commit0/harness/docker_utils.py

@@ -8,7 +8,6 @@
 import threading
 import time
 import traceback
-import pwd
 from pathlib import Path
 from io import BytesIO
 from typing import Optional, List, Union

commit0/harness/execution_context.py

@@ -7,7 +7,6 @@
 from abc import ABC, abstractmethod
 import docker
 import logging
-import os
 import modal
 import modal.io_streams
 from enum import StrEnum, auto
@@ -90,7 +89,9 @@ def delete_file_from_remote(self, remote_path: Path) -> None:
         """Delete"""
         raise NotImplementedError
 
-    def write_test_output(self, log_dir: Path, test_output: str, timed_out: bool) -> None:
+    def write_test_output(
+        self, log_dir: Path, test_output: str, timed_out: bool
+    ) -> None:
         """Write test output"""
         test_output_path = log_dir / "test_output.txt"
         with open(test_output_path, "w") as f:
@@ -145,9 +146,9 @@ def __init__(
         self.container.start()
         if files_to_copy:
             for _, f in files_to_copy.items():
-                copy_to_container(self.container, f['src'], Path(f['dest']))
+                copy_to_container(self.container, f["src"], f["dest"])  # type: ignore
 
-    def get_ssh_pubkey_from_remote(self, user: str) -> None:
+    def get_ssh_pubkey_from_remote(self, user: str) -> str:
         """Copy"""
         return get_ssh_pubkey_from_container(self.container, user)
 
@@ -195,7 +196,7 @@ def __init__(
         image = modal.Image.from_registry(image_name)
         if files_to_copy:
             for _, f in files_to_copy.items():
-                image = image.copy_local_file(f['src'], f['dest'])
+                image = image.copy_local_file(f["src"], f["dest"])  # type: ignore
 
         self.sandbox = modal.Sandbox.create(
             "sleep",
@@ -205,7 +206,7 @@ def __init__(
             timeout=timeout,
         )
 
-    def get_ssh_pubkey_from_remote(self, user: str) -> None:
+    def get_ssh_pubkey_from_remote(self, user: str) -> str:
         """Copy ssh pubkey"""
         process = self.sandbox.exec("bash", "-c", f"cat /{user}/.ssh/id_rsa.pub")
         public_key = "".join([line for line in process.stdout]).strip()

commit0/harness/run_pytest_ids.py

@@ -4,7 +4,12 @@
 
 from typing import Iterator
 from git import Repo
-from commit0.harness.constants import Files, RUN_PYTEST_LOG_DIR, RepoInstance
+from commit0.harness.constants import (
+    EVAL_BACKENDS,
+    Files,
+    RUN_PYTEST_LOG_DIR,
+    RepoInstance,
+)
 from commit0.harness.docker_build import (
     setup_logger,
 )
@@ -76,6 +81,10 @@ def main(
         execution_context = Modal
     elif ExecutionBackend(backend) == ExecutionBackend.LOCAL:
         execution_context = Docker
+    else:
+        raise ValueError(
+            f"Evaluation must be from {', '.join(EVAL_BACKENDS)}, but {backend} is provided."
+        )
 
     files_to_copy = Files(eval_script={"src": eval_file, "dest": Path("/eval.sh")})
 

commit0/harness/setup_git_user.py

@@ -1,7 +1,6 @@
 import json
 import logging
 import os
-import traceback
 
 from datasets import load_dataset
 
@@ -22,11 +21,18 @@
 logger = logging.getLogger(__name__)
 
 
-def main(dataset_name: str, dataset_split: str, repo_split: str, base_dir: str, git_user: str, key_path: str) -> None:
+def main(
+    dataset_name: str,
+    dataset_split: str,
+    repo_split: str,
+    base_dir: str,
+    git_user: str,
+    key_path: str,
+) -> None:
     dataset: Iterator[RepoInstance] = load_dataset(dataset_name, split=dataset_split)  # type: ignore
     setup_user(git_user, logger)
     setup_ssh_directory(git_user, logger)
-    with open(key_path, 'r') as f:
+    with open(key_path, "r") as f:
         public_keys = json.load(f)
 
     for example in dataset:

commit0/harness/utils.py

@@ -114,36 +114,42 @@ def setup_user(user: str, logger: logging.Logger) -> None:
 
 
 def chmod(path: str, mode: int, logger: logging.Logger) -> None:
-    """
-    A Python wrapper for the chmod command to change file or directory permissions.
+    """A Python wrapper for the chmod command to change file or directory permissions.
 
     Args:
+    ----
         path (str): The path to the file or directory.
         mode (int): The permission mode (octal), e.g., 0o755, 0o644, etc.
         logger (logging.Logger): The logger object.
+
     """
     try:
         os.chmod(path, mode)
         logger.info(f"Permissions for '{path}' changed to {oct(mode)}")
     except FileNotFoundError:
-        raise FileNotFoundError(f"Error: The file or directory '{path}' does not exist.")
+        raise FileNotFoundError(
+            f"Error: The file or directory '{path}' does not exist."
+        )
     except PermissionError:
-        raise PermissionError(f"Error: Permission denied when changing permissions for '{path}'")
+        raise PermissionError(
+            f"Error: Permission denied when changing permissions for '{path}'"
+        )
     except Exception as e:
         raise Exception(f"An error occurred: {e}")
 
 
 def setup_ssh_directory(user: str, logger: logging.Logger) -> None:
-    """
-    Sets up the .ssh directory for the user and sets appropriate permissions.
+    """Sets up the .ssh directory for the user and sets appropriate permissions.
 
     Args:
+    ----
         user (str): The name of the user.
         logger (logging.Logger): The logger object.
+
     """
     home = get_home_directory(user)
-    ssh_dir = os.path.join(home, '.ssh')
-    authorized_keys_file = os.path.join(ssh_dir, 'authorized_keys')
+    ssh_dir = os.path.join(home, ".ssh")
+    authorized_keys_file = os.path.join(ssh_dir, "authorized_keys")
 
     try:
         # Create the .ssh directory if it doesn't exist
@@ -156,19 +162,23 @@ def setup_ssh_directory(user: str, logger: logging.Logger) -> None:
 
         # Create the authorized_keys file if it doesn't exist
         if not os.path.exists(authorized_keys_file):
-            open(authorized_keys_file, 'a').close()
+            open(authorized_keys_file, "a").close()
             logger.info(f"Created file: {authorized_keys_file}")
     except Exception as e:
         raise e
 
 
 def add_key(user: str, public_key: str) -> None:
-    public_key = f"no-port-forwarding,no-X11-forwarding,no-agent-forwarding,no-pty {public_key}"
+    public_key = (
+        f"no-port-forwarding,no-X11-forwarding,no-agent-forwarding,no-pty {public_key}"
+    )
 
     home_directory = get_home_directory(user)
     authorized_keys_path = os.path.join(home_directory, ".ssh", "authorized_keys")
     if not os.path.exists(authorized_keys_path):
-        raise FileNotFoundError(f"f{authorized_keys_path} does not exists, please call setup_ssh_directory() before adding keys")
+        raise FileNotFoundError(
+            f"f{authorized_keys_path} does not exists, please call setup_ssh_directory() before adding keys"
+        )
     else:
         with open(authorized_keys_path, "r") as authorized_keys_file:
             content = authorized_keys_file.read()