Skip to content

Commit ff535a0

Browse files
justinchiu-coherejustinchiu-cohere
committed
pre-commit
1 parent 0b998b7 commit ff535a0

File tree

2 files changed

+279
-2
lines changed

2 files changed

+279
-2
lines changed

‎commit0/harness/execution_context.py

Lines changed: 218 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,218 @@
1+
""" Remote code execution contexts
2+
3+
Implements the interface for local docker containers, remote modal sandboxes,
4+
and HTTP servers.
5+
"""
6+
7+
8+
from abc import ABC
9+
import docker
10+
import logging
11+
import os
12+
import modal
13+
from pathlib import Path
14+
15+
from commit0.harness.spec import Spec
16+
from commit0.harness.docker_build import (
17+
close_logger,
18+
setup_logger,
19+
)
20+
from commit0.harness.docker_utils import (
21+
cleanup_container,
22+
create_container,
23+
copy_from_container,
24+
copy_to_container,
25+
copy_ssh_pubkey_from_container,
26+
delete_file_from_container,
27+
exec_run_with_timeout,
28+
)
29+
from commit0.harness.utils import (
30+
EvaluationError,
31+
extract_test_output,
32+
get_hash_string,
33+
get_ip,
34+
get_user,
35+
)
36+
37+
38+
class ExecutionContext(ABC):
39+
def __init__(
40+
self,
41+
spec: Spec,
42+
logger: logging.Logger,
43+
eval_file: Path,
44+
timeout: int,
45+
log_dir: Path,
46+
):
47+
""" Create the remote execution context
48+
49+
The execution context will persist for the lifetime of this object.
50+
The execution context can be a Docker container or Modal sandbox.
51+
"""
52+
raise NotImplementedError
53+
54+
def copy_ssh_pubkey_from_remote(self):
55+
raise NotImplementedError
56+
57+
def copy_to_remote(self, local_path, remote_path):
58+
raise NotImplementedError
59+
60+
def exec_run_with_timeout(self, command, timeout):
61+
raise NotImplementedError
62+
63+
def exec_run(self, command):
64+
raise NotImplementedError
65+
66+
def copy_from_remote(self, remote_path, local_path):
67+
raise NotImplementedError
68+
69+
def delete_file_from_remote(self, remote_path):
70+
raise NotImplementedError
71+
72+
def __enter__(self):
73+
raise NotImplementedError
74+
75+
def __exit__(self, exc_type, exc_value, exc_traceback):
76+
raise NotImplementedError
77+
78+
79+
class Docker(ExecutionContext):
80+
def __init__(
81+
self,
82+
spec: Spec,
83+
logger: logging.Logger,
84+
eval_file: Path,
85+
timeout: int,
86+
log_dir: Path,
87+
):
88+
client = docker.from_env()
89+
self.logger = logger
90+
self.container = create_container(
91+
client=client,
92+
image_name=spec.repo_image_key,
93+
container_name=spec.get_container_name(),
94+
logger=logger,
95+
)
96+
self.container.start()
97+
self.copy_ssh_pubkey_from_remote()
98+
99+
def copy_ssh_pubkey_from_remote(self) -> None:
100+
copy_ssh_pubkey_from_container(self.container)
101+
102+
def copy_to_remote(self, local_file: Path, remote_path: Path) -> None:
103+
copy_to_container(self.container, local_file, remote_path)
104+
105+
def exec_run_with_timeout(self, command: str, timeout: int) -> ():
106+
return exec_run_with_timeout(
107+
self.container, command, timeout
108+
)
109+
110+
def exec_run(self, command: str) -> None:
111+
return self.container.exec_run(command, demux=True)
112+
113+
def copy_from_remote(self, remote_path: Path, local_path: Path) -> None:
114+
copy_from_container(self.container, remote_path, local_path)
115+
116+
def delete_file_from_remote(self, remote_path: Path) -> None:
117+
delete_file_from_container(self.container, str(remote_path))
118+
119+
def __enter__(self):
120+
return self
121+
122+
def __exit__(self, exc_type, exc_value, exc_traceback):
123+
cleanup_container(self.client, self.container, self.logger)
124+
close_logger(self.logger)
125+
126+
127+
class Modal(ExecutionContext):
128+
def __init__(
129+
self,
130+
spec: Spec,
131+
logger: logging.Logger,
132+
eval_file: Path,
133+
timeout: int,
134+
log_dir: Path,
135+
):
136+
self.logger = logger
137+
# the image must exist on dockerhub
138+
reponame = spec.repo.split("/")[-1]
139+
image_name = f"wentingzhao/{reponame}"
140+
image = modal.Image.from_registry(image_name)
141+
142+
self.nfs = modal.NetworkFileSystem.ephemeral().__enter__()
143+
self.sandbox = modal.Sandbox.create(
144+
"sleep",
145+
"infinity",
146+
image=image,
147+
network_file_systems={
148+
"/vol": self.nfs,
149+
},
150+
)
151+
152+
self.copy_ssh_pubkey_from_remote()
153+
154+
def copy_ssh_pubkey_from_remote(self):
155+
process = self.sandbox.exec("bash", "-c", "cat /root/.ssh/id_rsa.pub")
156+
public_key = "".join([line for line in process.stdout]).strip()
157+
158+
# add to authorized keys locally. copy-pasted from utils
159+
local_authorized_keys_path = os.path.expanduser("~/.ssh/authorized_keys")
160+
os.makedirs(os.path.dirname(local_authorized_keys_path), exist_ok=True)
161+
if not os.path.exists(local_authorized_keys_path):
162+
# Since the file does not exist, create it
163+
open(local_authorized_keys_path, "a").close()
164+
write = True
165+
else:
166+
with open(local_authorized_keys_path, "r") as authorized_keys_file:
167+
content = authorized_keys_file.read()
168+
if public_key not in content:
169+
write = True
170+
else:
171+
write = False
172+
if write:
173+
with open(local_authorized_keys_path, "a") as authorized_keys_file:
174+
authorized_keys_file.write(public_key + "\n")
175+
176+
def copy_to_remote(self, local_path: Path, remote_path: Path) -> None:
177+
with local_path.open("rb") as f:
178+
self.nfs.write_file(str(local_path), f)
179+
self.sandbox.exec("bash", "-c", f"cp /vol/{str(local_path)} {str(remote_path)}")
180+
181+
def exec_run_with_timeout(self, command: str, timeout: int) -> None:
182+
"""Execute command on modal sandbox"""
183+
process = self.sandbox.exec("bash", "-c", command)
184+
stdout = []
185+
for line in process.stdout:
186+
stdout.append(line)
187+
stderr = []
188+
for line in process.stderr:
189+
stderr.append(line)
190+
return "\n".join(stdout), False, 1
191+
return "\n".join(stdout), "\n".join(stderr)
192+
193+
def exec_run(self, command: str) -> None:
194+
"""Execute command on modal sandbox"""
195+
process = self.sandbox.exec("bash", "-c", command)
196+
stdout = []
197+
for line in process.stdout:
198+
stdout.append(line)
199+
stderr = []
200+
for line in process.stderr:
201+
stderr.append(line)
202+
return 1, "\n".join(stdout)
203+
204+
def copy_from_remote(self, remote_path: Path, local_path: Path) -> None:
205+
"""Copy file from modal sandbox"""
206+
process = self.sandbox.exec("bash", "-c", f"cat {str(remote_path)}")
207+
output = "".join([line for line in process.stdout]).strip()
208+
with local_path.open("w") as f:
209+
f.write(output)
210+
211+
def delete_file_from_remote(src, remote_path):
212+
self.sandbox.exec("bash", "-c", f"rm {str(remote_path)}")
213+
214+
def __enter__(self):
215+
return self
216+
217+
def __exit__(self, exc_type, exc_value, exc_traceback):
218+
self.nfs.__exit__()

