Skip to content

Commit

Permalink
Add covalent aws ec2 launcher
Browse files Browse the repository at this point in the history
covalent is not compatible with milabench with requirement sqlalchemy<2.0.0
  • Loading branch information
satyaog committed Feb 27, 2024
1 parent 0bf6348 commit 53c00e3
Show file tree
Hide file tree
Showing 3 changed files with 172 additions and 0 deletions.
1 change: 1 addition & 0 deletions benchmarks/_template/requirements.cpu.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
voir>=0.2.9,<0.3
169 changes: 169 additions & 0 deletions milabench/cli/covalent/__main__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,169 @@
import asyncio
import os
import pathlib
import subprocess
import sys
import tempfile


def main(argv=None):
if argv is None:
argv = sys.argv[1:]

try:
import covalent as ct
from covalent._shared_files.util_classes import RESULT_STATUS
except ImportError:
cache_dir = pathlib.Path("/tmp/milabench/covalent_venv")
python3 = str(cache_dir / "bin/python3")
try:
subprocess.run([
python3,
"-c",
"import covalent ; from covalent.executor import EC2Executor"
], check=True)
except (FileNotFoundError, subprocess.CalledProcessError):
cache_dir.mkdir(parents=True, exist_ok=True)
subprocess.run([sys.executable, "-m", "virtualenv", str(cache_dir)], check=True)
subprocess.run([python3, "-m", "pip", "install", "-U", "pip"], check=True)
subprocess.run([
python3,
"-m",
"pip",
"install",
"-r",
str(pathlib.Path(__file__).resolve().parent / "requirements.txt")
], check=True)
subprocess.run([
python3,
"-c",
"import covalent ; from covalent.executor import EC2Executor"
], check=True)
exit(subprocess.call(
[python3, __file__, *argv],
))

def _popen(cmd, *args, _env=None, **kwargs):
_env = _env if _env is not None else {}

for envvar in _env.keys():
envvar_val = _env[envvar]

if not envvar_val:
continue

envvar_val = pathlib.Path(envvar_val).expanduser()
if str(envvar_val) != _env[envvar]:
_env[envvar] = str(envvar_val)

if "MILABENCH_CONFIG_CONTENT" in _env:
_config_dir = pathlib.Path(_env["MILABENCH_CONFIG"]).parent
with tempfile.NamedTemporaryFile("wt", dir=str(_config_dir), suffix=".yaml", delete=False) as _f:
_f.write(_env["MILABENCH_CONFIG_CONTENT"])
_env["MILABENCH_CONFIG"] = _f.name

try:
cmd = (str(pathlib.Path(cmd[0]).expanduser()), *cmd[1:])
except IndexError:
pass

cwd = kwargs.pop("cwd", None)
if cwd is not None:
cwd = str(pathlib.Path(cwd).expanduser())

_env = {**os.environ.copy(), **kwargs.pop("env", {}), **_env}

kwargs = {
**kwargs,
"cwd": cwd,
"env": _env,
"stdout": subprocess.PIPE,
"stderr": subprocess.PIPE,
}
p = subprocess.Popen(cmd, *args, **kwargs)

stdout_chunks = []
while True:
line = p.stdout.readline()
if not line:
break
line_str = line.decode("utf-8").strip()
stdout_chunks.append(line_str)
print(line_str)

_, stderr = p.communicate()
stderr = stderr.decode("utf-8").strip()
stdout = os.linesep.join(stdout_chunks)

if p.returncode != 0:
raise subprocess.CalledProcessError(
p.returncode,
(cmd, args, kwargs),
stdout,
stderr
)
return p.returncode, os.linesep.join([stdout, stderr])

executor:ct.executor.BaseExecutor = ct.executor.EC2Executor(
profile="mb_test_sog_1",
username="ubuntu",
instance_type="t2.micro",
volume_size="8",
region="us-east-2",
python_path="python3",
)

@ct.lattice
def lattice(argv=()):
venv = pathlib.Path("~/venv")
code = pathlib.Path("~/milabench")
_env = {
"MILABENCH_BASE": "~/benches",
}
env_config = os.environ.get("MILABENCH_CONFIG", None)

if env_config is not None:
env_config = pathlib.Path(env_config)
_env["MILABENCH_CONFIG"] = f"{code / (env_config.relative_to(env_config.parent.parent))}"
_env["MILABENCH_CONFIG_CONTENT"] = env_config.read_text()

return ct.electron(
_popen,
executor=executor,
deps_bash=ct.DepsBash([
f"[ -d {code} ] || git clone https://github.com/mila-iqia/milabench.git {code}",
f"git -C {code} checkout -B stable origin/stable",
f"python3 -m virtualenv {venv}",
f"{venv}/bin/python3 -m pip install -U pip",
f"{venv}/bin/python3 -m pip install -U -e {code}",
]),
)(
[f"{venv}/bin/python3", "-m", "milabench", *argv],
cwd=str(code),
_env=_env,
)

# lattice(argv)
dispatch_id = ct.dispatch(lattice, disable_run=False)(argv)
try:
result = ct.get_result(dispatch_id=dispatch_id, wait=True)
finally:
result = ct.get_result(dispatch_id=dispatch_id, wait=False)
if result.get_node_result(0)["status"] not in (
RESULT_STATUS.CANCELLED,
RESULT_STATUS.COMPLETED,
RESULT_STATUS.FAILED,
):
for coro in (
# This should not start a new instance but only retreive the
# necessary attributes to shutdown the running instance
executor.setup({'dispatch_id': dispatch_id, 'node_id': 0, 'results_dir': result.results_dir}),
executor.teardown({'dispatch_id': dispatch_id, 'node_id': 0, 'results_dir': result.results_dir}),
):
asyncio.run(coro)

sys.exit(result.result[0] if result.result is not None else 1)


if __name__ == "__main__":
main()
2 changes: 2 additions & 0 deletions milabench/cli/covalent/requirements.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
covalent
covalent-ec2-plugin @ git+https://github.com/satyaog/covalent-ec2-plugin.git@feature/milabench

0 comments on commit 53c00e3

Please sign in to comment.