Merge pull request #29 from commit-0/modal-fix

Get modal working + simplify ExecutionContext class

Author: wenting-zhao

commit0/__main__.py

@@ -78,6 +78,7 @@ def main() -> None:
             test_ids,
             config.backend,
             config.timeout,
+            config.num_cpus,
             stdout=True,
         )
     elif command == "evaluate" or command == "evaluate-reference":
@@ -91,6 +92,7 @@ def main() -> None:
             config.branch,
             config.backend,
             config.timeout,
+            config.num_cpus,
             config.num_workers,
         )
     elif command == "save":

commit0/configs/base.yaml

@@ -16,6 +16,7 @@ num_workers: 8
 backend: local
 branch: ai
 timeout: 1_800
+num_cpus: 1
 
 # save related
 github_token: null

commit0/configs/config_class.py

@@ -22,6 +22,7 @@ class Commit0Config:
     branch: str
     # timeout for running pytest
     timeout: int
+    num_cpus: int
 
     # save related
     github_token: Optional[str]

commit0/harness/docker_utils.py

@@ -116,30 +116,6 @@ def safe_extract(
         extracted_file_path.rename(dst)
 
 
-def delete_file_from_container(container: Container, file_path: str) -> None:
-    """Delete a file from a docker container.
-
-    Args:
-    ----
-        container (Container): Docker container to delete the file from
-        file_path (str): Path to the file in the container to be deleted
-
-    Raises:
-    ------
-        docker.errors.APIError: If there is an error calling the Docker API.
-        Exception: If the file deletion command fails with a non-zero exit code.
-
-    """
-    try:
-        exit_code, output = container.exec_run(f"rm -f {file_path}")
-        if exit_code != 0:
-            raise Exception(f"Error deleting file: {output.decode('utf-8').strip()}")
-    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:
     """Write a string to a file in a docker container"""
     # echo with heredoc to file
@@ -247,7 +223,7 @@ def create_container(
     image_name: str,
     container_name: Optional[str] = None,
     user: Optional[str] = None,
-    command: Optional[str] = None,
+    command: Optional[str] = "tail -f /dev/null",
     nano_cpus: Optional[int] = None,
     logger: Optional[Union[str, logging.Logger]] = None,
 ) -> Container:
@@ -312,7 +288,7 @@ def log_info(x: str) -> None:
             image=image_name,
             name=container_name,
             user=user,
-            command="tail -f /dev/null",
+            command=command,
             nano_cpus=nano_cpus,
             detach=True,
         )

commit0/harness/evaluate.py

@@ -26,6 +26,7 @@ def main(
     branch: str,
     backend: str,
     timeout: int,
+    num_cpus: int,
     num_workers: int,
 ) -> None:
     dataset: Iterator[RepoInstance] = load_dataset(dataset_name, split=dataset_split)  # type: ignore
@@ -52,6 +53,7 @@ def main(
                     test_dir,
                     backend,
                     timeout,
+                    num_cpus,
                     stdout=False,
                 ): None
                 for repo, test_dir in pairs

commit0/harness/execution_context.py

@@ -11,6 +11,7 @@
 import modal.io_streams
 from enum import StrEnum, auto
 from pathlib import Path
+import time
 from typing import Optional, Type
 from types import TracebackType
 
@@ -27,7 +28,6 @@
     create_container,
     copy_from_container,
     copy_to_container,
-    delete_file_from_container,
     exec_run_with_timeout,
 )
 
@@ -51,43 +51,29 @@ def __init__(
         spec: Spec,
         logger: logging.Logger,
         timeout: int,
+        num_cpus: int,
+        log_dir: Path,
+        files_to_copy: Optional[Files] = None,
     ):
         """Create the remote execution context
 
-        The execution context will persist for the lifetime of this object.
         The execution context can be a Docker container or Modal sandbox.
+        The execution context may not persist for the lifetime of this object.
         """
         self.spec = spec
         self.logger = logger
         self.timeout = timeout
+        self.num_cpus = num_cpus
+        self.log_dir = log_dir
 
     @abstractmethod
-    def exec_run_with_timeout(
-        self, command: str, timeout: int
-    ) -> tuple[str, bool, float]:
-        """Exec"""
-        raise NotImplementedError
-
-    @abstractmethod
-    def exec_run(self, command: str) -> tuple[int, str]:
-        """Exec"""
-        raise NotImplementedError
-
-    @abstractmethod
-    def copy_from_remote(self, remote_path: Path, local_path: Path) -> None:
-        """Copy"""
-        raise NotImplementedError
-
-    @abstractmethod
-    def delete_file_from_remote(self, remote_path: Path) -> None:
-        """Delete"""
+    def exec_run_with_timeout(self, command: str) -> tuple[str, bool, float]:
+        """Execute a test command"""
         raise NotImplementedError
 
-    def write_test_output(
-        self, log_dir: Path, test_output: str, timed_out: bool
-    ) -> None:
+    def write_test_output(self, test_output: str, timed_out: bool) -> None:
         """Write test output"""
-        test_output_path = log_dir / "test_output.txt"
+        test_output_path = self.log_dir / "test_output.txt"
         with open(test_output_path, "w") as f:
             f.write(test_output)
             if timed_out:
@@ -98,15 +84,6 @@ def write_test_output(
                     self.logger,
                 )
 
-        # copy back report.json if there is any
-        report_file = Path(self.spec.repo_directory) / "report.json"
-        # Run the test command inside the container to check if the file exists
-        exit_code, output = self.exec_run(f"test -e {report_file}")
-        # Check the exit code of the command
-        if exit_code == 0:
-            self.copy_from_remote(report_file, log_dir / "report.json")
-            self.delete_file_from_remote(report_file)
-
     def __enter__(self):
         return self
 
@@ -126,39 +103,41 @@ def __init__(
         spec: Spec,
         logger: logging.Logger,
         timeout: int,
+        num_cpus: int,
+        log_dir: Path,
         files_to_copy: Optional[Files] = None,
     ):
-        super().__init__(spec, logger, timeout)
+        super().__init__(spec, logger, timeout, num_cpus, log_dir)
 
         self.client = docker.from_env()
         self.container = create_container(
             client=self.client,
             image_name=spec.repo_image_key,
             container_name=spec.get_container_name(),
+            nano_cpus=num_cpus,
             logger=logger,
         )
         self.container.start()
         if files_to_copy:
             for _, f in files_to_copy.items():
                 copy_to_container(self.container, f["src"], f["dest"])  # type: ignore
 
-    def exec_run_with_timeout(
-        self, command: str, timeout: int
-    ) -> tuple[str, bool, float]:
-        """Exec"""
-        return exec_run_with_timeout(self.container, command, timeout)
-
-    def exec_run(self, command: str) -> tuple[int, str]:
+    def exec_run_with_timeout(self, command: str) -> tuple[str, bool, float]:
         """Exec"""
-        return self.container.exec_run(command, demux=True)
-
-    def copy_from_remote(self, remote_path: Path, local_path: Path) -> None:
-        """Copy"""
-        copy_from_container(self.container, remote_path, local_path)
+        output = exec_run_with_timeout(self.container, command, self.timeout)
 