‎commit0/harness/run_pytest_ids.py

Lines changed: 61 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,11 @@
2929
get_ip,
3030
get_user,
3131
)
32+
from commit0.harness.execution_context import (
33+
Docker,
34+
Modal,
35+
)
36+
3237

3338

3439
class ExecutionBackend(StrEnum):
@@ -148,8 +153,9 @@ def run_modal(
148153

149154
# execute tests
150155
output, error = execute_command(sandbox, "/bin/bash /eval.sh")
151-
print(output)
152-
print(error)
156+
# TODO: add timeout
157+
#print(output)
158+
#print(error)
153159

154160
timed_out = False
155161
test_output = extract_test_output(
@@ -169,6 +175,15 @@ def run_modal(
169175
logger,
170176
)
171177

178+
# copy back report.json if there is any
179+
report_file = Path(spec.repo_directory) / "report.json"
180+
# Run the test command inside the container to check if the file exists
181+
exit_code, output = container.exec_run(f"test -e {report_file}", demux=True)
182+
# Check the exit code of the command
183+
if exit_code == 0:
184+
copy_from_sandbox(container, report_file, Path(log_dir / "report.json"))
185+
delete_file_from_sandbox(container, str(report_file))
186+
172187

173188
def main(
174189
dataset_name: str,
@@ -220,10 +235,54 @@ def main(
220235
eval_file = Path(log_dir / "eval.sh")
221236
eval_file.write_text(eval_script)
222237

238+
backend = "modal"
239+
execution_context = None
240+
if ExecutionBackend(backend) == ExecutionBackend.MODAL:
241+
execution_context = Modal
242+
elif ExecutionBackend(backend) == ExecutionBackend.LOCAL:
243+
execution_context = Docker
244+
print(backend, execution_context)
245+
246+
with execution_context(spec, logger, eval_file, timeout, log_dir) as context:
247+
context.copy_to_remote(eval_file, "/eval.sh")
248+
output, timed_out, total_runtime = context.exec_run_with_timeout(
249+
"/bin/bash /eval.sh", timeout
250+
)
251+
logger.info(output)
252+
test_output = extract_test_output(
253+
output, "--json-report --json-report-file=report.json"
254+
)
255+
256+
# stdout might be more straightforward
257+
print(test_output)
258+
test_output_path = log_dir / "test_output.txt"
259+
with open(test_output_path, "w") as f:
260+
f.write(test_output)
261+
if timed_out:
262+
f.write(f"\n\nTimeout error: {timeout} seconds exceeded.")
263+
raise EvaluationError(
264+
spec.repo,
265+
f"Test timed out after {timeout} seconds.",
266+
logger,
267+
)
268+
269+
# copy back report.json if there is any
270+
report_file = Path(spec.repo_directory) / "report.json"
271+
# Run the test command inside the container to check if the file exists
272+
#exit_code, output = container.exec_run(f"test -e {report_file}", demux=True)
273+
exit_code, output = context.exec_run(f"test -e {report_file}")
274+
# Check the exit code of the command
275+
if exit_code == 0:
276+
context.copy_from_remote(report_file, Path(log_dir / "report.json"))
277+
context.delete_file_from_remote(report_file)
278+
279+
280+
"""
223281
if ExecutionBackend(backend) == ExecutionBackend.LOCAL:
224282
run_docker(spec, logger, eval_file, timeout, log_dir)
225283
elif ExecutionBackend(backend) == ExecutionBackend.MODAL:
226284
run_modal(spec, logger, eval_file, timeout, log_dir)
285+
"""
227286

228287

229288
__all__ = []

0 commit comments

Comments
 (0)