Merge pull request #26 from commit-0/ssh-modal

make ssh also work for modal

Author: wenting-zhao

.github/workflows/system.yml

@@ -19,9 +19,11 @@ jobs:
       - name: Install the project
         run: uv sync
       - name: Set up commit0
-        run: uv run commit0 setup simpy
+        run: uv run commit0 clone simpy
       - name: Build docker images
         run: uv run commit0 build simpy
+      - name: Set up git user
+        run: sudo "$(which uv)" run commit0 setup-git-user simpy
       - name: Get tests
         run: uv run commit0 get-tests simpy
       - name: Test

commit0/__main__.py

@@ -2,6 +2,7 @@
 import commit0.harness.get_pytest_ids
 import commit0.harness.build
 import commit0.harness.setup
+import commit0.harness.setup_git_user
 import commit0.harness.evaluate
 import commit0.harness.save
 import copy
@@ -30,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 ["setup", "build", "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(
@@ -39,7 +47,7 @@ def main() -> None:
             config.repo_split = sys.argv[2]
     config.base_dir = os.path.abspath(config.base_dir)
 
-    if command == "setup":
+    if command == "clone":
         commit0.harness.setup.main(
             config.dataset_name,
             config.dataset_split,
@@ -53,6 +61,17 @@ def main() -> None:
             config.dataset_split,
             config.repo_split,
             config.num_workers,
+            config.backend,
+            config.key_path,
+        )
+    elif command == "setup-git-user":
+        commit0.harness.setup_git_user.main(
+            config.dataset_name,
+            config.dataset_split,
+            config.repo_split,
+            config.base_dir,
+            config.git_user,
+            config.key_path,
         )
     elif command == "get-tests":
         repo = sys.argv[2]

commit0/configs/base.yaml

@@ -11,6 +11,10 @@ repo_split: all
 
 # build related
 num_workers: 8
+key_path: commit0/configs/public_keys.json
+
+# set up git user
+git_user: git  # by default, git user is called git
 
 # test related
 backend: local

commit0/configs/config_class.py

@@ -15,6 +15,8 @@ class Commit0Config:
     # build related
     # which repo to build, all or one repo
     num_workers: int
+    # path to store public keys from docker images
+    key_path: str
 
     # test related
     backend: str

commit0/harness/build.py

@@ -1,12 +1,20 @@
+import json
 import logging
+import traceback
 
 import docker
 from datasets import load_dataset
+from tqdm import tqdm
 from typing import Iterator
 
+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,
+    Docker,
+    Modal,
+)
 from commit0.harness.spec import make_spec
-from commit0.harness.constants import RepoInstance, SPLIT
 
 logging.basicConfig(
     level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
@@ -15,7 +23,12 @@
 
 
 def main(
-    dataset_name: str, dataset_split: str, repo_split: str, num_workers: int
+    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 = []
@@ -26,8 +39,29 @@ def main(
         spec = make_spec(example)
         specs.append(spec)
 
-    client = docker.from_env()
-    build_repo_images(client, specs, num_workers)
+    if ExecutionBackend(backend) == ExecutionBackend.MODAL:
+        execution_context = Modal
+    elif ExecutionBackend(backend) == ExecutionBackend.LOCAL:
+        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()
+    for spec in tqdm(specs, desc="Retrieving public keys from docker images"):
+        try:
+            with execution_context(spec, logger, timeout=60) as context:
+                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"
+            raise RuntimeError(error_msg)
+    with open(key_path, "w") as json_file:
+        json.dump(img2key, json_file, indent=4)
 
 
 __all__ = []

commit0/harness/constants.py

@@ -11,6 +11,10 @@ class RepoInstance(TypedDict):
     test: Dict[str, str]
 
 
+class Files(TypedDict):
+    eval_script: Dict[str, Path]
+
+
 # Constants - Evaluation Log Directories
 BASE_IMAGE_BUILD_DIR = Path("logs/build_images/base")
 REPO_IMAGE_BUILD_DIR = Path("logs/build_images/repo")
@@ -27,8 +31,9 @@ class RepoInstance(TypedDict):
 
 # available commands
 COMMANDS = [
-    "setup",
+    "clone",
     "build",
+    "setup-git-user",
     "test",
     "test-reference",
     "get-tests",

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
@@ -141,50 +140,31 @@ def delete_file_from_container(container: Container, file_path: str) -> None:
         raise Exception(f"General Error: {str(e)}")
 
 
-def copy_ssh_pubkey_from_container(container: Container) -> None:
+def get_ssh_pubkey_from_container(container: Container, user: str) -> str:
     """Copy the SSH public key from a Docker container to the local authorized_keys file.
 
     Args:
     ----
     container (Container): Docker container to copy the key from.
+    user (str): to get public key of which user
+
+    Returns:
+    -------
+    public_key (str): public key from docker container
 
     Raises:
     ------
     docker.errors.APIError: If there is an error calling the Docker API.
-    Exception: If the file reading or writing process fails.
 
     """
     try:
         exit_code, output = container.exec_run("cat /root/.ssh/id_rsa.pub")
         if exit_code != 0:
             raise Exception(f"Error reading file: {output.decode('utf-8').strip()}")
         public_key = output.decode("utf-8").strip()
-        public_key = f"no-port-forwarding,no-X11-forwarding,no-agent-forwarding,no-pty {public_key}"
-
-        user_info = pwd.getpwnam("git")
-        home_directory = user_info.pw_dir
-        authorized_keys_path = os.path.join(home_directory, ".ssh", "authorized_keys")
-        os.makedirs(os.path.dirname(authorized_keys_path), exist_ok=True)
-        if not os.path.exists(authorized_keys_path):
-            # Since the file does not exist, create it
-            open(authorized_keys_path, "a").close()
-            write = True
-        else:
-            with open(authorized_keys_path, "r") as authorized_keys_file:
-                content = authorized_keys_file.read()
-                if public_key not in content:
-                    write = True
-                else:
-                    write = False
-
-        if write:
-            with open(authorized_keys_path, "a") as authorized_keys_file:
-                authorized_keys_file.write(public_key + "\n")
-
+        return public_key
     except docker.errors.APIError as e:
         raise docker.errors.APIError(f"Docker API Error: {str(e)}")
-    except Exception as e:
-        raise Exception(f"General Error: {str(e)}")
 
 
 def write_to_container(container: Container, data: str, dst: Path) -> None: