diff --git a/caliban/cli.py b/caliban/cli.py index 86cfdf3..0a28191 100644 --- a/caliban/cli.py +++ b/caliban/cli.py @@ -31,6 +31,7 @@ import caliban.platform.gke.constants as gke_k import caliban.platform.gke.types as gke_t import caliban.platform.gke.util as gke_u +import caliban.platform.slurm as slurm import caliban.util as u import caliban.util.argparse as ua import caliban.util.schema as us @@ -413,6 +414,7 @@ def caliban_parser(): local_run_parser(subparser) cloud_parser(subparser) cluster_parser(subparser) + slurm_parser(subparser) status_parser(subparser) stop_parser(subparser) resubmit_parser(subparser) @@ -946,3 +948,60 @@ def max_jobs_arg(parser): f'then this specifies the total number of jobs to return, ordered ' f'by creation date, or all jobs if max_jobs==0.'), ) + + +# ---------------------------------------------------------------------------- +def slurm_parser(base): + """cli parser for slurm commands""" + + parser = base.add_parser("slurm", + description="slurm commands", + help="slurm-related commands") + + subparser = parser.add_subparsers(dest="slurm_cmd") + slurm_ls_cmd(subparser) + slurm_job_parser(subparser) + + +# ---------------------------------------------------------------------------- +def slurm_ls_cmd(base): + """caliban slurm ls""" + + parser = base.add_parser( + "ls", + description="list partitions", + help="list partitions", + formatter_class=argparse.ArgumentDefaultsHelpFormatter) + + +# ---------------------------------------------------------------------------- +def slurm_job_parser(base): + parser = base.add_parser( + "job", + description="job commands", + help="job commands", + formatter_class=argparse.ArgumentDefaultsHelpFormatter) + + subparser = parser.add_subparsers(dest="job_cmd") + slurm_job_ls_cmd(subparser) + cluster_job_submit_cmd(subparser) + + +# ---------------------------------------------------------------------------- +def slurm_job_ls_cmd(base): + parser = base.add_parser( + "ls", + description="list Slurm jobs", + help="list Slurm jobs", + formatter_class=argparse.ArgumentDefaultsHelpFormatter) + + +# ---------------------------------------------------------------------------- +def slurm_job_submit_cmd(base): + parser = base.add_parser( + "submit", + description="submit Slurm job(s)", + help="submit Slurm job(s)", + formatter_class=argparse.ArgumentDefaultsHelpFormatter) + + job_name_arg(parser) diff --git a/caliban/docker/build.py b/caliban/docker/build.py index a8ac0f2..09baf4a 100644 --- a/caliban/docker/build.py +++ b/caliban/docker/build.py @@ -37,7 +37,8 @@ t = Terminal() -DEV_CONTAINER_ROOT = "gcr.io/blueshift-playground/blueshift" +# DEV_CONTAINER_ROOT = "gcr.io/blueshift-playground/blueshift" +DEV_CONTAINER_ROOT = "docker.io/eschnett/carpetx-caliban" DEFAULT_GPU_TAG = "gpu-ubuntu1804-py37-cuda101" DEFAULT_CPU_TAG = "cpu-ubuntu1804-py37" TF_VERSIONS = {"2.2.0", "1.12.3", "1.14.0", "1.15.0"} diff --git a/caliban/main.py b/caliban/main.py index 3dada18..7607463 100644 --- a/caliban/main.py +++ b/caliban/main.py @@ -34,6 +34,8 @@ import caliban.platform.notebook as pn import caliban.platform.run as pr import caliban.platform.shell as ps +import caliban.platform.slurm as slurm +import caliban.platform.slurm.cli import caliban.util.schema as cs ll.getLogger('caliban.main').setLevel(logging.ERROR) @@ -53,6 +55,9 @@ def run_app(arg_input): if command == "cluster": return gke.cli.run_cli_command(args) + if command == "slurm": + return slurm.cli.run_cli_command(args) + job_mode = cli.resolve_job_mode(args) docker_args = cli.generate_docker_args(job_mode, args) docker_run_args = args.get("docker_run_args", []) diff --git a/caliban/platform/slurm/__init__.py b/caliban/platform/slurm/__init__.py new file mode 100644 index 0000000..79c6a2f --- /dev/null +++ b/caliban/platform/slurm/__init__.py @@ -0,0 +1,15 @@ +#!/usr/bin/python +# +# Copyright 2020 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/caliban/platform/slurm/cli.py b/caliban/platform/slurm/cli.py new file mode 100644 index 0000000..31f5222 --- /dev/null +++ b/caliban/platform/slurm/cli.py @@ -0,0 +1,260 @@ +#!/usr/bin/python +# +# Copyright 2020 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""slurm cli support""" + +import json +import logging +import os +import pprint as pp +import re +import shlex +import subprocess +from blessings import Terminal +from datetime import datetime +from pprint import pformat +from typing import Any, Dict, List, Optional + +import caliban.cli as cli +import caliban.config as conf +import caliban.docker.build as db +import caliban.util as u +from caliban.history.util import (create_experiments, generate_container_spec, + get_mem_engine, get_sql_engine, session_scope) +from caliban.platform.cloud.core import generate_image_tag + +t = Terminal() + +# ---------------------------------------------------------------------------- +# Remote access via ssh +_ssh = "ssh" +_hostname = "symmetry.pi.local" +_username = "eschnetter" +_logdir = "/home/eschnetter/caliban" +_resultdir_host = "/gpfs/eschnetter/caliban-simulations" +_resultdir_container = "/caliban-simulations" + +# Environment modules +_setup_cmds = [["source", "/etc/profile"], ["module", "load", "slurm"]] + +# Job management via Slurm +_sbatch = "sbatch" +_sinfo = "sinfo" +_squeue = "squeue" +_partition = "debugq" +_timelimit = "00:10:00" +_nodes = 1 +_ntasks_per_node = 40 +_nthreads_per_task = 1 + + +# ---------------------------------------------------------------------------- +def run_cli_command(args: dict) -> None: + """cli entrypoint for Slurm commands""" + SLURM_CMDS = { + 'ls': _partition_ls, + 'job': _job_commands, + } + SLURM_CMDS[args['slurm_cmd']](args) + + +# ---------------------------------------------------------------------------- +def _partition_ls(args: dict) -> None: + """list Slurm partitions""" + cmd = [_sinfo, '-h', '-o', '%P'] + process = subprocess.run(_with_ssh(cmd), + stdout=subprocess.PIPE, + universal_newlines=True) + partitions = process.stdout.splitlines() + logging.info("{} partitions found".format(len(partitions))) + for p in partitions: + logging.info(" " + p) + + +# ---------------------------------------------------------------------------- +def _job_commands(args: dict) -> None: + """job commands""" + JOB_CMDS = { + 'ls': _job_ls, + 'submit': _job_submit, + #TODO 'submit_file': _job_submit_file + } + JOB_CMDS[args['job_cmd']](args) + + +# ---------------------------------------------------------------------------- +def _job_ls(args: dict) -> None: + """list Slurm jobs""" + # TODO: Look also for jobs that have terminated a long time ago? + cmd = [_squeue, '-h', '-o', "%A", '-t', "all", '-u', _username] + process = subprocess.run(_with_ssh(cmd), + stdout=subprocess.PIPE, + universal_newlines=True) + jobs = process.stdout.splitlines() + logging.info("{} jobs found".format(len(jobs))) + for j in jobs: + # TODO: Look up and output job metadata + logging.info(" " + j) + + +# ---------------------------------------------------------------------------- +def _job_submit(args: dict) -> None: + """submits Slurm job(s) + + Args: + args: argument dictionary + """ + + script_args = conf.extract_script_args(args) + job_mode = cli.resolve_job_mode(args) + docker_args = cli.generate_docker_args(job_mode, args) + docker_run_args = args.get('docker_run_args', []) or [] + dry_run = args['dry_run'] + package = args['module'] + job_name = args.get('name') + gpu_spec = args.get('gpu_spec') + preemptible = not args['nonpreemptible'] + min_cpu = args.get('min_cpu') + min_mem = args.get('min_mem') + experiment_config = args.get('experiment_config') or [{}] + xgroup = args.get('xgroup') + image_tag = args.get('image_tag') + export = args.get('export', None) + + labels = args.get('label') + if labels is not None: + labels = dict(cu.sanitize_labels(args.get('label'))) + + # Generate job name + dt = datetime.now().astimezone() + job_prefix = f'caliban-{dt.strftime("%Y%m%d-%H%M%S")}' + # TODO: Sanitize job_name + if job_name is None: + job_name = job_prefix + log_filename = job_name + ".log" + else: + log_filename = job_prefix + "-" + job_name + ".log" + log_filename = _logdir + "/" + log_filename + + # Arguments to internally build the image required to submit to Slurm + docker_m = {'job_mode': job_mode, 'package': package, **docker_args} + + # -------------------------------------------------------------------------- + engine = get_mem_engine() if dry_run else get_sql_engine() + + logging.info("*** engine") + with session_scope(engine) as session: + container_spec = generate_container_spec(session, docker_m, image_tag) + + logging.info("*** image_tag") + if image_tag is None: + # Create Docker image, and push it to a Docker repository + logging.info("Generating Docker image with parameters:") + logging.info(t.yellow(pformat(docker_args))) + + if dry_run: + logging.info( + "Dry run - skipping actual 'docker build' and 'docker push'.") + image_tag = "dry_run_tag" + else: + image_id = db.build_image(**docker_m) + + project_id = args['project_id'] + assert project_id is not None + project_s = project_id.replace(":", "/") + # TODO: Use sub-project "caliban" or similar? + # base = f"docker.io/{project_s}/{image_id}" + base = f"{project_s}/{image_id}" + image_tag = f"{base}:latest" + subprocess.run(["docker", "tag", image_id, image_tag], check=True) + subprocess.run(["docker", "push", image_tag], check=True) + + logging.info("*** experiments") + experiments = create_experiments( + session=session, + container_spec=container_spec, + script_args=script_args, + experiment_config=experiment_config, + xgroup=xgroup, + ) + + for experiment in experiments: + logging.info("*** script") + # TODO: Use srun to start job, then use singularity inside srun + qargs = _quote_args(script_args) + script = f"""\ +#!/bin/bash +source /etc/profile +set -euxo pipefail +date +hostname +nproc +module load singularity +env PYTHONUNBUFFERED=1 singularity run --tmpdir /gpfs/eschnetter/singularity/tmp --no-home --bind {_resultdir_host}:{_resultdir_container} --pwd /usr/app docker://{image_tag} {qargs} +""" + logging.info(f"Script is {script}") + + logging.info("Submitting job...") + prepare_cmd = ["mkdir", "-p", _logdir, _resultdir_host] + env_cmd = [ + "export", "CALIBAN_NPROCS=" + str(_nodes * _ntasks_per_node), + "CALIBAN_NTHREADS_PER_PROC=" + str(_nthreads_per_task) + ] + cmd = [ + _sbatch, "--job-name", job_name, "--nodes", + str(_nodes), "--ntasks-per-node", + str(_ntasks_per_node), "--output", log_filename, "--partition", + _partition, "--time", _timelimit + ] + cmds = [prepare_cmd, env_cmd, cmd] + process = subprocess.run(_with_ssh(cmds), + input=script, + stdout=subprocess.PIPE, + universal_newlines=True, + check=True) + output = process.stdout.rstrip() + logging.info("Received response \"{}\"".format(output)) + # Output looks like "Submitted batch job 133417" + m = re.search(r"Submitted batch job (\d+)", output) + assert m.lastindex == 1 + job_id = m.group(1) + assert job_id != '' + logging.info("Submitted job {}".format(job_id)) + + +# ---------------------------------------------------------------------------- +def _quote_args(args: [str]) -> str: + """converts a command into a shell command string""" + qargs = [] + for arg in args: + qargs.append(shlex.quote(arg)) + return " ".join(qargs) + + +# ---------------------------------------------------------------------------- +def _join_cmds(cmds: [str]) -> str: + """joins several shell commands into a single one""" + return "; ".join(cmds) + + +# ---------------------------------------------------------------------------- +def _with_ssh(cmds: [[str]]) -> [str]: + """adds ssh command to execute command remotely""" + # Add module load commands + cmds = _setup_cmds + cmds + cmd = _join_cmds(map(_quote_args, cmds)) + # Add ssh command prefix + assert _hostname[0] != '-' + return [_ssh, '-l', _username, _hostname, cmd]