-    def delete_file_from_remote(self, remote_path: Path) -> None:
-        """Delete"""
-        delete_file_from_container(self.container, str(remote_path))
+        # copy back report.json if there is any
+        report_file = Path(self.spec.repo_directory) / "report.json"
+        # Run the test command inside the container to check if the file exists
+        exit_code, test_output = self.container.exec_run(
+            f"test -e {report_file}", demux=True
+        )
+        # Check the exit code of the command
+        if exit_code == 0:
+            copy_from_container(
+                self.container, report_file, self.log_dir / "report.json"
+            )
+        return output
 
     def __exit__(
         self,
@@ -176,64 +155,67 @@ def __init__(
         spec: Spec,
         logger: logging.Logger,
         timeout: int,
+        num_cpus: int,
+        log_dir: Path,
         files_to_copy: Optional[Files] = None,
     ):
-        super().__init__(spec, logger, timeout)
+        super().__init__(spec, logger, timeout, num_cpus, log_dir)
+
+        self.app = modal.App()
 
         # the image must exist on dockerhub
         reponame = spec.repo.split("/")[-1]
-        image_name = f"wentingzhao/{reponame}"
+        image_name = f"wentingzhao/{reponame}:latest"
         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"])  # type: ignore
+        self.image = image
 
-        self.sandbox = modal.Sandbox.create(
-            "sleep",
-            "infinity",
-            image=image,
-            cpu=4.0,
-            timeout=timeout,
-        )
-
-    def exec_run_with_timeout(
-        self, command: str, timeout: int
-    ) -> tuple[str, bool, float]:
-        """Execute command on modal sandbox"""
-        print("Executing:", command)
-        process = self.sandbox.exec("bash", "-c", command)
-        print("stdout")
-        stdout = read_stream(process.stdout)
-        print("stderr")
-        stderr = read_stream(process.stderr)
-        print(stderr)
-        return stdout, False, 1.0
-        return stdout, stderr
-
-    def exec_run(self, command: str) -> tuple[int, str]:
+    def exec_run_with_timeout(self, command: str) -> tuple[str, bool, float]:
         """Execute command on modal sandbox"""
-        process = self.sandbox.exec("bash", "-c", command)
-        stdout = read_stream(process.stdout)
-        stderr = read_stream(process.stderr)
-        print(stderr)
-        return 1, stdout
-
-    def copy_from_remote(self, remote_path: Path, local_path: Path) -> None:
-        """Copy file from modal sandbox"""
-        process = self.sandbox.exec("bash", "-c", f"cat {str(remote_path)}")
-        output = "".join([line for line in process.stdout]).strip()
-        with local_path.open("w") as f:
-            f.write(output)
-
-    def delete_file_from_remote(self, remote_path: Path) -> None:
-        """Delete"""
-        self.sandbox.exec("bash", "-c", f"rm {str(remote_path)}")
+        start_time = time.time()
+        with modal.Volume.ephemeral() as vol:
+            # copy back report.json if there is any
+            report_file = Path(self.spec.repo_directory) / "report.json"
+
+            self.sandbox = modal.Sandbox.create(
+                "bash",
+                "-c",
+                f"{command} && cp {str(report_file)} /vol/report.json",
+                image=self.image,
+                cpu=self.num_cpus,
+                timeout=self.timeout,
+                app=self.app,
+                volumes={"/vol": vol},
+            )
+            self.sandbox.wait()
+
+            # stdout has been redirected to stderr
+            stdout = read_stream(self.sandbox.stderr)
+
+            return_code = self.sandbox.returncode
+            # https://github.com/modal-labs/modal-client/blob/d577b2916b5c3bf4ebbcb58fadced84d85e1cf8c/modal/sandbox.py#L413
+            if return_code == 124:
+                timed_out = True
+            else:
+                timed_out = False
+
+            # copy over report.json from mount
+            with (self.log_dir / "report.json").open("wb") as f:
+                for data in vol.read_file("report.json"):
+                    f.write(data)
+
+            self.sandbox.terminate()
+
+            end_time = time.time()
+
+            return stdout, timed_out, end_time - start_time
 
     def __exit__(
         self,
         exctype: Optional[Type[BaseException]],
         excinst: Optional[BaseException],
         exctb: Optional[TracebackType],
     ) -> None:
-        self.sandbox.terminate()
         close_logger(self.logger)

commit0/harness/run_pytest_ids.py

@@ -36,6 +36,7 @@ def main(
     test_ids: str,
     backend: str,
     timeout: int,
+    num_cpus: int,
     stdout: bool,
 ) -> str:
     """Runs the pytests for repos in a dataset.
@@ -92,16 +93,19 @@ def main(
     )
 
     try:
-        with execution_context(spec, logger, timeout, files_to_copy) as context:
+        with execution_context(
+            spec, logger, timeout, num_cpus, log_dir, files_to_copy
+        ) as context:
             output, timed_out, total_runtime = context.exec_run_with_timeout(
-                "/bin/bash /eval.sh", timeout
+                "/bin/bash /eval.sh"
             )
             logger.info(output)
             test_output = extract_test_output(
                 output, "--json-report --json-report-file=report.json"
             )
-            context.write_test_output(log_dir, test_output, timed_out)
-            print(test_output)
+            context.write_test_output(test_output, timed_out)
+            if stdout:
+                print(test_output)
     except EvaluationError as e:
         error_msg = (
             f"Error in running pytest for {repo}: {e}\n"

commit0/harness/spec.py

@@ -157,6 +157,8 @@ def make_eval_script_list(instance: RepoInstance, repo_directory: str) -> list[s
         f"git reset --hard {instance['base_commit']}",
         "git status",
     ]
+    for i in range(len(eval_script_list)):
+        eval_script_list[i] = f"{eval_script_list[i]} 1>&2"
     return eval_script_list