From 53c00e3bc3b717e659c2c44fce1a559af234405a Mon Sep 17 00:00:00 2001 From: Satya Ortiz-Gagne Date: Tue, 27 Feb 2024 08:21:12 -0500 Subject: [PATCH] Add covalent aws ec2 launcher covalent is not compatible with milabench with requirement sqlalchemy<2.0.0 --- benchmarks/_template/requirements.cpu.txt | 1 + milabench/cli/covalent/__main__.py | 169 ++++++++++++++++++++++ milabench/cli/covalent/requirements.txt | 2 + 3 files changed, 172 insertions(+) create mode 100644 benchmarks/_template/requirements.cpu.txt create mode 100644 milabench/cli/covalent/__main__.py create mode 100644 milabench/cli/covalent/requirements.txt diff --git a/benchmarks/_template/requirements.cpu.txt b/benchmarks/_template/requirements.cpu.txt new file mode 100644 index 000000000..cd5826d4a --- /dev/null +++ b/benchmarks/_template/requirements.cpu.txt @@ -0,0 +1 @@ +voir>=0.2.9,<0.3 diff --git a/milabench/cli/covalent/__main__.py b/milabench/cli/covalent/__main__.py new file mode 100644 index 000000000..2ed921972 --- /dev/null +++ b/milabench/cli/covalent/__main__.py @@ -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() diff --git a/milabench/cli/covalent/requirements.txt b/milabench/cli/covalent/requirements.txt new file mode 100644 index 000000000..f810e6eaf --- /dev/null +++ b/milabench/cli/covalent/requirements.txt @@ -0,0 +1,2 @@ +covalent +covalent-ec2-plugin @ git+https://github.com/satyaog/covalent-ec2-plugin.git@feature/milabench \ No newline at end of file