From fbbb4e9f10edc8864eab3ec3f91ec89786def01c Mon Sep 17 00:00:00 2001 From: Sam Ritchie Date: Mon, 29 Jun 2020 11:32:07 -0600 Subject: [PATCH 01/15] serious restructure --- caliban/builder.py | 734 +++++++++++++ caliban/cli.py | 3 +- caliban/config/__init__.py | 181 ++++ caliban/{config.py => config/experiment.py} | 217 ++-- caliban/{cloud => docker}/__init__.py | 0 caliban/docker/build.py | 619 +++++++++++ caliban/docker/push.py | 734 +++++++++++++ caliban/expansion.py | 2 +- caliban/{gke => platform/cloud}/__init__.py | 0 caliban/{ => platform}/cloud/core.py | 0 caliban/{ => platform}/cloud/types.py | 0 caliban/platform/cloud/util.py | 119 +++ caliban/platform/gke/__init__.py | 15 + caliban/{ => platform}/gke/cli.py | 0 caliban/{ => platform}/gke/cluster.py | 0 caliban/{ => platform}/gke/constants.py | 0 caliban/{ => platform}/gke/types.py | 0 caliban/{ => platform}/gke/utils.py | 0 caliban/{docker.py => platform/notebook.py} | 0 caliban/platform/run.py | 1056 +++++++++++++++++++ caliban/platform/shell.py | 1056 +++++++++++++++++++ caliban/util.py | 752 ------------- caliban/util/__init__.py | 180 ++++ caliban/util/argparse.py | 129 +++ caliban/util/fs.py | 219 ++++ caliban/util/tqdm.py | 95 ++ tests/caliban/config/test_config.py | 44 + tests/caliban/config/test_experiment.py | 229 ++++ tests/caliban/test_config.py | 257 ----- tests/caliban/test_util.py | 16 - 30 files changed, 5485 insertions(+), 1172 deletions(-) create mode 100644 caliban/builder.py create mode 100644 caliban/config/__init__.py rename caliban/{config.py => config/experiment.py} (60%) rename caliban/{cloud => docker}/__init__.py (100%) create mode 100644 caliban/docker/build.py create mode 100644 caliban/docker/push.py rename caliban/{gke => platform/cloud}/__init__.py (100%) rename caliban/{ => platform}/cloud/core.py (100%) rename caliban/{ => platform}/cloud/types.py (100%) create mode 100644 caliban/platform/cloud/util.py create mode 100644 caliban/platform/gke/__init__.py rename caliban/{ => platform}/gke/cli.py (100%) rename caliban/{ => platform}/gke/cluster.py (100%) rename caliban/{ => platform}/gke/constants.py (100%) rename caliban/{ => platform}/gke/types.py (100%) rename caliban/{ => platform}/gke/utils.py (100%) rename caliban/{docker.py => platform/notebook.py} (100%) create mode 100644 caliban/platform/run.py create mode 100644 caliban/platform/shell.py delete mode 100644 caliban/util.py create mode 100644 caliban/util/__init__.py create mode 100644 caliban/util/argparse.py create mode 100644 caliban/util/fs.py create mode 100644 caliban/util/tqdm.py create mode 100644 tests/caliban/config/test_config.py create mode 100644 tests/caliban/config/test_experiment.py delete mode 100644 tests/caliban/test_config.py diff --git a/caliban/builder.py b/caliban/builder.py new file mode 100644 index 0000000..298aa1f --- /dev/null +++ b/caliban/builder.py @@ -0,0 +1,734 @@ +#!/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. +"""Functions required to interact with Docker to build and run images, shells +and notebooks in a Docker environment. + +""" + +from __future__ import absolute_import, division, print_function + +import json +import os +import subprocess +import sys +from enum import Enum +from pathlib import Path +from typing import (Any, Callable, Dict, Iterable, List, NamedTuple, NewType, + Optional, Union) + +import tqdm +from absl import logging +from blessings import Terminal +from tqdm.utils import _screen_shape_wrapper + +import caliban.config as c +import caliban.util as u +from caliban.history.types import Experiment, Job, JobSpec, JobStatus, Platform +from caliban.history.utils import (create_experiments, generate_container_spec, + get_mem_engine, get_sql_engine, + session_scope) + +t = Terminal() + +DEV_CONTAINER_ROOT = "gcr.io/blueshift-playground/blueshift" +TF_VERSIONS = {"2.2.0", "1.12.3", "1.14.0", "1.15.0"} +DEFAULT_WORKDIR = "/usr/app" +CREDS_DIR = "/.creds" +CONDA_BIN = "/opt/conda/bin/conda" + +ImageId = NewType('ImageId', str) +ArgSeq = NewType('ArgSeq', List[str]) + + +class DockerError(Exception): + """Exception that passes info on a failed Docker command.""" + + def __init__(self, message, cmd, ret_code): + super().__init__(message) + self.message = message + self.cmd = cmd + self.ret_code = ret_code + + @property + def command(self): + return " ".join(self.cmd) + + +class NotebookInstall(Enum): + """Flag to decide what to do .""" + none = 'none' + lab = 'lab' + jupyter = 'jupyter' + + def __str__(self) -> str: + return self.value + + +class Shell(Enum): + """Add new shells here and below, in SHELL_DICT.""" + bash = 'bash' + zsh = 'zsh' + + def __str__(self) -> str: + return self.value + + +# Tuple to track the information required to install and execute some custom +# shell into a container. +ShellData = NamedTuple("ShellData", [("executable", str), + ("packages", List[str])]) + + +def apt_install(*packages: str) -> str: + """Returns a command that will install the supplied list of packages without + requiring confirmation or any user interaction. + """ + package_str = ' '.join(packages) + no_prompt = "DEBIAN_FRONTEND=noninteractive" + return f"{no_prompt} apt-get install --yes --no-install-recommends {package_str}" + + +def apt_command(commands: List[str]) -> List[str]: + """Pre-and-ap-pends the supplied commands with the appropriate in-container and + cleanup command for aptitude. + + """ + update = ["apt-get update"] + cleanup = ["apt-get clean", "rm -rf /var/lib/apt/lists/*"] + return update + commands + cleanup + + +# Dict linking a particular supported shell to the data required to run and +# install the shell inside a container. +# +# : Dict[Shell, ShellData] +SHELL_DICT = { + Shell.bash: ShellData("/bin/bash", []), + Shell.zsh: ShellData("/bin/zsh", ["zsh"]) +} + + +def default_shell() -> Shell: + """Returns the shell to load into the container. Defaults to Shell.bash, but if + the user's SHELL variable refers to a supported sub-shell, returns that + instead. + + """ + ret = Shell.bash + + if "zsh" in os.environ.get("SHELL"): + ret = Shell.zsh + + return ret + + +def adc_location(home_dir: Optional[str] = None) -> str: + """Returns the location for application default credentials, INSIDE the + container (so, hardcoded unix separators), given the supplied home directory. + + """ + if home_dir is None: + home_dir = Path.home() + + return "{}/.config/gcloud/application_default_credentials.json".format( + home_dir) + + +def container_home(): + """Returns the location of the home directory inside the generated + container. + + """ + return "/home/{}".format(u.current_user()) + + +def tf_base_image(job_mode: c.JobMode, tensorflow_version: str) -> str: + """Returns the base image to use, depending on whether or not we're using a + GPU. This is JUST for building our base images for Blueshift; not for + actually using in a job. + + List of available tags: https://hub.docker.com/r/tensorflow/tensorflow/tags + + """ + if tensorflow_version not in TF_VERSIONS: + raise Exception("""{} is not a valid tensorflow version. + Try one of: {}""".format(tensorflow_version, TF_VERSIONS)) + + gpu = "-gpu" if c.gpu(job_mode) else "" + return "tensorflow/tensorflow:{}{}-py3".format(tensorflow_version, gpu) + + +def base_image_suffix(job_mode: c.JobMode) -> str: + return "gpu" if c.gpu(job_mode) else "cpu" + + +def base_image_id(job_mode: c.JobMode) -> str: + """Returns the default base image for all caliban Dockerfiles.""" + base_suffix = base_image_suffix(job_mode) + return "{}:{}".format(DEV_CONTAINER_ROOT, base_suffix) + + +def extras_string(extras: List[str]) -> str: + """Returns the argument passed to `pip install` to install a project from its + setup.py and target a specific set of extras_require dependencies. + + Args: + extras: (potentially empty) list of extra_requires deps. + """ + ret = "." + if len(extras) > 0: + ret += "[{}]".format(','.join(extras)) + return ret + + +def base_extras(job_mode: c.JobMode, path: str, + extras: Optional[List[str]]) -> Optional[List[str]]: + """Returns None if the supplied path doesn't exist (it's assumed it points to a + setup.py file). + + If the path DOES exist, generates a list of extras to install. gpu or cpu are + always added to the beginning of the list, depending on the mode. + + """ + ret = None + + if os.path.exists(path): + base = extras or [] + extra = 'gpu' if c.gpu(job_mode) else 'cpu' + ret = base if extra in base else [extra] + base + + return ret + + +def _dependency_entries(workdir: str, + user_id: int, + user_group: int, + requirements_path: Optional[str] = None, + conda_env_path: Optional[str] = None, + setup_extras: Optional[List[str]] = None) -> str: + """Returns the Dockerfile entries required to install dependencies from either: + + - a requirements.txt file, path supplied by requirements_path + - a conda environment.yml file, path supplied by conda_env_path. + - a setup.py file, if some sequence of dependencies is supplied. + + An empty list for setup_extras means, run `pip install -c .` with no extras. + None for this argument means do nothing. If a list of strings is supplied, + they'll be treated as extras dependency sets. + """ + ret = "" + + if setup_extras is not None: + ret += f""" +COPY --chown={user_id}:{user_group} setup.py {workdir} +RUN /bin/bash -c "pip install --no-cache-dir {extras_string(setup_extras)}" +""" + + if conda_env_path is not None: + ret += f""" +COPY --chown={user_id}:{user_group} {conda_env_path} {workdir} +RUN /bin/bash -c "{CONDA_BIN} env update \ + --quiet --name caliban \ + --file {conda_env_path} && \ + {CONDA_BIN} clean -y -q --all" +""" + + if requirements_path is not None: + ret += f""" +COPY --chown={user_id}:{user_group} {requirements_path} {workdir} +RUN /bin/bash -c "pip install --no-cache-dir -r {requirements_path}" +""" + + return ret + + +def _package_entries(workdir: str, user_id: int, user_group: int, + package: u.Package) -> str: + """Returns the Dockerfile entries required to: + + - copy a directory of code into a docker container + - inject an entrypoint that executes a python module inside that directory. + + Python code runs as modules vs scripts so that we can enforce import hygiene + between files inside a project. + + """ + owner = "{}:{}".format(user_id, user_group) + + arg = package.main_module or package.script_path + + # This needs to use json so that quotes print as double quotes, not single + # quotes. + entrypoint_s = json.dumps(package.executable + [arg]) + + return """ +# Copy project code into the docker container. +COPY --chown={owner} {package_path} {workdir}/{package_path} + +# Declare an entrypoint that actually runs the container. +ENTRYPOINT {entrypoint_s} + """.format_map({ + "owner": owner, + "package_path": package.package_path, + "workdir": workdir, + "entrypoint_s": entrypoint_s + }) + + +def _service_account_entry(user_id: int, user_group: int, credentials_path: str, + docker_credentials_dir: str, + write_adc_placeholder: bool): + """Generates the Dockerfile entries required to transfer a set of Cloud service + account credentials into the Docker container. + + NOTE the write_adc_placeholder variable is here because the "ctpu" script + that we use to interact with TPUs has a bug in it, as of 1/21/2020, where the + script will fail if the application_default_credentials.json file isn't + present, EVEN THOUGH it properly uses the service account credentials + registered with gcloud instead of ADC creds. + + If a service account is present, we write a placeholder string to get past + this problem. This shouldn't matter for anyone else since adc isn't used if a + service account is present. + + """ + container_creds = "{}/credentials.json".format(docker_credentials_dir) + ret = """ +COPY --chown={user_id}:{user_group} {credentials_path} {container_creds} + +# Use the credentials file to activate gcloud, gsutil inside the container. +RUN gcloud auth activate-service-account --key-file={container_creds} && \ + git config --global credential.'https://source.developers.google.com'.helper gcloud.sh + +ENV GOOGLE_APPLICATION_CREDENTIALS={container_creds} +""".format_map({ + "user_id": user_id, + "user_group": user_group, + "credentials_path": credentials_path, + "container_creds": container_creds + }) + + if write_adc_placeholder: + ret += """ +RUN echo "placeholder" >> {} +""".format(adc_location(container_home())) + + return ret + + +def _adc_entry(user_id: int, user_group: int, adc_path: str): + """Returns the Dockerfile line required to transfer the + application_default_credentials.json file into the container's home + directory. + + """ + return """ +COPY --chown={user_id}:{user_group} {adc_path} {adc_loc} + """.format_map({ + "user_id": user_id, + "user_group": user_group, + "adc_path": adc_path, + "adc_loc": adc_location(container_home()) + }) + + +def _credentials_entries(user_id: int, + user_group: int, + adc_path: Optional[str], + credentials_path: Optional[str], + docker_credentials_dir: Optional[str] = None) -> str: + """Returns the Dockerfile entries necessary to copy a user's Cloud credentials + into the Docker container. + + - adc_path is the relative path inside the current directory to an + application_default_credentials.json file containing... well, you get it. + - credentials_path is the relative path inside the current directory to a + JSON credentials file. + - docker_credentials_dir is the relative path inside the docker container + where the JSON file will be copied on build. + + """ + if docker_credentials_dir is None: + docker_credentials_dir = CREDS_DIR + + ret = "" + if credentials_path is not None: + ret += _service_account_entry(user_id, + user_group, + credentials_path, + docker_credentials_dir, + write_adc_placeholder=adc_path is None) + + if adc_path is not None: + ret += _adc_entry(user_id, user_group, adc_path) + + return ret + + +def _notebook_entries(lab: bool = False, version: Optional[str] = None) -> str: + """Returns the Dockerfile entries necessary to install Jupyter{lab}. + + Optionally takes a version string. + + """ + version_suffix = "" + + if version is not None: + version_suffix = "=={}".format(version) + + library = "jupyterlab" if lab else "jupyter" + + return """ +RUN pip install {}{} +""".format(library, version_suffix) + + +def _custom_packages( + user_id: int, + user_group: int, + packages: Optional[List[str]] = None, + shell: Optional[Shell] = None, +) -> str: + """Returns the Dockerfile entries necessary to install custom dependencies for + the supplied shell and sequence of aptitude packages. + + """ + if packages is None: + packages = [] + + if shell is None: + shell = Shell.bash + + ret = "" + + to_install = sorted(packages + SHELL_DICT[shell].packages) + + if len(to_install) != 0: + commands = apt_command([apt_install(*to_install)]) + ret = """ +USER root + +RUN {commands} + +USER {user_id}:{user_group} +""".format_map({ + "commands": " && ".join(commands), + "user_id": user_id, + "user_group": user_group + }) + + return ret + + +def _copy_dir_entry(workdir: str, user_id: int, user_group: int, + dirname: str) -> str: + """Returns the Dockerfile entry necessary to copy a single extra subdirectory + from the current directory into a docker container during build. + + """ + owner = "{}:{}".format(user_id, user_group) + return """# Copy {dirname} into the Docker container. +COPY --chown={owner} {dirname} {workdir}/{dirname} +""".format_map({ + "owner": owner, + "workdir": workdir, + "dirname": dirname + }) + + +def _extra_dir_entries(workdir: str, user_id: int, user_group: int, + extra_dirs: List[str]) -> str: + """Returns the Dockerfile entries necessary to copy all directories in the + extra_dirs list into a docker container during build. + + """ + ret = "" + for d in extra_dirs: + ret += "\n{}".format(_copy_dir_entry(workdir, user_id, user_group, d)) + return ret + + +def _dockerfile_template( + job_mode: c.JobMode, + workdir: Optional[str] = None, + base_image_fn: Optional[Callable[[c.JobMode], str]] = None, + package: Optional[Union[List, u.Package]] = None, + requirements_path: Optional[str] = None, + conda_env_path: Optional[str] = None, + setup_extras: Optional[List[str]] = None, + adc_path: Optional[str] = None, + credentials_path: Optional[str] = None, + jupyter_version: Optional[str] = None, + inject_notebook: NotebookInstall = NotebookInstall.none, + shell: Optional[Shell] = None, + extra_dirs: Optional[List[str]] = None, + caliban_config: Optional[Dict[str, Any]] = None) -> str: + """Returns a Dockerfile that builds on a local CPU or GPU base image (depending + on the value of job_mode) to create a container that: + + - installs any dependency specified in a requirements.txt file living at + requirements_path, a conda environment at conda_env_path, or any + dependencies in a setup.py file, including extra dependencies, if + setup_extras isn't None + - injects gcloud credentials into the container, so Cloud interaction works + just like it does locally + - potentially installs a custom shell, or jupyterlab for notebook support + - copies all source needed by the main module specified by package, and + potentially injects an entrypoint that, on run, will run that main module + + Most functions that call _dockerfile_template pass along any kwargs that they + receive. It should be enough to add kwargs here, then rely on that mechanism + to pass them along, vs adding kwargs all the way down the call chain. + + Supply a custom base_image_fn (function from job_mode -> image ID) to inject + more complex Docker commands into the Caliban environments by, for example, + building your own image on top of the TF base images, then using that. + + """ + uid = os.getuid() + gid = os.getgid() + username = u.current_user() + + if isinstance(package, list): + package = u.Package(*package) + + if workdir is None: + workdir = DEFAULT_WORKDIR + + if base_image_fn is None: + base_image_fn = base_image_id + + base_image = base_image_fn(job_mode) + + dockerfile = """ +FROM {base_image} + +# Create the same group we're using on the host machine. +RUN [ $(getent group {gid}) ] || groupadd --gid {gid} {gid} + +# Create the user by name. --no-log-init guards against a crash with large user +# IDs. +RUN useradd --no-log-init --no-create-home -u {uid} -g {gid} --shell /bin/bash {username} + +# The directory is created by root. This sets permissions so that any user can +# access the folder. +RUN mkdir -m 777 {workdir} {creds_dir} {c_home} + +ENV HOME={c_home} + +WORKDIR {workdir} + +USER {uid}:{gid} +""".format_map({ + "base_image": base_image, + "username": username, + "uid": uid, + "gid": gid, + "workdir": workdir, + "c_home": container_home(), + "creds_dir": CREDS_DIR + }) + dockerfile += _credentials_entries(uid, + gid, + adc_path=adc_path, + credentials_path=credentials_path) + + dockerfile += _dependency_entries(workdir, + uid, + gid, + requirements_path=requirements_path, + conda_env_path=conda_env_path, + setup_extras=setup_extras) + + if inject_notebook.value != 'none': + install_lab = inject_notebook == NotebookInstall.lab + dockerfile += _notebook_entries(lab=install_lab, version=jupyter_version) + + if extra_dirs is not None: + dockerfile += _extra_dir_entries(workdir, uid, gid, extra_dirs) + + dockerfile += _custom_packages(uid, + gid, + packages=c.apt_packages( + caliban_config, job_mode), + shell=shell) + + if package is not None: + # The actual entrypoint and final copied code. + dockerfile += _package_entries(workdir, uid, gid, package) + + return dockerfile + + +def docker_image_id(output: str) -> ImageId: + """Accepts a string containing the output of a successful `docker build` + command and parses the Docker image ID from the stream. + + NOTE this is probably quite brittle! I can imagine this breaking quite easily + on a Docker upgrade. + + """ + return ImageId(output.splitlines()[-1].split()[-1]) + + +def build_image(job_mode: c.JobMode, + build_path: str, + credentials_path: Optional[str] = None, + adc_path: Optional[str] = None, + no_cache: bool = False, + **kwargs) -> str: + """Builds a Docker image by generating a Dockerfile and passing it to `docker + build` via stdin. All output from the `docker build` process prints to + stdout. + + Returns the image ID of the new docker container; if the command fails, + throws on error with information about the command and any issues that caused + the problem. + + """ + with u.TempCopy(credentials_path, + tmp_name=".caliban_default_creds.json") as creds: + with u.TempCopy(adc_path, tmp_name=".caliban_adc_creds.json") as adc: + cache_args = ["--no-cache"] if no_cache else [] + cmd = ["docker", "build"] + cache_args + ["--rm", "-f-", build_path] + + dockerfile = _dockerfile_template(job_mode, + credentials_path=creds, + adc_path=adc, + **kwargs) + + joined_cmd = " ".join(cmd) + logging.info("Running command: {}".format(joined_cmd)) + + try: + output, ret_code = u.capture_stdout(cmd, input_str=dockerfile) + if ret_code == 0: + return docker_image_id(output) + else: + error_msg = "Docker failed with error code {}.".format(ret_code) + raise DockerError(error_msg, cmd, ret_code) + + except subprocess.CalledProcessError as e: + logging.error(e.output) + logging.error(e.stderr) + + +def _image_tag_for_project(project_id: str, image_id: str) -> str: + """Generate the GCR Docker image tag for the supplied pair of project_id and + image_id. + + This function properly handles "domain scoped projects", where the project ID + contains a domain name and project ID separated by : + https://cloud.google.com/container-registry/docs/overview#domain-scoped_projects. + + """ + project_s = project_id.replace(":", "/") + return "gcr.io/{}/{}:latest".format(project_s, image_id) + + +def push_uuid_tag(project_id: str, image_id: str) -> str: + """Takes a base image and tags it for upload, then pushes it to a remote Google + Container Registry. + + Returns the tag on a successful push. + + TODO should this just check first before attempting to push if the image + exists? Immutable names means that if the tag is up there, we're done. + Potentially use docker-py for this. + + """ + image_tag = _image_tag_for_project(project_id, image_id) + subprocess.run(["docker", "tag", image_id, image_tag], check=True) + subprocess.run(["docker", "push", image_tag], check=True) + return image_tag + + +def _run_cmd(job_mode: c.JobMode, + run_args: Optional[List[str]] = None) -> List[str]: + """Returns the sequence of commands for the subprocess run functions required + to execute `docker run`. in CPU or GPU mode, depending on the value of + job_mode. + + Keyword args: + - run_args: list of args to pass to docker run. + + """ + if run_args is None: + run_args = [] + + runtime = ["--runtime", "nvidia"] if c.gpu(job_mode) else [] + return ["docker", "run"] + runtime + ["--ipc", "host"] + run_args + + +def _home_mount_cmds(enable_home_mount: bool) -> List[str]: + """Returns the argument needed by Docker to mount a user's local home directory + into the home directory location inside their container. + + If enable_home_mount is false returns an empty list. + + """ + ret = [] + if enable_home_mount: + ret = ["-v", "{}:{}".format(Path.home(), container_home())] + return ret + + +def _interactive_opts(workdir: str) -> List[str]: + """Returns the basic arguments we want to run a docker process locally. + + """ + return [ + "-w", workdir, \ + "-u", "{}:{}".format(os.getuid(), os.getgid()), \ + "-v", "{}:{}".format(os.getcwd(), workdir) \ + ] + + +def log_job_spec_instance(job_spec: JobSpec, i: int) -> JobSpec: + """Prints logging as a side effect for the supplied sequence of job specs + generated from an experiment definition; returns the input job spec. + + """ + args = c.experiment_to_args(job_spec.experiment.kwargs, + job_spec.experiment.args) + logging.info("") + logging.info("Job {} - Experiment args: {}".format(i, t.yellow(str(args)))) + return job_spec + + +def logged_job_specs(job_specs: Iterable[JobSpec]) -> Iterable[JobSpec]: + """Generates an iterable of job specs that should be passed to `docker run` to + execute the experiments defined by the supplied iterable. + + """ + for i, s in enumerate(job_specs, 1): + yield log_job_spec_instance(s, i) + + +def execute_dry_run(job_specs: Iterable[JobSpec]) -> None: + """Expands the supplied sequence of experiments into sequences of args and logs + the jobs that WOULD have been executed, had the dry run flag not been + applied. + + """ + list(logged_job_specs(job_specs)) + + logging.info('') + logging.info( + t.yellow("To build your image and execute these jobs, \ +run your command again without {}.".format(c.DRY_RUN_FLAG))) + logging.info('') + return None diff --git a/caliban/cli.py b/caliban/cli.py index e499b51..1a2eeee 100644 --- a/caliban/cli.py +++ b/caliban/cli.py @@ -26,6 +26,7 @@ import caliban.cloud.types as ct import caliban.config as conf +import caliban.config.experiment as ce import caliban.docker as docker import caliban.gke as gke import caliban.gke.constants as gke_k @@ -352,7 +353,7 @@ def job_name_arg(parser): def experiment_config_arg(parser): parser.add_argument( "--experiment_config", - type=conf.load_experiment_config, + type=ce.load_experiment_config, help="Path to an experiment config, or 'stdin' to read from stdin.") diff --git a/caliban/config/__init__.py b/caliban/config/__init__.py new file mode 100644 index 0000000..28e00a4 --- /dev/null +++ b/caliban/config/__init__.py @@ -0,0 +1,181 @@ +#!/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. +""" +Utilities for our job runner, for working with configs. +""" + +from __future__ import absolute_import, division, print_function + +import argparse +import os +import sys +from enum import Enum +from typing import Any, Dict, List, Optional + +import commentjson +import yaml + +import caliban.cloud.types as ct + + +class JobMode(str, Enum): + CPU = 'CPU' + GPU = 'GPU' + + +# Special config for Caliban. +CalibanConfig = Dict[str, Any] + +DRY_RUN_FLAG = "--dry_run" +CALIBAN_CONFIG = ".calibanconfig.json" + +# Defaults for various input values that we can supply given some partial set +# of info from the CLI. +DEFAULT_REGION = ct.US.central1 + +# : Dict[JobMode, ct.MachineType] +DEFAULT_MACHINE_TYPE = { + JobMode.CPU: ct.MachineType.highcpu_32, + JobMode.GPU: ct.MachineType.standard_8 +} +DEFAULT_GPU = ct.GPU.P100 + +# Config to supply for CPU jobs. +DEFAULT_ACCELERATOR_CONFIG = { + "count": 0, + "type": "ACCELERATOR_TYPE_UNSPECIFIED" +} + + +def gpu(job_mode: JobMode) -> bool: + """Returns True if the supplied JobMode is JobMode.GPU, False otherwise. + + """ + return job_mode == JobMode.GPU + + +def load_yaml_config(path): + """returns the config parsed based on the info in the flags. + + Grabs the config file, written in yaml, slurps it in. + """ + with open(path) as f: + config = yaml.load(f, Loader=yaml.FullLoader) + + return config + + +def load_config(path, mode='yaml'): + """Load a JSON or YAML config. + + """ + if mode == 'json': + with open(path) as f: + return commentjson.load(f) + + return load_yaml_config(path) + + +def valid_json(path: str) -> Dict[str, Any]: + """Loads JSON if the path points to a valid JSON file; otherwise, throws an + exception that's picked up by argparse. + + """ + try: + return load_config(path, mode='json') + except commentjson.JSONLibraryException: + raise argparse.ArgumentTypeError( + """File '{}' doesn't seem to contain valid JSON. Try again!""".format( + path)) + + +def extract_script_args(m: Dict[str, Any]) -> List[str]: + """Strip off the "--" argument if it was passed in as a separator.""" + script_args = m.get("script_args") + if script_args is None or script_args == []: + return script_args + + head, *tail = script_args + + return tail if head == "--" else script_args + + +def extract_project_id(m: Dict[str, Any]) -> str: + """Attempts to extract the project_id from the args; falls back to an + environment variable, or exits if this isn't available. There's no sensible + default available. + + """ + project_id = m.get("project_id") or os.environ.get("PROJECT_ID") + + if project_id is None: + print() + print( + "\nNo project_id found. 'caliban cloud' requires that you either set a \n\ +$PROJECT_ID environment variable with the ID of your Cloud project, or pass one \n\ +explicitly via --project_id. Try again, please!") + print() + + sys.exit(1) + + return project_id + + +def extract_region(m: Dict[str, Any]) -> ct.Region: + """Returns the region specified in the args; defaults to an environment + variable. If that's not supplied defaults to the default cloud provider from + caliban.cloud. + + """ + region = m.get("region") or os.environ.get("REGION") + + if region: + return ct.parse_region(region) + + return DEFAULT_REGION + + +def extract_zone(m: Dict[str, Any]) -> str: + return "{}-a".format(extract_region(m)) + + +def extract_cloud_key(m: Dict[str, Any]) -> Optional[str]: + """Returns the Google service account key filepath specified in the args; + defaults to the $GOOGLE_APPLICATION_CREDENTIALS variable. + + """ + return m.get("cloud_key") or \ + os.environ.get("GOOGLE_APPLICATION_CREDENTIALS") + + +def apt_packages(conf: CalibanConfig, mode: JobMode) -> List[str]: + """Returns the list of aptitude packages that should be installed to satisfy + the requests in the config. + + """ + packages = conf.get("apt_packages") or {} + + if isinstance(packages, dict): + k = "gpu" if gpu(mode) else "cpu" + return packages.get(k, []) + + elif isinstance(packages, list): + return packages + + else: + raise argparse.ArgumentTypeError( + """{}'s "apt_packages" entry must be a dictionary or list, not '{}'""". + format(CALIBAN_CONFIG, packages)) diff --git a/caliban/config.py b/caliban/config/experiment.py similarity index 60% rename from caliban/config.py rename to caliban/config/experiment.py index a602c40..0fcc1bc 100644 --- a/caliban/config.py +++ b/caliban/config/experiment.py @@ -13,23 +13,21 @@ # 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. -""" -Utilities for our job runner, for working with configs. +"""Utilities for working with experiment.json files. + """ from __future__ import absolute_import, division, print_function import argparse import itertools -import os -import sys -from enum import Enum -from typing import Any, Dict, List, Union, Optional import re +import sys +from collections import ChainMap +from typing import Any, Dict, List, Optional, Tuple, Union + import commentjson -import yaml -import caliban.cloud.types as ct import caliban.util as u # int, str and bool are allowed in a final experiment; lists are markers for @@ -48,168 +46,97 @@ Experiment = Dict[str, ExpValue] -# Mode -class JobMode(str, Enum): - CPU = 'CPU' - GPU = 'GPU' - - -# Special config for Caliban. -CalibanConfig = Dict[str, Any] - -DRY_RUN_FLAG = "--dry_run" -CALIBAN_CONFIG = ".calibanconfig.json" - -# Defaults for various input values that we can supply given some partial set -# of info from the CLI. -DEFAULT_REGION = ct.US.central1 - -# : Dict[JobMode, ct.MachineType] -DEFAULT_MACHINE_TYPE = { - JobMode.CPU: ct.MachineType.highcpu_32, - JobMode.GPU: ct.MachineType.standard_8 -} -DEFAULT_GPU = ct.GPU.P100 - -# Config to supply for CPU jobs. -DEFAULT_ACCELERATOR_CONFIG = { - "count": 0, - "type": "ACCELERATOR_TYPE_UNSPECIFIED" -} - - -def gpu(job_mode: JobMode) -> bool: - """Returns True if the supplied JobMode is JobMode.GPU, False otherwise. - - """ - return job_mode == JobMode.GPU - - -def load_yaml_config(path): - """returns the config parsed based on the info in the flags. - - Grabs the config file, written in yaml, slurps it in. - """ - with open(path) as f: - config = yaml.load(f, Loader=yaml.FullLoader) - - return config - - -def load_config(path, mode='yaml'): - """Load a JSON or YAML config. - - """ - if mode == 'json': - with open(path) as f: - return commentjson.load(f) - - return load_yaml_config(path) - - -def valid_json(path: str) -> Dict[str, Any]: - """Loads JSON if the path points to a valid JSON file; otherwise, throws an - exception that's picked up by argparse. - - """ - try: - return load_config(path, mode='json') - except commentjson.JSONLibraryException: - raise argparse.ArgumentTypeError( - """File '{}' doesn't seem to contain valid JSON. Try again!""".format( - path)) - - -def extract_script_args(m: Dict[str, Any]) -> List[str]: - """Strip off the "--" argument if it was passed in as a separator.""" - script_args = m.get("script_args") - if script_args is None or script_args == []: - return script_args - - head, *tail = script_args - - return tail if head == "--" else script_args - - -def extract_project_id(m: Dict[str, Any]) -> str: - """Attempts to extract the project_id from the args; falls back to an - environment variable, or exits if this isn't available. There's no sensible - default available. - +def _is_compound_key(s: Any) -> bool: + """ compound key is defined as a string which uses square brackets to enclose + a comma-separated list, e.g. "[batch_size,learning_rate]" or "[a,b,c]" """ - project_id = m.get("project_id") or os.environ.get("PROJECT_ID") - if project_id is None: - print() - print( - "\nNo project_id found. 'caliban cloud' requires that you either set a \n\ -$PROJECT_ID environment variable with the ID of your Cloud project, or pass one \n\ -explicitly via --project_id. Try again, please!") - print() + if type(s) is not str or len(s) <= 2: + return False + else: + return s[0] == '[' and s[-1] == ']' - sys.exit(1) - return project_id +def _tupleize_compound_key(k: str) -> List[str]: + """ converts a JSON-input compound key into a tuple """ + assert _is_compound_key(k), "{} must be a valid compound key".format(k) + return tuple([x.strip() for x in k.strip('][').split(',')]) -def extract_region(m: Dict[str, Any]) -> ct.Region: - """Returns the region specified in the args; defaults to an environment - variable. If that's not supplied defaults to the default cloud provider from - caliban.cloud. +def _tupleize_compound_value( + v: Union[List, bool, str, int, float]) -> Union[List, Tuple]: + """ list of lists -> list of tuples + list of primitives -> tuple of primitives + single primitive -> length-1 tuple of that primitive + E.g., [[0,1],[3,4]] -> [(0,1),(3,4)] + [0,1] -> (0,1) + 0 -> (0, ) """ - region = m.get("region") or os.environ.get("REGION") - - if region: - return ct.parse_region(region) - - return DEFAULT_REGION + if isinstance(v, list): + if isinstance(v[0], list): + # v is list of lists + return [tuple(vi) for vi in v] + else: + # v is list of primitives + return tuple(v) + else: + # v is a single primitive (bool, str, int, float) + return tuple([v]) -def extract_zone(m: Dict[str, Any]) -> str: - return "{}-a".format(extract_region(m)) +def _tupleize_compound_item(k: Union[Tuple, str], v: Any) -> Dict: + """ converts a JSON-input compound key/value pair into a dictionary of tuples """ + if _is_compound_key(k): + return {_tupleize_compound_key(k): _tupleize_compound_value(v)} + else: + return {k: v} -def extract_cloud_key(m: Dict[str, Any]) -> Optional[str]: - """Returns the Google service account key filepath specified in the args; - defaults to the $GOOGLE_APPLICATION_CREDENTIALS variable. +def tupleize_dict(m: Dict) -> Dict: + """ given a dictionary with compound keys, converts those keys to tuples, and + converts the corresponding values to a tuple or list of tuples + Compound key: a string which uses square brackets to enclose + a comma-separated list, e.g. "[batch_size,learning_rate]" or "[a,b,c]" """ - return m.get("cloud_key") or \ - os.environ.get("GOOGLE_APPLICATION_CREDENTIALS") + formatted_items = [_tupleize_compound_item(k, v) for k, v in m.items()] + return dict(ChainMap(*formatted_items)) -def apt_packages(conf: CalibanConfig, mode: JobMode) -> List[str]: - """Returns the list of aptitude packages that should be installed to satisfy - the requests in the config. +def _expand_compound_pair(k: Union[Tuple, str], v: Any) -> Dict: + """ given a key-value pair k v, where k is either: + a) a primitive representing a single, e.g. k = 'key', v = 'value', or + b) a tuple of primitives representing multiple keys, e.g. k = ('key1','key2'), v = ('value1', 'value2') + this function returns the corresponding dictionary without compound keys """ - packages = conf.get("apt_packages") or {} - - if isinstance(packages, dict): - k = "gpu" if gpu(mode) else "cpu" - return packages.get(k, []) - - elif isinstance(packages, list): - return packages + if isinstance(k, tuple): + if not isinstance(v, tuple): + raise argparse.ArgumentTypeError( + """function _expand_compound_pair(k, v) requires that if type(k) is tuple, + type(v) must also be tuple.""") + else: + return dict(zip(k, v)) else: - raise argparse.ArgumentTypeError( - """{}'s "apt_packages" entry must be a dictionary or list, not '{}'""". - format(CALIBAN_CONFIG, packages)) + return {k: v} -def caliban_config() -> CalibanConfig: - """Returns a dict that represents a `.calibanconfig.json` file if present, - empty dictionary otherwise. +def expand_compound_dict(m: Union[Dict, List]) -> Union[Dict, List]: + """ given a dictionary with some compound keys, aka tuples, + returns a dictionary which each compound key separated into primitives + given a list of such dictionaries, will apply the transformation + described above to each dictionary and return the list, maintaining + structure """ - if not os.path.isfile(CALIBAN_CONFIG): - return {} - with open(CALIBAN_CONFIG) as f: - conf = commentjson.load(f) - return conf + if isinstance(m, list): + return [expand_compound_dict(mi) for mi in m] + else: + expanded_dicts = [_expand_compound_pair(k, v) for k, v in m.items()] + return dict(ChainMap(*expanded_dicts)) def expand_experiment_config(items: ExpConf) -> List[Experiment]: diff --git a/caliban/cloud/__init__.py b/caliban/docker/__init__.py similarity index 100% rename from caliban/cloud/__init__.py rename to caliban/docker/__init__.py diff --git a/caliban/docker/build.py b/caliban/docker/build.py new file mode 100644 index 0000000..1d16aa5 --- /dev/null +++ b/caliban/docker/build.py @@ -0,0 +1,619 @@ +#!/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. +"""Functions required to interact with Docker to build and run images, shells +and notebooks in a Docker environment. + +""" + +from __future__ import absolute_import, division, print_function + +import json +import os +import subprocess +from enum import Enum +from pathlib import Path +from typing import (Any, Callable, Dict, List, NamedTuple, NewType, Optional, + Union) + +from absl import logging +from blessings import Terminal + +import caliban.config as c +import caliban.util as u + +t = Terminal() + +DEV_CONTAINER_ROOT = "gcr.io/blueshift-playground/blueshift" +TF_VERSIONS = {"2.2.0", "1.12.3", "1.14.0", "1.15.0"} +DEFAULT_WORKDIR = "/usr/app" +CREDS_DIR = "/.creds" +CONDA_BIN = "/opt/conda/bin/conda" + +ImageId = NewType('ImageId', str) +ArgSeq = NewType('ArgSeq', List[str]) + + +class DockerError(Exception): + """Exception that passes info on a failed Docker command.""" + + def __init__(self, message, cmd, ret_code): + super().__init__(message) + self.message = message + self.cmd = cmd + self.ret_code = ret_code + + @property + def command(self): + return " ".join(self.cmd) + + +class NotebookInstall(Enum): + """Flag to decide what to do .""" + none = 'none' + lab = 'lab' + jupyter = 'jupyter' + + def __str__(self) -> str: + return self.value + + +class Shell(Enum): + """Add new shells here and below, in SHELL_DICT.""" + bash = 'bash' + zsh = 'zsh' + + def __str__(self) -> str: + return self.value + + +# Tuple to track the information required to install and execute some custom +# shell into a container. +ShellData = NamedTuple("ShellData", [("executable", str), + ("packages", List[str])]) + + +def apt_install(*packages: str) -> str: + """Returns a command that will install the supplied list of packages without + requiring confirmation or any user interaction. + """ + package_str = ' '.join(packages) + no_prompt = "DEBIAN_FRONTEND=noninteractive" + return f"{no_prompt} apt-get install --yes --no-install-recommends {package_str}" + + +def apt_command(commands: List[str]) -> List[str]: + """Pre-and-ap-pends the supplied commands with the appropriate in-container and + cleanup command for aptitude. + + """ + update = ["apt-get update"] + cleanup = ["apt-get clean", "rm -rf /var/lib/apt/lists/*"] + return update + commands + cleanup + + +# Dict linking a particular supported shell to the data required to run and +# install the shell inside a container. +# +# : Dict[Shell, ShellData] +SHELL_DICT = { + Shell.bash: ShellData("/bin/bash", []), + Shell.zsh: ShellData("/bin/zsh", ["zsh"]) +} + + +def default_shell() -> Shell: + """Returns the shell to load into the container. Defaults to Shell.bash, but if + the user's SHELL variable refers to a supported sub-shell, returns that + instead. + + """ + ret = Shell.bash + + if "zsh" in os.environ.get("SHELL"): + ret = Shell.zsh + + return ret + + +def adc_location(home_dir: Optional[str] = None) -> str: + """Returns the location for application default credentials, INSIDE the + container (so, hardcoded unix separators), given the supplied home directory. + + """ + if home_dir is None: + home_dir = Path.home() + + return "{}/.config/gcloud/application_default_credentials.json".format( + home_dir) + + +def container_home(): + """Returns the location of the home directory inside the generated + container. + + """ + return "/home/{}".format(u.current_user()) + + +def tf_base_image(job_mode: c.JobMode, tensorflow_version: str) -> str: + """Returns the base image to use, depending on whether or not we're using a + GPU. This is JUST for building our base images for Blueshift; not for + actually using in a job. + + List of available tags: https://hub.docker.com/r/tensorflow/tensorflow/tags + + """ + if tensorflow_version not in TF_VERSIONS: + raise Exception("""{} is not a valid tensorflow version. + Try one of: {}""".format(tensorflow_version, TF_VERSIONS)) + + gpu = "-gpu" if c.gpu(job_mode) else "" + return "tensorflow/tensorflow:{}{}-py3".format(tensorflow_version, gpu) + + +def base_image_suffix(job_mode: c.JobMode) -> str: + return "gpu" if c.gpu(job_mode) else "cpu" + + +def base_image_id(job_mode: c.JobMode) -> str: + """Returns the default base image for all caliban Dockerfiles.""" + base_suffix = base_image_suffix(job_mode) + return "{}:{}".format(DEV_CONTAINER_ROOT, base_suffix) + + +def extras_string(extras: List[str]) -> str: + """Returns the argument passed to `pip install` to install a project from its + setup.py and target a specific set of extras_require dependencies. + + Args: + extras: (potentially empty) list of extra_requires deps. + """ + ret = "." + if len(extras) > 0: + ret += "[{}]".format(','.join(extras)) + return ret + + +def base_extras(job_mode: c.JobMode, path: str, + extras: Optional[List[str]]) -> Optional[List[str]]: + """Returns None if the supplied path doesn't exist (it's assumed it points to a + setup.py file). + + If the path DOES exist, generates a list of extras to install. gpu or cpu are + always added to the beginning of the list, depending on the mode. + + """ + ret = None + + if os.path.exists(path): + base = extras or [] + extra = 'gpu' if c.gpu(job_mode) else 'cpu' + ret = base if extra in base else [extra] + base + + return ret + + +def _dependency_entries(workdir: str, + user_id: int, + user_group: int, + requirements_path: Optional[str] = None, + conda_env_path: Optional[str] = None, + setup_extras: Optional[List[str]] = None) -> str: + """Returns the Dockerfile entries required to install dependencies from either: + + - a requirements.txt file, path supplied by requirements_path + - a conda environment.yml file, path supplied by conda_env_path. + - a setup.py file, if some sequence of dependencies is supplied. + + An empty list for setup_extras means, run `pip install -c .` with no extras. + None for this argument means do nothing. If a list of strings is supplied, + they'll be treated as extras dependency sets. + """ + ret = "" + + if setup_extras is not None: + ret += f""" +COPY --chown={user_id}:{user_group} setup.py {workdir} +RUN /bin/bash -c "pip install --no-cache-dir {extras_string(setup_extras)}" +""" + + if conda_env_path is not None: + ret += f""" +COPY --chown={user_id}:{user_group} {conda_env_path} {workdir} +RUN /bin/bash -c "{CONDA_BIN} env update \ + --quiet --name caliban \ + --file {conda_env_path} && \ + {CONDA_BIN} clean -y -q --all" +""" + + if requirements_path is not None: + ret += f""" +COPY --chown={user_id}:{user_group} {requirements_path} {workdir} +RUN /bin/bash -c "pip install --no-cache-dir -r {requirements_path}" +""" + + return ret + + +def _package_entries(workdir: str, user_id: int, user_group: int, + package: u.Package) -> str: + """Returns the Dockerfile entries required to: + + - copy a directory of code into a docker container + - inject an entrypoint that executes a python module inside that directory. + + Python code runs as modules vs scripts so that we can enforce import hygiene + between files inside a project. + + """ + owner = "{}:{}".format(user_id, user_group) + + arg = package.main_module or package.script_path + + # This needs to use json so that quotes print as double quotes, not single + # quotes. + entrypoint_s = json.dumps(package.executable + [arg]) + + return """ +# Copy project code into the docker container. +COPY --chown={owner} {package_path} {workdir}/{package_path} + +# Declare an entrypoint that actually runs the container. +ENTRYPOINT {entrypoint_s} + """.format_map({ + "owner": owner, + "package_path": package.package_path, + "workdir": workdir, + "entrypoint_s": entrypoint_s + }) + + +def _service_account_entry(user_id: int, user_group: int, credentials_path: str, + docker_credentials_dir: str, + write_adc_placeholder: bool): + """Generates the Dockerfile entries required to transfer a set of Cloud service + account credentials into the Docker container. + + NOTE the write_adc_placeholder variable is here because the "ctpu" script + that we use to interact with TPUs has a bug in it, as of 1/21/2020, where the + script will fail if the application_default_credentials.json file isn't + present, EVEN THOUGH it properly uses the service account credentials + registered with gcloud instead of ADC creds. + + If a service account is present, we write a placeholder string to get past + this problem. This shouldn't matter for anyone else since adc isn't used if a + service account is present. + + """ + container_creds = "{}/credentials.json".format(docker_credentials_dir) + ret = """ +COPY --chown={user_id}:{user_group} {credentials_path} {container_creds} + +# Use the credentials file to activate gcloud, gsutil inside the container. +RUN gcloud auth activate-service-account --key-file={container_creds} && \ + git config --global credential.'https://source.developers.google.com'.helper gcloud.sh + +ENV GOOGLE_APPLICATION_CREDENTIALS={container_creds} +""".format_map({ + "user_id": user_id, + "user_group": user_group, + "credentials_path": credentials_path, + "container_creds": container_creds + }) + + if write_adc_placeholder: + ret += """ +RUN echo "placeholder" >> {} +""".format(adc_location(container_home())) + + return ret + + +def _adc_entry(user_id: int, user_group: int, adc_path: str): + """Returns the Dockerfile line required to transfer the + application_default_credentials.json file into the container's home + directory. + + """ + return """ +COPY --chown={user_id}:{user_group} {adc_path} {adc_loc} + """.format_map({ + "user_id": user_id, + "user_group": user_group, + "adc_path": adc_path, + "adc_loc": adc_location(container_home()) + }) + + +def _credentials_entries(user_id: int, + user_group: int, + adc_path: Optional[str], + credentials_path: Optional[str], + docker_credentials_dir: Optional[str] = None) -> str: + """Returns the Dockerfile entries necessary to copy a user's Cloud credentials + into the Docker container. + + - adc_path is the relative path inside the current directory to an + application_default_credentials.json file containing... well, you get it. + - credentials_path is the relative path inside the current directory to a + JSON credentials file. + - docker_credentials_dir is the relative path inside the docker container + where the JSON file will be copied on build. + + """ + if docker_credentials_dir is None: + docker_credentials_dir = CREDS_DIR + + ret = "" + if credentials_path is not None: + ret += _service_account_entry(user_id, + user_group, + credentials_path, + docker_credentials_dir, + write_adc_placeholder=adc_path is None) + + if adc_path is not None: + ret += _adc_entry(user_id, user_group, adc_path) + + return ret + + +def _notebook_entries(lab: bool = False, version: Optional[str] = None) -> str: + """Returns the Dockerfile entries necessary to install Jupyter{lab}. + + Optionally takes a version string. + + """ + version_suffix = "" + + if version is not None: + version_suffix = "=={}".format(version) + + library = "jupyterlab" if lab else "jupyter" + + return """ +RUN pip install {}{} +""".format(library, version_suffix) + + +def _custom_packages( + user_id: int, + user_group: int, + packages: Optional[List[str]] = None, + shell: Optional[Shell] = None, +) -> str: + """Returns the Dockerfile entries necessary to install custom dependencies for + the supplied shell and sequence of aptitude packages. + + """ + if packages is None: + packages = [] + + if shell is None: + shell = Shell.bash + + ret = "" + + to_install = sorted(packages + SHELL_DICT[shell].packages) + + if len(to_install) != 0: + commands = apt_command([apt_install(*to_install)]) + ret = """ +USER root + +RUN {commands} + +USER {user_id}:{user_group} +""".format_map({ + "commands": " && ".join(commands), + "user_id": user_id, + "user_group": user_group + }) + + return ret + + +def _copy_dir_entry(workdir: str, user_id: int, user_group: int, + dirname: str) -> str: + """Returns the Dockerfile entry necessary to copy a single extra subdirectory + from the current directory into a docker container during build. + + """ + owner = "{}:{}".format(user_id, user_group) + return """# Copy {dirname} into the Docker container. +COPY --chown={owner} {dirname} {workdir}/{dirname} +""".format_map({ + "owner": owner, + "workdir": workdir, + "dirname": dirname + }) + + +def _extra_dir_entries(workdir: str, user_id: int, user_group: int, + extra_dirs: List[str]) -> str: + """Returns the Dockerfile entries necessary to copy all directories in the + extra_dirs list into a docker container during build. + + """ + ret = "" + for d in extra_dirs: + ret += "\n{}".format(_copy_dir_entry(workdir, user_id, user_group, d)) + return ret + + +def _dockerfile_template( + job_mode: c.JobMode, + workdir: Optional[str] = None, + base_image_fn: Optional[Callable[[c.JobMode], str]] = None, + package: Optional[Union[List, u.Package]] = None, + requirements_path: Optional[str] = None, + conda_env_path: Optional[str] = None, + setup_extras: Optional[List[str]] = None, + adc_path: Optional[str] = None, + credentials_path: Optional[str] = None, + jupyter_version: Optional[str] = None, + inject_notebook: NotebookInstall = NotebookInstall.none, + shell: Optional[Shell] = None, + extra_dirs: Optional[List[str]] = None, + caliban_config: Optional[Dict[str, Any]] = None) -> str: + """Returns a Dockerfile that builds on a local CPU or GPU base image (depending + on the value of job_mode) to create a container that: + + - installs any dependency specified in a requirements.txt file living at + requirements_path, a conda environment at conda_env_path, or any + dependencies in a setup.py file, including extra dependencies, if + setup_extras isn't None + - injects gcloud credentials into the container, so Cloud interaction works + just like it does locally + - potentially installs a custom shell, or jupyterlab for notebook support + - copies all source needed by the main module specified by package, and + potentially injects an entrypoint that, on run, will run that main module + + Most functions that call _dockerfile_template pass along any kwargs that they + receive. It should be enough to add kwargs here, then rely on that mechanism + to pass them along, vs adding kwargs all the way down the call chain. + + Supply a custom base_image_fn (function from job_mode -> image ID) to inject + more complex Docker commands into the Caliban environments by, for example, + building your own image on top of the TF base images, then using that. + + """ + uid = os.getuid() + gid = os.getgid() + username = u.current_user() + + if isinstance(package, list): + package = u.Package(*package) + + if workdir is None: + workdir = DEFAULT_WORKDIR + + if base_image_fn is None: + base_image_fn = base_image_id + + base_image = base_image_fn(job_mode) + + dockerfile = """ +FROM {base_image} + +# Create the same group we're using on the host machine. +RUN [ $(getent group {gid}) ] || groupadd --gid {gid} {gid} + +# Create the user by name. --no-log-init guards against a crash with large user +# IDs. +RUN useradd --no-log-init --no-create-home -u {uid} -g {gid} --shell /bin/bash {username} + +# The directory is created by root. This sets permissions so that any user can +# access the folder. +RUN mkdir -m 777 {workdir} {creds_dir} {c_home} + +ENV HOME={c_home} + +WORKDIR {workdir} + +USER {uid}:{gid} +""".format_map({ + "base_image": base_image, + "username": username, + "uid": uid, + "gid": gid, + "workdir": workdir, + "c_home": container_home(), + "creds_dir": CREDS_DIR + }) + dockerfile += _credentials_entries(uid, + gid, + adc_path=adc_path, + credentials_path=credentials_path) + + dockerfile += _dependency_entries(workdir, + uid, + gid, + requirements_path=requirements_path, + conda_env_path=conda_env_path, + setup_extras=setup_extras) + + if inject_notebook.value != 'none': + install_lab = inject_notebook == NotebookInstall.lab + dockerfile += _notebook_entries(lab=install_lab, version=jupyter_version) + + if extra_dirs is not None: + dockerfile += _extra_dir_entries(workdir, uid, gid, extra_dirs) + + dockerfile += _custom_packages(uid, + gid, + packages=c.apt_packages( + caliban_config, job_mode), + shell=shell) + + if package is not None: + # The actual entrypoint and final copied code. + dockerfile += _package_entries(workdir, uid, gid, package) + + return dockerfile + + +def docker_image_id(output: str) -> ImageId: + """Accepts a string containing the output of a successful `docker build` + command and parses the Docker image ID from the stream. + + NOTE this is probably quite brittle! I can imagine this breaking quite easily + on a Docker upgrade. + + """ + return ImageId(output.splitlines()[-1].split()[-1]) + + +def build_image(job_mode: c.JobMode, + build_path: str, + credentials_path: Optional[str] = None, + adc_path: Optional[str] = None, + no_cache: bool = False, + **kwargs) -> str: + """Builds a Docker image by generating a Dockerfile and passing it to `docker + build` via stdin. All output from the `docker build` process prints to + stdout. + + Returns the image ID of the new docker container; if the command fails, + throws on error with information about the command and any issues that caused + the problem. + + """ + with u.TempCopy(credentials_path, + tmp_name=".caliban_default_creds.json") as creds: + with u.TempCopy(adc_path, tmp_name=".caliban_adc_creds.json") as adc: + cache_args = ["--no-cache"] if no_cache else [] + cmd = ["docker", "build"] + cache_args + ["--rm", "-f-", build_path] + + dockerfile = _dockerfile_template(job_mode, + credentials_path=creds, + adc_path=adc, + **kwargs) + + joined_cmd = " ".join(cmd) + logging.info("Running command: {}".format(joined_cmd)) + + try: + output, ret_code = u.capture_stdout(cmd, input_str=dockerfile) + if ret_code == 0: + return docker_image_id(output) + else: + error_msg = "Docker failed with error code {}.".format(ret_code) + raise DockerError(error_msg, cmd, ret_code) + + except subprocess.CalledProcessError as e: + logging.error(e.output) + logging.error(e.stderr) diff --git a/caliban/docker/push.py b/caliban/docker/push.py new file mode 100644 index 0000000..298aa1f --- /dev/null +++ b/caliban/docker/push.py @@ -0,0 +1,734 @@ +#!/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. +"""Functions required to interact with Docker to build and run images, shells +and notebooks in a Docker environment. + +""" + +from __future__ import absolute_import, division, print_function + +import json +import os +import subprocess +import sys +from enum import Enum +from pathlib import Path +from typing import (Any, Callable, Dict, Iterable, List, NamedTuple, NewType, + Optional, Union) + +import tqdm +from absl import logging +from blessings import Terminal +from tqdm.utils import _screen_shape_wrapper + +import caliban.config as c +import caliban.util as u +from caliban.history.types import Experiment, Job, JobSpec, JobStatus, Platform +from caliban.history.utils import (create_experiments, generate_container_spec, + get_mem_engine, get_sql_engine, + session_scope) + +t = Terminal() + +DEV_CONTAINER_ROOT = "gcr.io/blueshift-playground/blueshift" +TF_VERSIONS = {"2.2.0", "1.12.3", "1.14.0", "1.15.0"} +DEFAULT_WORKDIR = "/usr/app" +CREDS_DIR = "/.creds" +CONDA_BIN = "/opt/conda/bin/conda" + +ImageId = NewType('ImageId', str) +ArgSeq = NewType('ArgSeq', List[str]) + + +class DockerError(Exception): + """Exception that passes info on a failed Docker command.""" + + def __init__(self, message, cmd, ret_code): + super().__init__(message) + self.message = message + self.cmd = cmd + self.ret_code = ret_code + + @property + def command(self): + return " ".join(self.cmd) + + +class NotebookInstall(Enum): + """Flag to decide what to do .""" + none = 'none' + lab = 'lab' + jupyter = 'jupyter' + + def __str__(self) -> str: + return self.value + + +class Shell(Enum): + """Add new shells here and below, in SHELL_DICT.""" + bash = 'bash' + zsh = 'zsh' + + def __str__(self) -> str: + return self.value + + +# Tuple to track the information required to install and execute some custom +# shell into a container. +ShellData = NamedTuple("ShellData", [("executable", str), + ("packages", List[str])]) + + +def apt_install(*packages: str) -> str: + """Returns a command that will install the supplied list of packages without + requiring confirmation or any user interaction. + """ + package_str = ' '.join(packages) + no_prompt = "DEBIAN_FRONTEND=noninteractive" + return f"{no_prompt} apt-get install --yes --no-install-recommends {package_str}" + + +def apt_command(commands: List[str]) -> List[str]: + """Pre-and-ap-pends the supplied commands with the appropriate in-container and + cleanup command for aptitude. + + """ + update = ["apt-get update"] + cleanup = ["apt-get clean", "rm -rf /var/lib/apt/lists/*"] + return update + commands + cleanup + + +# Dict linking a particular supported shell to the data required to run and +# install the shell inside a container. +# +# : Dict[Shell, ShellData] +SHELL_DICT = { + Shell.bash: ShellData("/bin/bash", []), + Shell.zsh: ShellData("/bin/zsh", ["zsh"]) +} + + +def default_shell() -> Shell: + """Returns the shell to load into the container. Defaults to Shell.bash, but if + the user's SHELL variable refers to a supported sub-shell, returns that + instead. + + """ + ret = Shell.bash + + if "zsh" in os.environ.get("SHELL"): + ret = Shell.zsh + + return ret + + +def adc_location(home_dir: Optional[str] = None) -> str: + """Returns the location for application default credentials, INSIDE the + container (so, hardcoded unix separators), given the supplied home directory. + + """ + if home_dir is None: + home_dir = Path.home() + + return "{}/.config/gcloud/application_default_credentials.json".format( + home_dir) + + +def container_home(): + """Returns the location of the home directory inside the generated + container. + + """ + return "/home/{}".format(u.current_user()) + + +def tf_base_image(job_mode: c.JobMode, tensorflow_version: str) -> str: + """Returns the base image to use, depending on whether or not we're using a + GPU. This is JUST for building our base images for Blueshift; not for + actually using in a job. + + List of available tags: https://hub.docker.com/r/tensorflow/tensorflow/tags + + """ + if tensorflow_version not in TF_VERSIONS: + raise Exception("""{} is not a valid tensorflow version. + Try one of: {}""".format(tensorflow_version, TF_VERSIONS)) + + gpu = "-gpu" if c.gpu(job_mode) else "" + return "tensorflow/tensorflow:{}{}-py3".format(tensorflow_version, gpu) + + +def base_image_suffix(job_mode: c.JobMode) -> str: + return "gpu" if c.gpu(job_mode) else "cpu" + + +def base_image_id(job_mode: c.JobMode) -> str: + """Returns the default base image for all caliban Dockerfiles.""" + base_suffix = base_image_suffix(job_mode) + return "{}:{}".format(DEV_CONTAINER_ROOT, base_suffix) + + +def extras_string(extras: List[str]) -> str: + """Returns the argument passed to `pip install` to install a project from its + setup.py and target a specific set of extras_require dependencies. + + Args: + extras: (potentially empty) list of extra_requires deps. + """ + ret = "." + if len(extras) > 0: + ret += "[{}]".format(','.join(extras)) + return ret + + +def base_extras(job_mode: c.JobMode, path: str, + extras: Optional[List[str]]) -> Optional[List[str]]: + """Returns None if the supplied path doesn't exist (it's assumed it points to a + setup.py file). + + If the path DOES exist, generates a list of extras to install. gpu or cpu are + always added to the beginning of the list, depending on the mode. + + """ + ret = None + + if os.path.exists(path): + base = extras or [] + extra = 'gpu' if c.gpu(job_mode) else 'cpu' + ret = base if extra in base else [extra] + base + + return ret + + +def _dependency_entries(workdir: str, + user_id: int, + user_group: int, + requirements_path: Optional[str] = None, + conda_env_path: Optional[str] = None, + setup_extras: Optional[List[str]] = None) -> str: + """Returns the Dockerfile entries required to install dependencies from either: + + - a requirements.txt file, path supplied by requirements_path + - a conda environment.yml file, path supplied by conda_env_path. + - a setup.py file, if some sequence of dependencies is supplied. + + An empty list for setup_extras means, run `pip install -c .` with no extras. + None for this argument means do nothing. If a list of strings is supplied, + they'll be treated as extras dependency sets. + """ + ret = "" + + if setup_extras is not None: + ret += f""" +COPY --chown={user_id}:{user_group} setup.py {workdir} +RUN /bin/bash -c "pip install --no-cache-dir {extras_string(setup_extras)}" +""" + + if conda_env_path is not None: + ret += f""" +COPY --chown={user_id}:{user_group} {conda_env_path} {workdir} +RUN /bin/bash -c "{CONDA_BIN} env update \ + --quiet --name caliban \ + --file {conda_env_path} && \ + {CONDA_BIN} clean -y -q --all" +""" + + if requirements_path is not None: + ret += f""" +COPY --chown={user_id}:{user_group} {requirements_path} {workdir} +RUN /bin/bash -c "pip install --no-cache-dir -r {requirements_path}" +""" + + return ret + + +def _package_entries(workdir: str, user_id: int, user_group: int, + package: u.Package) -> str: + """Returns the Dockerfile entries required to: + + - copy a directory of code into a docker container + - inject an entrypoint that executes a python module inside that directory. + + Python code runs as modules vs scripts so that we can enforce import hygiene + between files inside a project. + + """ + owner = "{}:{}".format(user_id, user_group) + + arg = package.main_module or package.script_path + + # This needs to use json so that quotes print as double quotes, not single + # quotes. + entrypoint_s = json.dumps(package.executable + [arg]) + + return """ +# Copy project code into the docker container. +COPY --chown={owner} {package_path} {workdir}/{package_path} + +# Declare an entrypoint that actually runs the container. +ENTRYPOINT {entrypoint_s} + """.format_map({ + "owner": owner, + "package_path": package.package_path, + "workdir": workdir, + "entrypoint_s": entrypoint_s + }) + + +def _service_account_entry(user_id: int, user_group: int, credentials_path: str, + docker_credentials_dir: str, + write_adc_placeholder: bool): + """Generates the Dockerfile entries required to transfer a set of Cloud service + account credentials into the Docker container. + + NOTE the write_adc_placeholder variable is here because the "ctpu" script + that we use to interact with TPUs has a bug in it, as of 1/21/2020, where the + script will fail if the application_default_credentials.json file isn't + present, EVEN THOUGH it properly uses the service account credentials + registered with gcloud instead of ADC creds. + + If a service account is present, we write a placeholder string to get past + this problem. This shouldn't matter for anyone else since adc isn't used if a + service account is present. + + """ + container_creds = "{}/credentials.json".format(docker_credentials_dir) + ret = """ +COPY --chown={user_id}:{user_group} {credentials_path} {container_creds} + +# Use the credentials file to activate gcloud, gsutil inside the container. +RUN gcloud auth activate-service-account --key-file={container_creds} && \ + git config --global credential.'https://source.developers.google.com'.helper gcloud.sh + +ENV GOOGLE_APPLICATION_CREDENTIALS={container_creds} +""".format_map({ + "user_id": user_id, + "user_group": user_group, + "credentials_path": credentials_path, + "container_creds": container_creds + }) + + if write_adc_placeholder: + ret += """ +RUN echo "placeholder" >> {} +""".format(adc_location(container_home())) + + return ret + + +def _adc_entry(user_id: int, user_group: int, adc_path: str): + """Returns the Dockerfile line required to transfer the + application_default_credentials.json file into the container's home + directory. + + """ + return """ +COPY --chown={user_id}:{user_group} {adc_path} {adc_loc} + """.format_map({ + "user_id": user_id, + "user_group": user_group, + "adc_path": adc_path, + "adc_loc": adc_location(container_home()) + }) + + +def _credentials_entries(user_id: int, + user_group: int, + adc_path: Optional[str], + credentials_path: Optional[str], + docker_credentials_dir: Optional[str] = None) -> str: + """Returns the Dockerfile entries necessary to copy a user's Cloud credentials + into the Docker container. + + - adc_path is the relative path inside the current directory to an + application_default_credentials.json file containing... well, you get it. + - credentials_path is the relative path inside the current directory to a + JSON credentials file. + - docker_credentials_dir is the relative path inside the docker container + where the JSON file will be copied on build. + + """ + if docker_credentials_dir is None: + docker_credentials_dir = CREDS_DIR + + ret = "" + if credentials_path is not None: + ret += _service_account_entry(user_id, + user_group, + credentials_path, + docker_credentials_dir, + write_adc_placeholder=adc_path is None) + + if adc_path is not None: + ret += _adc_entry(user_id, user_group, adc_path) + + return ret + + +def _notebook_entries(lab: bool = False, version: Optional[str] = None) -> str: + """Returns the Dockerfile entries necessary to install Jupyter{lab}. + + Optionally takes a version string. + + """ + version_suffix = "" + + if version is not None: + version_suffix = "=={}".format(version) + + library = "jupyterlab" if lab else "jupyter" + + return """ +RUN pip install {}{} +""".format(library, version_suffix) + + +def _custom_packages( + user_id: int, + user_group: int, + packages: Optional[List[str]] = None, + shell: Optional[Shell] = None, +) -> str: + """Returns the Dockerfile entries necessary to install custom dependencies for + the supplied shell and sequence of aptitude packages. + + """ + if packages is None: + packages = [] + + if shell is None: + shell = Shell.bash + + ret = "" + + to_install = sorted(packages + SHELL_DICT[shell].packages) + + if len(to_install) != 0: + commands = apt_command([apt_install(*to_install)]) + ret = """ +USER root + +RUN {commands} + +USER {user_id}:{user_group} +""".format_map({ + "commands": " && ".join(commands), + "user_id": user_id, + "user_group": user_group + }) + + return ret + + +def _copy_dir_entry(workdir: str, user_id: int, user_group: int, + dirname: str) -> str: + """Returns the Dockerfile entry necessary to copy a single extra subdirectory + from the current directory into a docker container during build. + + """ + owner = "{}:{}".format(user_id, user_group) + return """# Copy {dirname} into the Docker container. +COPY --chown={owner} {dirname} {workdir}/{dirname} +""".format_map({ + "owner": owner, + "workdir": workdir, + "dirname": dirname + }) + + +def _extra_dir_entries(workdir: str, user_id: int, user_group: int, + extra_dirs: List[str]) -> str: + """Returns the Dockerfile entries necessary to copy all directories in the + extra_dirs list into a docker container during build. + + """ + ret = "" + for d in extra_dirs: + ret += "\n{}".format(_copy_dir_entry(workdir, user_id, user_group, d)) + return ret + + +def _dockerfile_template( + job_mode: c.JobMode, + workdir: Optional[str] = None, + base_image_fn: Optional[Callable[[c.JobMode], str]] = None, + package: Optional[Union[List, u.Package]] = None, + requirements_path: Optional[str] = None, + conda_env_path: Optional[str] = None, + setup_extras: Optional[List[str]] = None, + adc_path: Optional[str] = None, + credentials_path: Optional[str] = None, + jupyter_version: Optional[str] = None, + inject_notebook: NotebookInstall = NotebookInstall.none, + shell: Optional[Shell] = None, + extra_dirs: Optional[List[str]] = None, + caliban_config: Optional[Dict[str, Any]] = None) -> str: + """Returns a Dockerfile that builds on a local CPU or GPU base image (depending + on the value of job_mode) to create a container that: + + - installs any dependency specified in a requirements.txt file living at + requirements_path, a conda environment at conda_env_path, or any + dependencies in a setup.py file, including extra dependencies, if + setup_extras isn't None + - injects gcloud credentials into the container, so Cloud interaction works + just like it does locally + - potentially installs a custom shell, or jupyterlab for notebook support + - copies all source needed by the main module specified by package, and + potentially injects an entrypoint that, on run, will run that main module + + Most functions that call _dockerfile_template pass along any kwargs that they + receive. It should be enough to add kwargs here, then rely on that mechanism + to pass them along, vs adding kwargs all the way down the call chain. + + Supply a custom base_image_fn (function from job_mode -> image ID) to inject + more complex Docker commands into the Caliban environments by, for example, + building your own image on top of the TF base images, then using that. + + """ + uid = os.getuid() + gid = os.getgid() + username = u.current_user() + + if isinstance(package, list): + package = u.Package(*package) + + if workdir is None: + workdir = DEFAULT_WORKDIR + + if base_image_fn is None: + base_image_fn = base_image_id + + base_image = base_image_fn(job_mode) + + dockerfile = """ +FROM {base_image} + +# Create the same group we're using on the host machine. +RUN [ $(getent group {gid}) ] || groupadd --gid {gid} {gid} + +# Create the user by name. --no-log-init guards against a crash with large user +# IDs. +RUN useradd --no-log-init --no-create-home -u {uid} -g {gid} --shell /bin/bash {username} + +# The directory is created by root. This sets permissions so that any user can +# access the folder. +RUN mkdir -m 777 {workdir} {creds_dir} {c_home} + +ENV HOME={c_home} + +WORKDIR {workdir} + +USER {uid}:{gid} +""".format_map({ + "base_image": base_image, + "username": username, + "uid": uid, + "gid": gid, + "workdir": workdir, + "c_home": container_home(), + "creds_dir": CREDS_DIR + }) + dockerfile += _credentials_entries(uid, + gid, + adc_path=adc_path, + credentials_path=credentials_path) + + dockerfile += _dependency_entries(workdir, + uid, + gid, + requirements_path=requirements_path, + conda_env_path=conda_env_path, + setup_extras=setup_extras) + + if inject_notebook.value != 'none': + install_lab = inject_notebook == NotebookInstall.lab + dockerfile += _notebook_entries(lab=install_lab, version=jupyter_version) + + if extra_dirs is not None: + dockerfile += _extra_dir_entries(workdir, uid, gid, extra_dirs) + + dockerfile += _custom_packages(uid, + gid, + packages=c.apt_packages( + caliban_config, job_mode), + shell=shell) + + if package is not None: + # The actual entrypoint and final copied code. + dockerfile += _package_entries(workdir, uid, gid, package) + + return dockerfile + + +def docker_image_id(output: str) -> ImageId: + """Accepts a string containing the output of a successful `docker build` + command and parses the Docker image ID from the stream. + + NOTE this is probably quite brittle! I can imagine this breaking quite easily + on a Docker upgrade. + + """ + return ImageId(output.splitlines()[-1].split()[-1]) + + +def build_image(job_mode: c.JobMode, + build_path: str, + credentials_path: Optional[str] = None, + adc_path: Optional[str] = None, + no_cache: bool = False, + **kwargs) -> str: + """Builds a Docker image by generating a Dockerfile and passing it to `docker + build` via stdin. All output from the `docker build` process prints to + stdout. + + Returns the image ID of the new docker container; if the command fails, + throws on error with information about the command and any issues that caused + the problem. + + """ + with u.TempCopy(credentials_path, + tmp_name=".caliban_default_creds.json") as creds: + with u.TempCopy(adc_path, tmp_name=".caliban_adc_creds.json") as adc: + cache_args = ["--no-cache"] if no_cache else [] + cmd = ["docker", "build"] + cache_args + ["--rm", "-f-", build_path] + + dockerfile = _dockerfile_template(job_mode, + credentials_path=creds, + adc_path=adc, + **kwargs) + + joined_cmd = " ".join(cmd) + logging.info("Running command: {}".format(joined_cmd)) + + try: + output, ret_code = u.capture_stdout(cmd, input_str=dockerfile) + if ret_code == 0: + return docker_image_id(output) + else: + error_msg = "Docker failed with error code {}.".format(ret_code) + raise DockerError(error_msg, cmd, ret_code) + + except subprocess.CalledProcessError as e: + logging.error(e.output) + logging.error(e.stderr) + + +def _image_tag_for_project(project_id: str, image_id: str) -> str: + """Generate the GCR Docker image tag for the supplied pair of project_id and + image_id. + + This function properly handles "domain scoped projects", where the project ID + contains a domain name and project ID separated by : + https://cloud.google.com/container-registry/docs/overview#domain-scoped_projects. + + """ + project_s = project_id.replace(":", "/") + return "gcr.io/{}/{}:latest".format(project_s, image_id) + + +def push_uuid_tag(project_id: str, image_id: str) -> str: + """Takes a base image and tags it for upload, then pushes it to a remote Google + Container Registry. + + Returns the tag on a successful push. + + TODO should this just check first before attempting to push if the image + exists? Immutable names means that if the tag is up there, we're done. + Potentially use docker-py for this. + + """ + image_tag = _image_tag_for_project(project_id, image_id) + subprocess.run(["docker", "tag", image_id, image_tag], check=True) + subprocess.run(["docker", "push", image_tag], check=True) + return image_tag + + +def _run_cmd(job_mode: c.JobMode, + run_args: Optional[List[str]] = None) -> List[str]: + """Returns the sequence of commands for the subprocess run functions required + to execute `docker run`. in CPU or GPU mode, depending on the value of + job_mode. + + Keyword args: + - run_args: list of args to pass to docker run. + + """ + if run_args is None: + run_args = [] + + runtime = ["--runtime", "nvidia"] if c.gpu(job_mode) else [] + return ["docker", "run"] + runtime + ["--ipc", "host"] + run_args + + +def _home_mount_cmds(enable_home_mount: bool) -> List[str]: + """Returns the argument needed by Docker to mount a user's local home directory + into the home directory location inside their container. + + If enable_home_mount is false returns an empty list. + + """ + ret = [] + if enable_home_mount: + ret = ["-v", "{}:{}".format(Path.home(), container_home())] + return ret + + +def _interactive_opts(workdir: str) -> List[str]: + """Returns the basic arguments we want to run a docker process locally. + + """ + return [ + "-w", workdir, \ + "-u", "{}:{}".format(os.getuid(), os.getgid()), \ + "-v", "{}:{}".format(os.getcwd(), workdir) \ + ] + + +def log_job_spec_instance(job_spec: JobSpec, i: int) -> JobSpec: + """Prints logging as a side effect for the supplied sequence of job specs + generated from an experiment definition; returns the input job spec. + + """ + args = c.experiment_to_args(job_spec.experiment.kwargs, + job_spec.experiment.args) + logging.info("") + logging.info("Job {} - Experiment args: {}".format(i, t.yellow(str(args)))) + return job_spec + + +def logged_job_specs(job_specs: Iterable[JobSpec]) -> Iterable[JobSpec]: + """Generates an iterable of job specs that should be passed to `docker run` to + execute the experiments defined by the supplied iterable. + + """ + for i, s in enumerate(job_specs, 1): + yield log_job_spec_instance(s, i) + + +def execute_dry_run(job_specs: Iterable[JobSpec]) -> None: + """Expands the supplied sequence of experiments into sequences of args and logs + the jobs that WOULD have been executed, had the dry run flag not been + applied. + + """ + list(logged_job_specs(job_specs)) + + logging.info('') + logging.info( + t.yellow("To build your image and execute these jobs, \ +run your command again without {}.".format(c.DRY_RUN_FLAG))) + logging.info('') + return None diff --git a/caliban/expansion.py b/caliban/expansion.py index ebdeeb4..7659aba 100644 --- a/caliban/expansion.py +++ b/caliban/expansion.py @@ -24,7 +24,7 @@ from absl import app, logging from absl.flags import argparse_flags -import caliban.config as c +import caliban.config.experiment as c from caliban import __version__ ll.getLogger('caliban.expansion').setLevel(logging.ERROR) diff --git a/caliban/gke/__init__.py b/caliban/platform/cloud/__init__.py similarity index 100% rename from caliban/gke/__init__.py rename to caliban/platform/cloud/__init__.py diff --git a/caliban/cloud/core.py b/caliban/platform/cloud/core.py similarity index 100% rename from caliban/cloud/core.py rename to caliban/platform/cloud/core.py diff --git a/caliban/cloud/types.py b/caliban/platform/cloud/types.py similarity index 100% rename from caliban/cloud/types.py rename to caliban/platform/cloud/types.py diff --git a/caliban/platform/cloud/util.py b/caliban/platform/cloud/util.py new file mode 100644 index 0000000..c9ad813 --- /dev/null +++ b/caliban/platform/cloud/util.py @@ -0,0 +1,119 @@ +#!/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. +""" +Utilities relevant to AI Platform. +""" +import re +from typing import Dict, List, Optional, Tuple, Union + +import caliban.util as u +import caliban.util.argparse as ua + +# key and value for labels can be at most this-many-characters long. +AI_PLATFORM_MAX_LABEL_LENGTH = 63 + + +def _truncate(s: str, max_length: int) -> str: + """Returns the input string s truncated to be at most max_length characters + long. + + """ + return s if len(s) <= max_length else s[0:max_length] + + +def _clean_label(s: Optional[str], is_key: bool) -> str: + """Processes the string into the sanitized format required by AI platform + labels. + + https://cloud.google.com/ml-engine/docs/resource-labels + + """ + if s is None: + return "" + + # periods are not allowed by AI Platform labels, but often occur in, + # e.g., learning rates + DECIMAL_REPLACEMENT = '_' + s = s.replace('.', DECIMAL_REPLACEMENT) + + # lowercase, letters, - and _ are valid, so strip the leading dashes, make + # everything lowercase and then kill any remaining unallowed characters. + cleaned = re.sub(r'[^a-z0-9_-]', '', s.lower()).lstrip("-") + + # Keys must start with a letter. If is_key is set and the cleaned version + # starts with something else, append `k`. + if is_key and cleaned != "" and not cleaned[0].isalpha(): + cleaned = "k" + cleaned + + return _truncate(cleaned, AI_PLATFORM_MAX_LABEL_LENGTH) + + +def key_label(k: Optional[str]) -> str: + """converts the argument into a valid label, suitable for submission as a label + key to Cloud. + + """ + return _clean_label(k, True) + + +def value_label(v: Optional[str]) -> str: + """converts the argument into a valid label, suitable for submission as a label + value to Cloud. + + """ + return _clean_label(v, False) + + +def script_args_to_labels(script_args: Optional[List[str]]) -> Dict[str, str]: + """Converts the arguments supplied to our scripts into a dictionary usable as + labels valid for Cloud submission. + + """ + ret = {} + + def process_pair(k, v): + if ua.is_key(k): + clean_k = key_label(k) + if clean_k != "": + ret[clean_k] = "" if ua.is_key(v) else value_label(v) + + if script_args is None or len(script_args) == 0: + return ret + + elif len(script_args) == 1: + process_pair(script_args[0], None) + + # Handle the case where the final argument in the list is a boolean flag. + # This won't get picked up by partition. + elif len(script_args) > 1: + for k, v in u.partition(script_args, 2): + process_pair(k, v) + + process_pair(script_args[-1], None) + + return ret + + +def sanitize_labels( + pairs: Union[Dict[str, str], List[Tuple[str, str]]]) -> Dict[str, str]: + """Turns a dict, or a list of unsanitized key-value pairs (each represented by + a tuple) into a dictionary suitable to submit to Cloud as a label dict. + + """ + if isinstance(pairs, dict): + return sanitize_labels(pairs.items()) + + return {key_label(k): value_label(v) for (k, v) in pairs if key_label(k)} diff --git a/caliban/platform/gke/__init__.py b/caliban/platform/gke/__init__.py new file mode 100644 index 0000000..79c6a2f --- /dev/null +++ b/caliban/platform/gke/__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/gke/cli.py b/caliban/platform/gke/cli.py similarity index 100% rename from caliban/gke/cli.py rename to caliban/platform/gke/cli.py diff --git a/caliban/gke/cluster.py b/caliban/platform/gke/cluster.py similarity index 100% rename from caliban/gke/cluster.py rename to caliban/platform/gke/cluster.py diff --git a/caliban/gke/constants.py b/caliban/platform/gke/constants.py similarity index 100% rename from caliban/gke/constants.py rename to caliban/platform/gke/constants.py diff --git a/caliban/gke/types.py b/caliban/platform/gke/types.py similarity index 100% rename from caliban/gke/types.py rename to caliban/platform/gke/types.py diff --git a/caliban/gke/utils.py b/caliban/platform/gke/utils.py similarity index 100% rename from caliban/gke/utils.py rename to caliban/platform/gke/utils.py diff --git a/caliban/docker.py b/caliban/platform/notebook.py similarity index 100% rename from caliban/docker.py rename to caliban/platform/notebook.py diff --git a/caliban/platform/run.py b/caliban/platform/run.py new file mode 100644 index 0000000..88c4d02 --- /dev/null +++ b/caliban/platform/run.py @@ -0,0 +1,1056 @@ +#!/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. +"""Functions required to interact with Docker to build and run images, shells +and notebooks in a Docker environment. + +""" + +from __future__ import absolute_import, division, print_function + +import json +import os +import subprocess +import sys +from enum import Enum +from pathlib import Path +from typing import (Any, Callable, Dict, Iterable, List, NamedTuple, NewType, + Optional, Union) + +import tqdm +from absl import logging +from blessings import Terminal +from tqdm.utils import _screen_shape_wrapper + +import caliban.config as c +import caliban.util as u +from caliban.history.types import Experiment, Job, JobSpec, JobStatus, Platform +from caliban.history.utils import (create_experiments, generate_container_spec, + get_mem_engine, get_sql_engine, + session_scope) + +t = Terminal() + +DEV_CONTAINER_ROOT = "gcr.io/blueshift-playground/blueshift" +TF_VERSIONS = {"2.2.0", "1.12.3", "1.14.0", "1.15.0"} +DEFAULT_WORKDIR = "/usr/app" +CREDS_DIR = "/.creds" +CONDA_BIN = "/opt/conda/bin/conda" + +ImageId = NewType('ImageId', str) +ArgSeq = NewType('ArgSeq', List[str]) + + +class DockerError(Exception): + """Exception that passes info on a failed Docker command.""" + + def __init__(self, message, cmd, ret_code): + super().__init__(message) + self.message = message + self.cmd = cmd + self.ret_code = ret_code + + @property + def command(self): + return " ".join(self.cmd) + + +class NotebookInstall(Enum): + """Flag to decide what to do .""" + none = 'none' + lab = 'lab' + jupyter = 'jupyter' + + def __str__(self) -> str: + return self.value + + +class Shell(Enum): + """Add new shells here and below, in SHELL_DICT.""" + bash = 'bash' + zsh = 'zsh' + + def __str__(self) -> str: + return self.value + + +# Tuple to track the information required to install and execute some custom +# shell into a container. +ShellData = NamedTuple("ShellData", [("executable", str), + ("packages", List[str])]) + + +def apt_install(*packages: str) -> str: + """Returns a command that will install the supplied list of packages without + requiring confirmation or any user interaction. + """ + package_str = ' '.join(packages) + no_prompt = "DEBIAN_FRONTEND=noninteractive" + return f"{no_prompt} apt-get install --yes --no-install-recommends {package_str}" + + +def apt_command(commands: List[str]) -> List[str]: + """Pre-and-ap-pends the supplied commands with the appropriate in-container and + cleanup command for aptitude. + + """ + update = ["apt-get update"] + cleanup = ["apt-get clean", "rm -rf /var/lib/apt/lists/*"] + return update + commands + cleanup + + +# Dict linking a particular supported shell to the data required to run and +# install the shell inside a container. +# +# : Dict[Shell, ShellData] +SHELL_DICT = { + Shell.bash: ShellData("/bin/bash", []), + Shell.zsh: ShellData("/bin/zsh", ["zsh"]) +} + + +def default_shell() -> Shell: + """Returns the shell to load into the container. Defaults to Shell.bash, but if + the user's SHELL variable refers to a supported sub-shell, returns that + instead. + + """ + ret = Shell.bash + + if "zsh" in os.environ.get("SHELL"): + ret = Shell.zsh + + return ret + + +def adc_location(home_dir: Optional[str] = None) -> str: + """Returns the location for application default credentials, INSIDE the + container (so, hardcoded unix separators), given the supplied home directory. + + """ + if home_dir is None: + home_dir = Path.home() + + return "{}/.config/gcloud/application_default_credentials.json".format( + home_dir) + + +def container_home(): + """Returns the location of the home directory inside the generated + container. + + """ + return "/home/{}".format(u.current_user()) + + +def tf_base_image(job_mode: c.JobMode, tensorflow_version: str) -> str: + """Returns the base image to use, depending on whether or not we're using a + GPU. This is JUST for building our base images for Blueshift; not for + actually using in a job. + + List of available tags: https://hub.docker.com/r/tensorflow/tensorflow/tags + + """ + if tensorflow_version not in TF_VERSIONS: + raise Exception("""{} is not a valid tensorflow version. + Try one of: {}""".format(tensorflow_version, TF_VERSIONS)) + + gpu = "-gpu" if c.gpu(job_mode) else "" + return "tensorflow/tensorflow:{}{}-py3".format(tensorflow_version, gpu) + + +def base_image_suffix(job_mode: c.JobMode) -> str: + return "gpu" if c.gpu(job_mode) else "cpu" + + +def base_image_id(job_mode: c.JobMode) -> str: + """Returns the default base image for all caliban Dockerfiles.""" + base_suffix = base_image_suffix(job_mode) + return "{}:{}".format(DEV_CONTAINER_ROOT, base_suffix) + + +def extras_string(extras: List[str]) -> str: + """Returns the argument passed to `pip install` to install a project from its + setup.py and target a specific set of extras_require dependencies. + + Args: + extras: (potentially empty) list of extra_requires deps. + """ + ret = "." + if len(extras) > 0: + ret += "[{}]".format(','.join(extras)) + return ret + + +def base_extras(job_mode: c.JobMode, path: str, + extras: Optional[List[str]]) -> Optional[List[str]]: + """Returns None if the supplied path doesn't exist (it's assumed it points to a + setup.py file). + + If the path DOES exist, generates a list of extras to install. gpu or cpu are + always added to the beginning of the list, depending on the mode. + + """ + ret = None + + if os.path.exists(path): + base = extras or [] + extra = 'gpu' if c.gpu(job_mode) else 'cpu' + ret = base if extra in base else [extra] + base + + return ret + + +def _dependency_entries(workdir: str, + user_id: int, + user_group: int, + requirements_path: Optional[str] = None, + conda_env_path: Optional[str] = None, + setup_extras: Optional[List[str]] = None) -> str: + """Returns the Dockerfile entries required to install dependencies from either: + + - a requirements.txt file, path supplied by requirements_path + - a conda environment.yml file, path supplied by conda_env_path. + - a setup.py file, if some sequence of dependencies is supplied. + + An empty list for setup_extras means, run `pip install -c .` with no extras. + None for this argument means do nothing. If a list of strings is supplied, + they'll be treated as extras dependency sets. + """ + ret = "" + + if setup_extras is not None: + ret += f""" +COPY --chown={user_id}:{user_group} setup.py {workdir} +RUN /bin/bash -c "pip install --no-cache-dir {extras_string(setup_extras)}" +""" + + if conda_env_path is not None: + ret += f""" +COPY --chown={user_id}:{user_group} {conda_env_path} {workdir} +RUN /bin/bash -c "{CONDA_BIN} env update \ + --quiet --name caliban \ + --file {conda_env_path} && \ + {CONDA_BIN} clean -y -q --all" +""" + + if requirements_path is not None: + ret += f""" +COPY --chown={user_id}:{user_group} {requirements_path} {workdir} +RUN /bin/bash -c "pip install --no-cache-dir -r {requirements_path}" +""" + + return ret + + +def _package_entries(workdir: str, user_id: int, user_group: int, + package: u.Package) -> str: + """Returns the Dockerfile entries required to: + + - copy a directory of code into a docker container + - inject an entrypoint that executes a python module inside that directory. + + Python code runs as modules vs scripts so that we can enforce import hygiene + between files inside a project. + + """ + owner = "{}:{}".format(user_id, user_group) + + arg = package.main_module or package.script_path + + # This needs to use json so that quotes print as double quotes, not single + # quotes. + entrypoint_s = json.dumps(package.executable + [arg]) + + return """ +# Copy project code into the docker container. +COPY --chown={owner} {package_path} {workdir}/{package_path} + +# Declare an entrypoint that actually runs the container. +ENTRYPOINT {entrypoint_s} + """.format_map({ + "owner": owner, + "package_path": package.package_path, + "workdir": workdir, + "entrypoint_s": entrypoint_s + }) + + +def _service_account_entry(user_id: int, user_group: int, credentials_path: str, + docker_credentials_dir: str, + write_adc_placeholder: bool): + """Generates the Dockerfile entries required to transfer a set of Cloud service + account credentials into the Docker container. + + NOTE the write_adc_placeholder variable is here because the "ctpu" script + that we use to interact with TPUs has a bug in it, as of 1/21/2020, where the + script will fail if the application_default_credentials.json file isn't + present, EVEN THOUGH it properly uses the service account credentials + registered with gcloud instead of ADC creds. + + If a service account is present, we write a placeholder string to get past + this problem. This shouldn't matter for anyone else since adc isn't used if a + service account is present. + + """ + container_creds = "{}/credentials.json".format(docker_credentials_dir) + ret = """ +COPY --chown={user_id}:{user_group} {credentials_path} {container_creds} + +# Use the credentials file to activate gcloud, gsutil inside the container. +RUN gcloud auth activate-service-account --key-file={container_creds} && \ + git config --global credential.'https://source.developers.google.com'.helper gcloud.sh + +ENV GOOGLE_APPLICATION_CREDENTIALS={container_creds} +""".format_map({ + "user_id": user_id, + "user_group": user_group, + "credentials_path": credentials_path, + "container_creds": container_creds + }) + + if write_adc_placeholder: + ret += """ +RUN echo "placeholder" >> {} +""".format(adc_location(container_home())) + + return ret + + +def _adc_entry(user_id: int, user_group: int, adc_path: str): + """Returns the Dockerfile line required to transfer the + application_default_credentials.json file into the container's home + directory. + + """ + return """ +COPY --chown={user_id}:{user_group} {adc_path} {adc_loc} + """.format_map({ + "user_id": user_id, + "user_group": user_group, + "adc_path": adc_path, + "adc_loc": adc_location(container_home()) + }) + + +def _credentials_entries(user_id: int, + user_group: int, + adc_path: Optional[str], + credentials_path: Optional[str], + docker_credentials_dir: Optional[str] = None) -> str: + """Returns the Dockerfile entries necessary to copy a user's Cloud credentials + into the Docker container. + + - adc_path is the relative path inside the current directory to an + application_default_credentials.json file containing... well, you get it. + - credentials_path is the relative path inside the current directory to a + JSON credentials file. + - docker_credentials_dir is the relative path inside the docker container + where the JSON file will be copied on build. + + """ + if docker_credentials_dir is None: + docker_credentials_dir = CREDS_DIR + + ret = "" + if credentials_path is not None: + ret += _service_account_entry(user_id, + user_group, + credentials_path, + docker_credentials_dir, + write_adc_placeholder=adc_path is None) + + if adc_path is not None: + ret += _adc_entry(user_id, user_group, adc_path) + + return ret + + +def _notebook_entries(lab: bool = False, version: Optional[str] = None) -> str: + """Returns the Dockerfile entries necessary to install Jupyter{lab}. + + Optionally takes a version string. + + """ + version_suffix = "" + + if version is not None: + version_suffix = "=={}".format(version) + + library = "jupyterlab" if lab else "jupyter" + + return """ +RUN pip install {}{} +""".format(library, version_suffix) + + +def _custom_packages( + user_id: int, + user_group: int, + packages: Optional[List[str]] = None, + shell: Optional[Shell] = None, +) -> str: + """Returns the Dockerfile entries necessary to install custom dependencies for + the supplied shell and sequence of aptitude packages. + + """ + if packages is None: + packages = [] + + if shell is None: + shell = Shell.bash + + ret = "" + + to_install = sorted(packages + SHELL_DICT[shell].packages) + + if len(to_install) != 0: + commands = apt_command([apt_install(*to_install)]) + ret = """ +USER root + +RUN {commands} + +USER {user_id}:{user_group} +""".format_map({ + "commands": " && ".join(commands), + "user_id": user_id, + "user_group": user_group + }) + + return ret + + +def _copy_dir_entry(workdir: str, user_id: int, user_group: int, + dirname: str) -> str: + """Returns the Dockerfile entry necessary to copy a single extra subdirectory + from the current directory into a docker container during build. + + """ + owner = "{}:{}".format(user_id, user_group) + return """# Copy {dirname} into the Docker container. +COPY --chown={owner} {dirname} {workdir}/{dirname} +""".format_map({ + "owner": owner, + "workdir": workdir, + "dirname": dirname + }) + + +def _extra_dir_entries(workdir: str, user_id: int, user_group: int, + extra_dirs: List[str]) -> str: + """Returns the Dockerfile entries necessary to copy all directories in the + extra_dirs list into a docker container during build. + + """ + ret = "" + for d in extra_dirs: + ret += "\n{}".format(_copy_dir_entry(workdir, user_id, user_group, d)) + return ret + + +def _dockerfile_template( + job_mode: c.JobMode, + workdir: Optional[str] = None, + base_image_fn: Optional[Callable[[c.JobMode], str]] = None, + package: Optional[Union[List, u.Package]] = None, + requirements_path: Optional[str] = None, + conda_env_path: Optional[str] = None, + setup_extras: Optional[List[str]] = None, + adc_path: Optional[str] = None, + credentials_path: Optional[str] = None, + jupyter_version: Optional[str] = None, + inject_notebook: NotebookInstall = NotebookInstall.none, + shell: Optional[Shell] = None, + extra_dirs: Optional[List[str]] = None, + caliban_config: Optional[Dict[str, Any]] = None) -> str: + """Returns a Dockerfile that builds on a local CPU or GPU base image (depending + on the value of job_mode) to create a container that: + + - installs any dependency specified in a requirements.txt file living at + requirements_path, a conda environment at conda_env_path, or any + dependencies in a setup.py file, including extra dependencies, if + setup_extras isn't None + - injects gcloud credentials into the container, so Cloud interaction works + just like it does locally + - potentially installs a custom shell, or jupyterlab for notebook support + - copies all source needed by the main module specified by package, and + potentially injects an entrypoint that, on run, will run that main module + + Most functions that call _dockerfile_template pass along any kwargs that they + receive. It should be enough to add kwargs here, then rely on that mechanism + to pass them along, vs adding kwargs all the way down the call chain. + + Supply a custom base_image_fn (function from job_mode -> image ID) to inject + more complex Docker commands into the Caliban environments by, for example, + building your own image on top of the TF base images, then using that. + + """ + uid = os.getuid() + gid = os.getgid() + username = u.current_user() + + if isinstance(package, list): + package = u.Package(*package) + + if workdir is None: + workdir = DEFAULT_WORKDIR + + if base_image_fn is None: + base_image_fn = base_image_id + + base_image = base_image_fn(job_mode) + + dockerfile = """ +FROM {base_image} + +# Create the same group we're using on the host machine. +RUN [ $(getent group {gid}) ] || groupadd --gid {gid} {gid} + +# Create the user by name. --no-log-init guards against a crash with large user +# IDs. +RUN useradd --no-log-init --no-create-home -u {uid} -g {gid} --shell /bin/bash {username} + +# The directory is created by root. This sets permissions so that any user can +# access the folder. +RUN mkdir -m 777 {workdir} {creds_dir} {c_home} + +ENV HOME={c_home} + +WORKDIR {workdir} + +USER {uid}:{gid} +""".format_map({ + "base_image": base_image, + "username": username, + "uid": uid, + "gid": gid, + "workdir": workdir, + "c_home": container_home(), + "creds_dir": CREDS_DIR + }) + dockerfile += _credentials_entries(uid, + gid, + adc_path=adc_path, + credentials_path=credentials_path) + + dockerfile += _dependency_entries(workdir, + uid, + gid, + requirements_path=requirements_path, + conda_env_path=conda_env_path, + setup_extras=setup_extras) + + if inject_notebook.value != 'none': + install_lab = inject_notebook == NotebookInstall.lab + dockerfile += _notebook_entries(lab=install_lab, version=jupyter_version) + + if extra_dirs is not None: + dockerfile += _extra_dir_entries(workdir, uid, gid, extra_dirs) + + dockerfile += _custom_packages(uid, + gid, + packages=c.apt_packages( + caliban_config, job_mode), + shell=shell) + + if package is not None: + # The actual entrypoint and final copied code. + dockerfile += _package_entries(workdir, uid, gid, package) + + return dockerfile + + +def docker_image_id(output: str) -> ImageId: + """Accepts a string containing the output of a successful `docker build` + command and parses the Docker image ID from the stream. + + NOTE this is probably quite brittle! I can imagine this breaking quite easily + on a Docker upgrade. + + """ + return ImageId(output.splitlines()[-1].split()[-1]) + + +def build_image(job_mode: c.JobMode, + build_path: str, + credentials_path: Optional[str] = None, + adc_path: Optional[str] = None, + no_cache: bool = False, + **kwargs) -> str: + """Builds a Docker image by generating a Dockerfile and passing it to `docker + build` via stdin. All output from the `docker build` process prints to + stdout. + + Returns the image ID of the new docker container; if the command fails, + throws on error with information about the command and any issues that caused + the problem. + + """ + with u.TempCopy(credentials_path, + tmp_name=".caliban_default_creds.json") as creds: + with u.TempCopy(adc_path, tmp_name=".caliban_adc_creds.json") as adc: + cache_args = ["--no-cache"] if no_cache else [] + cmd = ["docker", "build"] + cache_args + ["--rm", "-f-", build_path] + + dockerfile = _dockerfile_template(job_mode, + credentials_path=creds, + adc_path=adc, + **kwargs) + + joined_cmd = " ".join(cmd) + logging.info("Running command: {}".format(joined_cmd)) + + try: + output, ret_code = u.capture_stdout(cmd, input_str=dockerfile) + if ret_code == 0: + return docker_image_id(output) + else: + error_msg = "Docker failed with error code {}.".format(ret_code) + raise DockerError(error_msg, cmd, ret_code) + + except subprocess.CalledProcessError as e: + logging.error(e.output) + logging.error(e.stderr) + + +def _image_tag_for_project(project_id: str, image_id: str) -> str: + """Generate the GCR Docker image tag for the supplied pair of project_id and + image_id. + + This function properly handles "domain scoped projects", where the project ID + contains a domain name and project ID separated by : + https://cloud.google.com/container-registry/docs/overview#domain-scoped_projects. + + """ + project_s = project_id.replace(":", "/") + return "gcr.io/{}/{}:latest".format(project_s, image_id) + + +def push_uuid_tag(project_id: str, image_id: str) -> str: + """Takes a base image and tags it for upload, then pushes it to a remote Google + Container Registry. + + Returns the tag on a successful push. + + TODO should this just check first before attempting to push if the image + exists? Immutable names means that if the tag is up there, we're done. + Potentially use docker-py for this. + + """ + image_tag = _image_tag_for_project(project_id, image_id) + subprocess.run(["docker", "tag", image_id, image_tag], check=True) + subprocess.run(["docker", "push", image_tag], check=True) + return image_tag + + +def _run_cmd(job_mode: c.JobMode, + run_args: Optional[List[str]] = None) -> List[str]: + """Returns the sequence of commands for the subprocess run functions required + to execute `docker run`. in CPU or GPU mode, depending on the value of + job_mode. + + Keyword args: + - run_args: list of args to pass to docker run. + + """ + if run_args is None: + run_args = [] + + runtime = ["--runtime", "nvidia"] if c.gpu(job_mode) else [] + return ["docker", "run"] + runtime + ["--ipc", "host"] + run_args + + +def _home_mount_cmds(enable_home_mount: bool) -> List[str]: + """Returns the argument needed by Docker to mount a user's local home directory + into the home directory location inside their container. + + If enable_home_mount is false returns an empty list. + + """ + ret = [] + if enable_home_mount: + ret = ["-v", "{}:{}".format(Path.home(), container_home())] + return ret + + +def _interactive_opts(workdir: str) -> List[str]: + """Returns the basic arguments we want to run a docker process locally. + + """ + return [ + "-w", workdir, \ + "-u", "{}:{}".format(os.getuid(), os.getgid()), \ + "-v", "{}:{}".format(os.getcwd(), workdir) \ + ] + + +def log_job_spec_instance(job_spec: JobSpec, i: int) -> JobSpec: + """Prints logging as a side effect for the supplied sequence of job specs + generated from an experiment definition; returns the input job spec. + + """ + args = c.experiment_to_args(job_spec.experiment.kwargs, + job_spec.experiment.args) + logging.info("") + logging.info("Job {} - Experiment args: {}".format(i, t.yellow(str(args)))) + return job_spec + + +def logged_job_specs(job_specs: Iterable[JobSpec]) -> Iterable[JobSpec]: + """Generates an iterable of job specs that should be passed to `docker run` to + execute the experiments defined by the supplied iterable. + + """ + for i, s in enumerate(job_specs, 1): + yield log_job_spec_instance(s, i) + + +def execute_dry_run(job_specs: Iterable[JobSpec]) -> None: + """Expands the supplied sequence of experiments into sequences of args and logs + the jobs that WOULD have been executed, had the dry run flag not been + applied. + + """ + list(logged_job_specs(job_specs)) + + logging.info('') + logging.info( + t.yellow("To build your image and execute these jobs, \ +run your command again without {}.".format(c.DRY_RUN_FLAG))) + logging.info('') + return None + + +def local_callback(idx: int, job: Job) -> None: + """Provides logging feedback for jobs run locally. If the return code is 0, + logs success; else, logs the failure as an error and logs the script args + that provided the failure. + + """ + if job.status == JobStatus.SUCCEEDED: + logging.info(t.green(f'Job {idx} succeeded!')) + else: + logging.error( + t.red(f'Job {idx} failed with return code {job.details["ret_code"]}.')) + args = c.experiment_to_args(job.spec.experiment.kwargs, + job.spec.experiment.args) + logging.error(t.red(f'Failing args for job {idx}: {args}')) + + +def window_size_env_cmds(): + """Returns a sequence of `docker run` arguments that will internally configure + the terminal columns and lines, so that progress bars and other terminal + interactions will work properly. + + These aren't required for interactive Docker commands like those triggered by + `caliban shell`. + + """ + ret = [] + cols, lines = _screen_shape_wrapper()(0) + if cols: + ret += ["-e", f"COLUMNS={cols}"] + if lines: + ret += ["-e", f"LINES={lines}"] + return ret + + +# ---------------------------------------------------------------------------- +def _create_job_spec_dict( + experiment: Experiment, + job_mode: c.JobMode, + image_id: str, + run_args: Optional[List[str]] = None, +) -> Dict[str, Any]: + '''creates a job spec dictionary for a local job''' + + # Without the unbuffered environment variable, stderr and stdout won't be + # emitted in the proper order from inside the container. + terminal_cmds = ["-e" "PYTHONUNBUFFERED=1"] + window_size_env_cmds() + + base_cmd = _run_cmd(job_mode, run_args) + terminal_cmds + [image_id] + command = base_cmd + c.experiment_to_args(experiment.kwargs, experiment.args) + return {'command': command, 'container': image_id} + + +# ---------------------------------------------------------------------------- +def execute_jobs( + job_specs: Iterable[JobSpec], + dry_run: bool = False, +): + '''executes a sequence of jobs based on job specs + + Arg: + job_specs: specifications for jobs to be executed + dry_run: if True, only print what would be done + ''' + + with u.tqdm_logging() as orig_stream: + pbar = tqdm.tqdm(logged_job_specs(job_specs), + file=orig_stream, + total=len(job_specs), + ascii=True, + unit="experiment", + desc="Executing") + for idx, job_spec in enumerate(pbar, 1): + command = job_spec.spec['command'] + logging.info(f'Running command: {" ".join(command)}') + if not dry_run: + _, ret_code = u.capture_stdout(command, "", u.TqdmFile(sys.stderr)) + else: + ret_code = 0 + j = Job(spec=job_spec, + container=job_spec.spec['container'], + details={'ret_code': ret_code}, + status=JobStatus.SUCCEEDED if ret_code == 0 else JobStatus.FAILED) + local_callback(idx=idx, job=j) + + if dry_run: + logging.info( + t.yellow(f'\nTo build your image and execute these jobs, ' + f'run your command again without {c.DRY_RUN_FLAG}\n')) + + return None + + +def run_experiments(job_mode: c.JobMode, + run_args: Optional[List[str]] = None, + script_args: Optional[List[str]] = None, + image_id: Optional[str] = None, + dry_run: bool = False, + experiment_config: Optional[c.ExpConf] = None, + xgroup: Optional[str] = None, + **build_image_kwargs) -> None: + """Builds an image using the supplied **build_image_kwargs and calls `docker + run` on the resulting image using sensible defaults. + + Keyword args: + + - job_mode: c.JobMode. + + - run_args: extra arguments to supply to `docker run` after our defaults. + - script_args: extra arguments to supply to the entrypoint. (You can + - override the default container entrypoint by supplying a new one inside + run_args.) + - image_id: ID of the image to run. Supplying this will skip an image build. + - experiment_config: dict of string to list, boolean, string or int. Any + lists will trigger a cartesian product out with the rest of the config. A + job will be executed for every combination of parameters in the experiment + config. + - dry_run: if True, no actual jobs will be executed and docker won't + actually build; logging side effects will show the user what will happen + without dry_run=True. + + any extra kwargs supplied are passed through to build_image. + """ + if run_args is None: + run_args = [] + + if script_args is None: + script_args = [] + + if experiment_config is None: + experiment_config = {} + + docker_args = {k: v for k, v in build_image_kwargs.items()} + docker_args['job_mode'] = job_mode + + engine = get_mem_engine() if dry_run else get_sql_engine() + + with session_scope(engine) as session: + container_spec = generate_container_spec(session, docker_args, image_id) + + if image_id is None: + if dry_run: + logging.info("Dry run - skipping actual 'docker build'.") + image_id = 'dry_run_tag' + else: + image_id = build_image(**docker_args) + + experiments = create_experiments( + session=session, + container_spec=container_spec, + script_args=script_args, + experiment_config=experiment_config, + xgroup=xgroup, + ) + + job_specs = [ + JobSpec.get_or_create( + experiment=x, + spec=_create_job_spec_dict( + experiment=x, + job_mode=job_mode, + run_args=run_args, + image_id=image_id, + ), + platform=Platform.LOCAL, + ) for x in experiments + ] + + try: + execute_jobs(job_specs=job_specs, dry_run=dry_run) + except Exception as e: + logging.error(f'exception: {e}') + session.commit() # commit here, otherwise will be rolled back + + +def run(job_mode: c.JobMode, + run_args: Optional[List[str]] = None, + script_args: Optional[List[str]] = None, + image_id: Optional[str] = None, + **build_image_kwargs) -> None: + """Builds an image using the supplied **build_image_kwargs and calls `docker + run` on the resulting image using sensible defaults. + Keyword args: + - job_mode: c.JobMode. + - run_args: extra arguments to supply to `docker run` after our defaults. + - script_args: extra arguments to supply to the entrypoint. (You can + - override the default container entrypoint by supplying a new one inside + run_args.) + - image_id: ID of the image to run. Supplying this will skip an image build. + any extra kwargs supplied are passed through to build_image. + """ + if run_args is None: + run_args = [] + + if script_args is None: + script_args = [] + + if image_id is None: + image_id = build_image(job_mode, **build_image_kwargs) + + base_cmd = _run_cmd(job_mode, run_args) + + command = base_cmd + [image_id] + script_args + + logging.info("Running command: {}".format(' '.join(command))) + subprocess.call(command) + return None + + +def run_interactive(job_mode: c.JobMode, + workdir: Optional[str] = None, + image_id: Optional[str] = None, + run_args: Optional[List[str]] = None, + mount_home: Optional[bool] = None, + shell: Optional[Shell] = None, + entrypoint: Optional[str] = None, + entrypoint_args: Optional[List[str]] = None, + **build_image_kwargs) -> None: + """Start a live shell in the terminal, with all dependencies installed and the + current working directory (and optionally the user's home directory) mounted. + + Keyword args: + + - job_mode: c.JobMode. + - image_id: ID of the image to run. Supplying this will skip an image build. + - run_args: extra arguments to supply to `docker run`. + - mount_home: if true, mounts the user's $HOME directory into the container + to `/home/$USERNAME`. If False, nothing. + - shell: name of the shell to install into the container. Also configures the + entrypoint if that's not supplied. + - entrypoint: command to run. Defaults to the executable command for the + supplied shell. + - entrypoint_args: extra arguments to supply to the entrypoint. + + any extra kwargs supplied are passed through to build_image. + + """ + if workdir is None: + workdir = DEFAULT_WORKDIR + + if run_args is None: + run_args = [] + + if entrypoint_args is None: + entrypoint_args = [] + + if mount_home is None: + mount_home = True + + if shell is None: + # Only set a default shell if we're also mounting the home volume. + # Otherwise a custom shell won't have access to the user's profile. + shell = default_shell() if mount_home else Shell.bash + + if entrypoint is None: + entrypoint = SHELL_DICT[shell].executable + + interactive_run_args = _interactive_opts(workdir) + [ + "-it", \ + "--entrypoint", entrypoint + ] + _home_mount_cmds(mount_home) + run_args + + run(job_mode=job_mode, + run_args=interactive_run_args, + script_args=entrypoint_args, + image_id=image_id, + shell=shell, + workdir=workdir, + **build_image_kwargs) + + +def run_notebook(job_mode: c.JobMode, + port: Optional[int] = None, + lab: Optional[bool] = None, + version: Optional[bool] = None, + run_args: Optional[List[str]] = None, + **run_interactive_kwargs) -> None: + """Start a notebook in the current working directory; the process will run + inside of a Docker container that's identical to the environment available to + Cloud jobs that are submitted by `caliban cloud`, or local jobs run with + `caliban run.` + + if you pass mount_home=True your jupyter settings will persist across calls. + + Keyword args: + + - port: the port to pass to Jupyter when it boots, useful if you have + multiple instances running on one machine. + - lab: if True, starts jupyter lab, else jupyter notebook. + - version: explicit Jupyter version to install. + + run_interactive_kwargs are all extra arguments taken by run_interactive. + + """ + + if port is None: + port = u.next_free_port(8888) + + if lab is None: + lab = False + + if run_args is None: + run_args = [] + + inject_arg = NotebookInstall.lab if lab else NotebookInstall.jupyter + jupyter_cmd = "lab" if lab else "notebook" + jupyter_args = [ + "-m", "jupyter", jupyter_cmd, \ + "--ip=0.0.0.0", \ + "--port={}".format(port), \ + "--no-browser" + ] + docker_args = ["-p", "{}:{}".format(port, port)] + run_args + + run_interactive(job_mode, + entrypoint="/opt/conda/envs/caliban/bin/python", + entrypoint_args=jupyter_args, + run_args=docker_args, + inject_notebook=inject_arg, + jupyter_version=version, + **run_interactive_kwargs) diff --git a/caliban/platform/shell.py b/caliban/platform/shell.py new file mode 100644 index 0000000..88c4d02 --- /dev/null +++ b/caliban/platform/shell.py @@ -0,0 +1,1056 @@ +#!/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. +"""Functions required to interact with Docker to build and run images, shells +and notebooks in a Docker environment. + +""" + +from __future__ import absolute_import, division, print_function + +import json +import os +import subprocess +import sys +from enum import Enum +from pathlib import Path +from typing import (Any, Callable, Dict, Iterable, List, NamedTuple, NewType, + Optional, Union) + +import tqdm +from absl import logging +from blessings import Terminal +from tqdm.utils import _screen_shape_wrapper + +import caliban.config as c +import caliban.util as u +from caliban.history.types import Experiment, Job, JobSpec, JobStatus, Platform +from caliban.history.utils import (create_experiments, generate_container_spec, + get_mem_engine, get_sql_engine, + session_scope) + +t = Terminal() + +DEV_CONTAINER_ROOT = "gcr.io/blueshift-playground/blueshift" +TF_VERSIONS = {"2.2.0", "1.12.3", "1.14.0", "1.15.0"} +DEFAULT_WORKDIR = "/usr/app" +CREDS_DIR = "/.creds" +CONDA_BIN = "/opt/conda/bin/conda" + +ImageId = NewType('ImageId', str) +ArgSeq = NewType('ArgSeq', List[str]) + + +class DockerError(Exception): + """Exception that passes info on a failed Docker command.""" + + def __init__(self, message, cmd, ret_code): + super().__init__(message) + self.message = message + self.cmd = cmd + self.ret_code = ret_code + + @property + def command(self): + return " ".join(self.cmd) + + +class NotebookInstall(Enum): + """Flag to decide what to do .""" + none = 'none' + lab = 'lab' + jupyter = 'jupyter' + + def __str__(self) -> str: + return self.value + + +class Shell(Enum): + """Add new shells here and below, in SHELL_DICT.""" + bash = 'bash' + zsh = 'zsh' + + def __str__(self) -> str: + return self.value + + +# Tuple to track the information required to install and execute some custom +# shell into a container. +ShellData = NamedTuple("ShellData", [("executable", str), + ("packages", List[str])]) + + +def apt_install(*packages: str) -> str: + """Returns a command that will install the supplied list of packages without + requiring confirmation or any user interaction. + """ + package_str = ' '.join(packages) + no_prompt = "DEBIAN_FRONTEND=noninteractive" + return f"{no_prompt} apt-get install --yes --no-install-recommends {package_str}" + + +def apt_command(commands: List[str]) -> List[str]: + """Pre-and-ap-pends the supplied commands with the appropriate in-container and + cleanup command for aptitude. + + """ + update = ["apt-get update"] + cleanup = ["apt-get clean", "rm -rf /var/lib/apt/lists/*"] + return update + commands + cleanup + + +# Dict linking a particular supported shell to the data required to run and +# install the shell inside a container. +# +# : Dict[Shell, ShellData] +SHELL_DICT = { + Shell.bash: ShellData("/bin/bash", []), + Shell.zsh: ShellData("/bin/zsh", ["zsh"]) +} + + +def default_shell() -> Shell: + """Returns the shell to load into the container. Defaults to Shell.bash, but if + the user's SHELL variable refers to a supported sub-shell, returns that + instead. + + """ + ret = Shell.bash + + if "zsh" in os.environ.get("SHELL"): + ret = Shell.zsh + + return ret + + +def adc_location(home_dir: Optional[str] = None) -> str: + """Returns the location for application default credentials, INSIDE the + container (so, hardcoded unix separators), given the supplied home directory. + + """ + if home_dir is None: + home_dir = Path.home() + + return "{}/.config/gcloud/application_default_credentials.json".format( + home_dir) + + +def container_home(): + """Returns the location of the home directory inside the generated + container. + + """ + return "/home/{}".format(u.current_user()) + + +def tf_base_image(job_mode: c.JobMode, tensorflow_version: str) -> str: + """Returns the base image to use, depending on whether or not we're using a + GPU. This is JUST for building our base images for Blueshift; not for + actually using in a job. + + List of available tags: https://hub.docker.com/r/tensorflow/tensorflow/tags + + """ + if tensorflow_version not in TF_VERSIONS: + raise Exception("""{} is not a valid tensorflow version. + Try one of: {}""".format(tensorflow_version, TF_VERSIONS)) + + gpu = "-gpu" if c.gpu(job_mode) else "" + return "tensorflow/tensorflow:{}{}-py3".format(tensorflow_version, gpu) + + +def base_image_suffix(job_mode: c.JobMode) -> str: + return "gpu" if c.gpu(job_mode) else "cpu" + + +def base_image_id(job_mode: c.JobMode) -> str: + """Returns the default base image for all caliban Dockerfiles.""" + base_suffix = base_image_suffix(job_mode) + return "{}:{}".format(DEV_CONTAINER_ROOT, base_suffix) + + +def extras_string(extras: List[str]) -> str: + """Returns the argument passed to `pip install` to install a project from its + setup.py and target a specific set of extras_require dependencies. + + Args: + extras: (potentially empty) list of extra_requires deps. + """ + ret = "." + if len(extras) > 0: + ret += "[{}]".format(','.join(extras)) + return ret + + +def base_extras(job_mode: c.JobMode, path: str, + extras: Optional[List[str]]) -> Optional[List[str]]: + """Returns None if the supplied path doesn't exist (it's assumed it points to a + setup.py file). + + If the path DOES exist, generates a list of extras to install. gpu or cpu are + always added to the beginning of the list, depending on the mode. + + """ + ret = None + + if os.path.exists(path): + base = extras or [] + extra = 'gpu' if c.gpu(job_mode) else 'cpu' + ret = base if extra in base else [extra] + base + + return ret + + +def _dependency_entries(workdir: str, + user_id: int, + user_group: int, + requirements_path: Optional[str] = None, + conda_env_path: Optional[str] = None, + setup_extras: Optional[List[str]] = None) -> str: + """Returns the Dockerfile entries required to install dependencies from either: + + - a requirements.txt file, path supplied by requirements_path + - a conda environment.yml file, path supplied by conda_env_path. + - a setup.py file, if some sequence of dependencies is supplied. + + An empty list for setup_extras means, run `pip install -c .` with no extras. + None for this argument means do nothing. If a list of strings is supplied, + they'll be treated as extras dependency sets. + """ + ret = "" + + if setup_extras is not None: + ret += f""" +COPY --chown={user_id}:{user_group} setup.py {workdir} +RUN /bin/bash -c "pip install --no-cache-dir {extras_string(setup_extras)}" +""" + + if conda_env_path is not None: + ret += f""" +COPY --chown={user_id}:{user_group} {conda_env_path} {workdir} +RUN /bin/bash -c "{CONDA_BIN} env update \ + --quiet --name caliban \ + --file {conda_env_path} && \ + {CONDA_BIN} clean -y -q --all" +""" + + if requirements_path is not None: + ret += f""" +COPY --chown={user_id}:{user_group} {requirements_path} {workdir} +RUN /bin/bash -c "pip install --no-cache-dir -r {requirements_path}" +""" + + return ret + + +def _package_entries(workdir: str, user_id: int, user_group: int, + package: u.Package) -> str: + """Returns the Dockerfile entries required to: + + - copy a directory of code into a docker container + - inject an entrypoint that executes a python module inside that directory. + + Python code runs as modules vs scripts so that we can enforce import hygiene + between files inside a project. + + """ + owner = "{}:{}".format(user_id, user_group) + + arg = package.main_module or package.script_path + + # This needs to use json so that quotes print as double quotes, not single + # quotes. + entrypoint_s = json.dumps(package.executable + [arg]) + + return """ +# Copy project code into the docker container. +COPY --chown={owner} {package_path} {workdir}/{package_path} + +# Declare an entrypoint that actually runs the container. +ENTRYPOINT {entrypoint_s} + """.format_map({ + "owner": owner, + "package_path": package.package_path, + "workdir": workdir, + "entrypoint_s": entrypoint_s + }) + + +def _service_account_entry(user_id: int, user_group: int, credentials_path: str, + docker_credentials_dir: str, + write_adc_placeholder: bool): + """Generates the Dockerfile entries required to transfer a set of Cloud service + account credentials into the Docker container. + + NOTE the write_adc_placeholder variable is here because the "ctpu" script + that we use to interact with TPUs has a bug in it, as of 1/21/2020, where the + script will fail if the application_default_credentials.json file isn't + present, EVEN THOUGH it properly uses the service account credentials + registered with gcloud instead of ADC creds. + + If a service account is present, we write a placeholder string to get past + this problem. This shouldn't matter for anyone else since adc isn't used if a + service account is present. + + """ + container_creds = "{}/credentials.json".format(docker_credentials_dir) + ret = """ +COPY --chown={user_id}:{user_group} {credentials_path} {container_creds} + +# Use the credentials file to activate gcloud, gsutil inside the container. +RUN gcloud auth activate-service-account --key-file={container_creds} && \ + git config --global credential.'https://source.developers.google.com'.helper gcloud.sh + +ENV GOOGLE_APPLICATION_CREDENTIALS={container_creds} +""".format_map({ + "user_id": user_id, + "user_group": user_group, + "credentials_path": credentials_path, + "container_creds": container_creds + }) + + if write_adc_placeholder: + ret += """ +RUN echo "placeholder" >> {} +""".format(adc_location(container_home())) + + return ret + + +def _adc_entry(user_id: int, user_group: int, adc_path: str): + """Returns the Dockerfile line required to transfer the + application_default_credentials.json file into the container's home + directory. + + """ + return """ +COPY --chown={user_id}:{user_group} {adc_path} {adc_loc} + """.format_map({ + "user_id": user_id, + "user_group": user_group, + "adc_path": adc_path, + "adc_loc": adc_location(container_home()) + }) + + +def _credentials_entries(user_id: int, + user_group: int, + adc_path: Optional[str], + credentials_path: Optional[str], + docker_credentials_dir: Optional[str] = None) -> str: + """Returns the Dockerfile entries necessary to copy a user's Cloud credentials + into the Docker container. + + - adc_path is the relative path inside the current directory to an + application_default_credentials.json file containing... well, you get it. + - credentials_path is the relative path inside the current directory to a + JSON credentials file. + - docker_credentials_dir is the relative path inside the docker container + where the JSON file will be copied on build. + + """ + if docker_credentials_dir is None: + docker_credentials_dir = CREDS_DIR + + ret = "" + if credentials_path is not None: + ret += _service_account_entry(user_id, + user_group, + credentials_path, + docker_credentials_dir, + write_adc_placeholder=adc_path is None) + + if adc_path is not None: + ret += _adc_entry(user_id, user_group, adc_path) + + return ret + + +def _notebook_entries(lab: bool = False, version: Optional[str] = None) -> str: + """Returns the Dockerfile entries necessary to install Jupyter{lab}. + + Optionally takes a version string. + + """ + version_suffix = "" + + if version is not None: + version_suffix = "=={}".format(version) + + library = "jupyterlab" if lab else "jupyter" + + return """ +RUN pip install {}{} +""".format(library, version_suffix) + + +def _custom_packages( + user_id: int, + user_group: int, + packages: Optional[List[str]] = None, + shell: Optional[Shell] = None, +) -> str: + """Returns the Dockerfile entries necessary to install custom dependencies for + the supplied shell and sequence of aptitude packages. + + """ + if packages is None: + packages = [] + + if shell is None: + shell = Shell.bash + + ret = "" + + to_install = sorted(packages + SHELL_DICT[shell].packages) + + if len(to_install) != 0: + commands = apt_command([apt_install(*to_install)]) + ret = """ +USER root + +RUN {commands} + +USER {user_id}:{user_group} +""".format_map({ + "commands": " && ".join(commands), + "user_id": user_id, + "user_group": user_group + }) + + return ret + + +def _copy_dir_entry(workdir: str, user_id: int, user_group: int, + dirname: str) -> str: + """Returns the Dockerfile entry necessary to copy a single extra subdirectory + from the current directory into a docker container during build. + + """ + owner = "{}:{}".format(user_id, user_group) + return """# Copy {dirname} into the Docker container. +COPY --chown={owner} {dirname} {workdir}/{dirname} +""".format_map({ + "owner": owner, + "workdir": workdir, + "dirname": dirname + }) + + +def _extra_dir_entries(workdir: str, user_id: int, user_group: int, + extra_dirs: List[str]) -> str: + """Returns the Dockerfile entries necessary to copy all directories in the + extra_dirs list into a docker container during build. + + """ + ret = "" + for d in extra_dirs: + ret += "\n{}".format(_copy_dir_entry(workdir, user_id, user_group, d)) + return ret + + +def _dockerfile_template( + job_mode: c.JobMode, + workdir: Optional[str] = None, + base_image_fn: Optional[Callable[[c.JobMode], str]] = None, + package: Optional[Union[List, u.Package]] = None, + requirements_path: Optional[str] = None, + conda_env_path: Optional[str] = None, + setup_extras: Optional[List[str]] = None, + adc_path: Optional[str] = None, + credentials_path: Optional[str] = None, + jupyter_version: Optional[str] = None, + inject_notebook: NotebookInstall = NotebookInstall.none, + shell: Optional[Shell] = None, + extra_dirs: Optional[List[str]] = None, + caliban_config: Optional[Dict[str, Any]] = None) -> str: + """Returns a Dockerfile that builds on a local CPU or GPU base image (depending + on the value of job_mode) to create a container that: + + - installs any dependency specified in a requirements.txt file living at + requirements_path, a conda environment at conda_env_path, or any + dependencies in a setup.py file, including extra dependencies, if + setup_extras isn't None + - injects gcloud credentials into the container, so Cloud interaction works + just like it does locally + - potentially installs a custom shell, or jupyterlab for notebook support + - copies all source needed by the main module specified by package, and + potentially injects an entrypoint that, on run, will run that main module + + Most functions that call _dockerfile_template pass along any kwargs that they + receive. It should be enough to add kwargs here, then rely on that mechanism + to pass them along, vs adding kwargs all the way down the call chain. + + Supply a custom base_image_fn (function from job_mode -> image ID) to inject + more complex Docker commands into the Caliban environments by, for example, + building your own image on top of the TF base images, then using that. + + """ + uid = os.getuid() + gid = os.getgid() + username = u.current_user() + + if isinstance(package, list): + package = u.Package(*package) + + if workdir is None: + workdir = DEFAULT_WORKDIR + + if base_image_fn is None: + base_image_fn = base_image_id + + base_image = base_image_fn(job_mode) + + dockerfile = """ +FROM {base_image} + +# Create the same group we're using on the host machine. +RUN [ $(getent group {gid}) ] || groupadd --gid {gid} {gid} + +# Create the user by name. --no-log-init guards against a crash with large user +# IDs. +RUN useradd --no-log-init --no-create-home -u {uid} -g {gid} --shell /bin/bash {username} + +# The directory is created by root. This sets permissions so that any user can +# access the folder. +RUN mkdir -m 777 {workdir} {creds_dir} {c_home} + +ENV HOME={c_home} + +WORKDIR {workdir} + +USER {uid}:{gid} +""".format_map({ + "base_image": base_image, + "username": username, + "uid": uid, + "gid": gid, + "workdir": workdir, + "c_home": container_home(), + "creds_dir": CREDS_DIR + }) + dockerfile += _credentials_entries(uid, + gid, + adc_path=adc_path, + credentials_path=credentials_path) + + dockerfile += _dependency_entries(workdir, + uid, + gid, + requirements_path=requirements_path, + conda_env_path=conda_env_path, + setup_extras=setup_extras) + + if inject_notebook.value != 'none': + install_lab = inject_notebook == NotebookInstall.lab + dockerfile += _notebook_entries(lab=install_lab, version=jupyter_version) + + if extra_dirs is not None: + dockerfile += _extra_dir_entries(workdir, uid, gid, extra_dirs) + + dockerfile += _custom_packages(uid, + gid, + packages=c.apt_packages( + caliban_config, job_mode), + shell=shell) + + if package is not None: + # The actual entrypoint and final copied code. + dockerfile += _package_entries(workdir, uid, gid, package) + + return dockerfile + + +def docker_image_id(output: str) -> ImageId: + """Accepts a string containing the output of a successful `docker build` + command and parses the Docker image ID from the stream. + + NOTE this is probably quite brittle! I can imagine this breaking quite easily + on a Docker upgrade. + + """ + return ImageId(output.splitlines()[-1].split()[-1]) + + +def build_image(job_mode: c.JobMode, + build_path: str, + credentials_path: Optional[str] = None, + adc_path: Optional[str] = None, + no_cache: bool = False, + **kwargs) -> str: + """Builds a Docker image by generating a Dockerfile and passing it to `docker + build` via stdin. All output from the `docker build` process prints to + stdout. + + Returns the image ID of the new docker container; if the command fails, + throws on error with information about the command and any issues that caused + the problem. + + """ + with u.TempCopy(credentials_path, + tmp_name=".caliban_default_creds.json") as creds: + with u.TempCopy(adc_path, tmp_name=".caliban_adc_creds.json") as adc: + cache_args = ["--no-cache"] if no_cache else [] + cmd = ["docker", "build"] + cache_args + ["--rm", "-f-", build_path] + + dockerfile = _dockerfile_template(job_mode, + credentials_path=creds, + adc_path=adc, + **kwargs) + + joined_cmd = " ".join(cmd) + logging.info("Running command: {}".format(joined_cmd)) + + try: + output, ret_code = u.capture_stdout(cmd, input_str=dockerfile) + if ret_code == 0: + return docker_image_id(output) + else: + error_msg = "Docker failed with error code {}.".format(ret_code) + raise DockerError(error_msg, cmd, ret_code) + + except subprocess.CalledProcessError as e: + logging.error(e.output) + logging.error(e.stderr) + + +def _image_tag_for_project(project_id: str, image_id: str) -> str: + """Generate the GCR Docker image tag for the supplied pair of project_id and + image_id. + + This function properly handles "domain scoped projects", where the project ID + contains a domain name and project ID separated by : + https://cloud.google.com/container-registry/docs/overview#domain-scoped_projects. + + """ + project_s = project_id.replace(":", "/") + return "gcr.io/{}/{}:latest".format(project_s, image_id) + + +def push_uuid_tag(project_id: str, image_id: str) -> str: + """Takes a base image and tags it for upload, then pushes it to a remote Google + Container Registry. + + Returns the tag on a successful push. + + TODO should this just check first before attempting to push if the image + exists? Immutable names means that if the tag is up there, we're done. + Potentially use docker-py for this. + + """ + image_tag = _image_tag_for_project(project_id, image_id) + subprocess.run(["docker", "tag", image_id, image_tag], check=True) + subprocess.run(["docker", "push", image_tag], check=True) + return image_tag + + +def _run_cmd(job_mode: c.JobMode, + run_args: Optional[List[str]] = None) -> List[str]: + """Returns the sequence of commands for the subprocess run functions required + to execute `docker run`. in CPU or GPU mode, depending on the value of + job_mode. + + Keyword args: + - run_args: list of args to pass to docker run. + + """ + if run_args is None: + run_args = [] + + runtime = ["--runtime", "nvidia"] if c.gpu(job_mode) else [] + return ["docker", "run"] + runtime + ["--ipc", "host"] + run_args + + +def _home_mount_cmds(enable_home_mount: bool) -> List[str]: + """Returns the argument needed by Docker to mount a user's local home directory + into the home directory location inside their container. + + If enable_home_mount is false returns an empty list. + + """ + ret = [] + if enable_home_mount: + ret = ["-v", "{}:{}".format(Path.home(), container_home())] + return ret + + +def _interactive_opts(workdir: str) -> List[str]: + """Returns the basic arguments we want to run a docker process locally. + + """ + return [ + "-w", workdir, \ + "-u", "{}:{}".format(os.getuid(), os.getgid()), \ + "-v", "{}:{}".format(os.getcwd(), workdir) \ + ] + + +def log_job_spec_instance(job_spec: JobSpec, i: int) -> JobSpec: + """Prints logging as a side effect for the supplied sequence of job specs + generated from an experiment definition; returns the input job spec. + + """ + args = c.experiment_to_args(job_spec.experiment.kwargs, + job_spec.experiment.args) + logging.info("") + logging.info("Job {} - Experiment args: {}".format(i, t.yellow(str(args)))) + return job_spec + + +def logged_job_specs(job_specs: Iterable[JobSpec]) -> Iterable[JobSpec]: + """Generates an iterable of job specs that should be passed to `docker run` to + execute the experiments defined by the supplied iterable. + + """ + for i, s in enumerate(job_specs, 1): + yield log_job_spec_instance(s, i) + + +def execute_dry_run(job_specs: Iterable[JobSpec]) -> None: + """Expands the supplied sequence of experiments into sequences of args and logs + the jobs that WOULD have been executed, had the dry run flag not been + applied. + + """ + list(logged_job_specs(job_specs)) + + logging.info('') + logging.info( + t.yellow("To build your image and execute these jobs, \ +run your command again without {}.".format(c.DRY_RUN_FLAG))) + logging.info('') + return None + + +def local_callback(idx: int, job: Job) -> None: + """Provides logging feedback for jobs run locally. If the return code is 0, + logs success; else, logs the failure as an error and logs the script args + that provided the failure. + + """ + if job.status == JobStatus.SUCCEEDED: + logging.info(t.green(f'Job {idx} succeeded!')) + else: + logging.error( + t.red(f'Job {idx} failed with return code {job.details["ret_code"]}.')) + args = c.experiment_to_args(job.spec.experiment.kwargs, + job.spec.experiment.args) + logging.error(t.red(f'Failing args for job {idx}: {args}')) + + +def window_size_env_cmds(): + """Returns a sequence of `docker run` arguments that will internally configure + the terminal columns and lines, so that progress bars and other terminal + interactions will work properly. + + These aren't required for interactive Docker commands like those triggered by + `caliban shell`. + + """ + ret = [] + cols, lines = _screen_shape_wrapper()(0) + if cols: + ret += ["-e", f"COLUMNS={cols}"] + if lines: + ret += ["-e", f"LINES={lines}"] + return ret + + +# ---------------------------------------------------------------------------- +def _create_job_spec_dict( + experiment: Experiment, + job_mode: c.JobMode, + image_id: str, + run_args: Optional[List[str]] = None, +) -> Dict[str, Any]: + '''creates a job spec dictionary for a local job''' + + # Without the unbuffered environment variable, stderr and stdout won't be + # emitted in the proper order from inside the container. + terminal_cmds = ["-e" "PYTHONUNBUFFERED=1"] + window_size_env_cmds() + + base_cmd = _run_cmd(job_mode, run_args) + terminal_cmds + [image_id] + command = base_cmd + c.experiment_to_args(experiment.kwargs, experiment.args) + return {'command': command, 'container': image_id} + + +# ---------------------------------------------------------------------------- +def execute_jobs( + job_specs: Iterable[JobSpec], + dry_run: bool = False, +): + '''executes a sequence of jobs based on job specs + + Arg: + job_specs: specifications for jobs to be executed + dry_run: if True, only print what would be done + ''' + + with u.tqdm_logging() as orig_stream: + pbar = tqdm.tqdm(logged_job_specs(job_specs), + file=orig_stream, + total=len(job_specs), + ascii=True, + unit="experiment", + desc="Executing") + for idx, job_spec in enumerate(pbar, 1): + command = job_spec.spec['command'] + logging.info(f'Running command: {" ".join(command)}') + if not dry_run: + _, ret_code = u.capture_stdout(command, "", u.TqdmFile(sys.stderr)) + else: + ret_code = 0 + j = Job(spec=job_spec, + container=job_spec.spec['container'], + details={'ret_code': ret_code}, + status=JobStatus.SUCCEEDED if ret_code == 0 else JobStatus.FAILED) + local_callback(idx=idx, job=j) + + if dry_run: + logging.info( + t.yellow(f'\nTo build your image and execute these jobs, ' + f'run your command again without {c.DRY_RUN_FLAG}\n')) + + return None + + +def run_experiments(job_mode: c.JobMode, + run_args: Optional[List[str]] = None, + script_args: Optional[List[str]] = None, + image_id: Optional[str] = None, + dry_run: bool = False, + experiment_config: Optional[c.ExpConf] = None, + xgroup: Optional[str] = None, + **build_image_kwargs) -> None: + """Builds an image using the supplied **build_image_kwargs and calls `docker + run` on the resulting image using sensible defaults. + + Keyword args: + + - job_mode: c.JobMode. + + - run_args: extra arguments to supply to `docker run` after our defaults. + - script_args: extra arguments to supply to the entrypoint. (You can + - override the default container entrypoint by supplying a new one inside + run_args.) + - image_id: ID of the image to run. Supplying this will skip an image build. + - experiment_config: dict of string to list, boolean, string or int. Any + lists will trigger a cartesian product out with the rest of the config. A + job will be executed for every combination of parameters in the experiment + config. + - dry_run: if True, no actual jobs will be executed and docker won't + actually build; logging side effects will show the user what will happen + without dry_run=True. + + any extra kwargs supplied are passed through to build_image. + """ + if run_args is None: + run_args = [] + + if script_args is None: + script_args = [] + + if experiment_config is None: + experiment_config = {} + + docker_args = {k: v for k, v in build_image_kwargs.items()} + docker_args['job_mode'] = job_mode + + engine = get_mem_engine() if dry_run else get_sql_engine() + + with session_scope(engine) as session: + container_spec = generate_container_spec(session, docker_args, image_id) + + if image_id is None: + if dry_run: + logging.info("Dry run - skipping actual 'docker build'.") + image_id = 'dry_run_tag' + else: + image_id = build_image(**docker_args) + + experiments = create_experiments( + session=session, + container_spec=container_spec, + script_args=script_args, + experiment_config=experiment_config, + xgroup=xgroup, + ) + + job_specs = [ + JobSpec.get_or_create( + experiment=x, + spec=_create_job_spec_dict( + experiment=x, + job_mode=job_mode, + run_args=run_args, + image_id=image_id, + ), + platform=Platform.LOCAL, + ) for x in experiments + ] + + try: + execute_jobs(job_specs=job_specs, dry_run=dry_run) + except Exception as e: + logging.error(f'exception: {e}') + session.commit() # commit here, otherwise will be rolled back + + +def run(job_mode: c.JobMode, + run_args: Optional[List[str]] = None, + script_args: Optional[List[str]] = None, + image_id: Optional[str] = None, + **build_image_kwargs) -> None: + """Builds an image using the supplied **build_image_kwargs and calls `docker + run` on the resulting image using sensible defaults. + Keyword args: + - job_mode: c.JobMode. + - run_args: extra arguments to supply to `docker run` after our defaults. + - script_args: extra arguments to supply to the entrypoint. (You can + - override the default container entrypoint by supplying a new one inside + run_args.) + - image_id: ID of the image to run. Supplying this will skip an image build. + any extra kwargs supplied are passed through to build_image. + """ + if run_args is None: + run_args = [] + + if script_args is None: + script_args = [] + + if image_id is None: + image_id = build_image(job_mode, **build_image_kwargs) + + base_cmd = _run_cmd(job_mode, run_args) + + command = base_cmd + [image_id] + script_args + + logging.info("Running command: {}".format(' '.join(command))) + subprocess.call(command) + return None + + +def run_interactive(job_mode: c.JobMode, + workdir: Optional[str] = None, + image_id: Optional[str] = None, + run_args: Optional[List[str]] = None, + mount_home: Optional[bool] = None, + shell: Optional[Shell] = None, + entrypoint: Optional[str] = None, + entrypoint_args: Optional[List[str]] = None, + **build_image_kwargs) -> None: + """Start a live shell in the terminal, with all dependencies installed and the + current working directory (and optionally the user's home directory) mounted. + + Keyword args: + + - job_mode: c.JobMode. + - image_id: ID of the image to run. Supplying this will skip an image build. + - run_args: extra arguments to supply to `docker run`. + - mount_home: if true, mounts the user's $HOME directory into the container + to `/home/$USERNAME`. If False, nothing. + - shell: name of the shell to install into the container. Also configures the + entrypoint if that's not supplied. + - entrypoint: command to run. Defaults to the executable command for the + supplied shell. + - entrypoint_args: extra arguments to supply to the entrypoint. + + any extra kwargs supplied are passed through to build_image. + + """ + if workdir is None: + workdir = DEFAULT_WORKDIR + + if run_args is None: + run_args = [] + + if entrypoint_args is None: + entrypoint_args = [] + + if mount_home is None: + mount_home = True + + if shell is None: + # Only set a default shell if we're also mounting the home volume. + # Otherwise a custom shell won't have access to the user's profile. + shell = default_shell() if mount_home else Shell.bash + + if entrypoint is None: + entrypoint = SHELL_DICT[shell].executable + + interactive_run_args = _interactive_opts(workdir) + [ + "-it", \ + "--entrypoint", entrypoint + ] + _home_mount_cmds(mount_home) + run_args + + run(job_mode=job_mode, + run_args=interactive_run_args, + script_args=entrypoint_args, + image_id=image_id, + shell=shell, + workdir=workdir, + **build_image_kwargs) + + +def run_notebook(job_mode: c.JobMode, + port: Optional[int] = None, + lab: Optional[bool] = None, + version: Optional[bool] = None, + run_args: Optional[List[str]] = None, + **run_interactive_kwargs) -> None: + """Start a notebook in the current working directory; the process will run + inside of a Docker container that's identical to the environment available to + Cloud jobs that are submitted by `caliban cloud`, or local jobs run with + `caliban run.` + + if you pass mount_home=True your jupyter settings will persist across calls. + + Keyword args: + + - port: the port to pass to Jupyter when it boots, useful if you have + multiple instances running on one machine. + - lab: if True, starts jupyter lab, else jupyter notebook. + - version: explicit Jupyter version to install. + + run_interactive_kwargs are all extra arguments taken by run_interactive. + + """ + + if port is None: + port = u.next_free_port(8888) + + if lab is None: + lab = False + + if run_args is None: + run_args = [] + + inject_arg = NotebookInstall.lab if lab else NotebookInstall.jupyter + jupyter_cmd = "lab" if lab else "notebook" + jupyter_args = [ + "-m", "jupyter", jupyter_cmd, \ + "--ip=0.0.0.0", \ + "--port={}".format(port), \ + "--no-browser" + ] + docker_args = ["-p", "{}:{}".format(port, port)] + run_args + + run_interactive(job_mode, + entrypoint="/opt/conda/envs/caliban/bin/python", + entrypoint_args=jupyter_args, + run_args=docker_args, + inject_notebook=inject_arg, + jupyter_version=version, + **run_interactive_kwargs) diff --git a/caliban/util.py b/caliban/util.py deleted file mode 100644 index 8d5f3a7..0000000 --- a/caliban/util.py +++ /dev/null @@ -1,752 +0,0 @@ -#!/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. -""" -Utilities for our job runner. -""" -import argparse -import contextlib -import getpass -import io -import itertools as it -import os -import platform -import re -import shutil -import socket -import subprocess -import sys -import time -import uuid -from collections import ChainMap -from enum import Enum -from typing import (Any, Callable, Dict, Iterable, List, NamedTuple, Optional, - Set, Tuple, Union) - -import tqdm -from absl import logging -from blessings import Terminal -from tqdm._utils import _term_move_up - -t = Terminal() - -# key and value for labels can be at most this-many-characters long. -AI_PLATFORM_MAX_LABEL_LENGTH = 63 - -Package = NamedTuple("Package", [("executable", List[str]), - ("package_path", str), ("script_path", str), - ("main_module", Optional[str])]) - - -def module_package(main_module: str) -> Package: - """Generates a Package instance for a python module executable that should be - executed with python -m. - - """ - script_path = module_to_path(main_module) - root = extract_root_directory(script_path) - return Package(["python", "-m"], - package_path=root, - script_path=script_path, - main_module=main_module) - - -def script_package(path: str, executable: str = "/bin/bash") -> Package: - """Generates a Package instance for a non-python-module executable.""" - root = extract_root_directory(path) - return Package([executable], - package_path=root, - script_path=path, - main_module=None) - - -def err(s: str) -> None: - """Prints the supplied string to stderr in red text.""" - sys.stderr.write(t.red(s)) - - -def current_user() -> str: - return getpass.getuser() - - -def is_mac() -> bool: - """Returns True if the current code is executing on a Mac, False otherwise. - - """ - return platform.system() == "Darwin" - - -def is_linux() -> bool: - """Returns True if the current code is executing on a Linux system, False - otherwise. - - """ - return platform.system() == "Darwin" - - -def enum_vals(enum: Enum) -> List[str]: - """Returns the list of all values for a specific enum.""" - return [v.value for v in enum] - - -def any_of(value_s: str, union_type: Union) -> Any: - """Attempts to parse the supplied string into one of the components of the - supplied Union. Returns the value if possible, else raises a value error. - - union_type must be a union of enums! - - """ - - def attempt(s: str, enum_type: Enum) -> Optional[Any]: - try: - return enum_type(s) - except ValueError: - return None - - enums = union_type.__args__ - ret = None - - for enum_type in enums: - ret = attempt(value_s, enum_type) - if ret is not None: - break - - if ret is None: - raise ValueError("{} isn't a value of any of {}".format(value_s, enums)) - - return ret - - -def _expand_compound_pair(k: Union[Tuple, str], v: Any) -> Dict: - """ given a key-value pair k v, where k is either: - a) a primitive representing a single, e.g. k = 'key', v = 'value', or - b) a tuple of primitives representing multiple keys, e.g. k = ('key1','key2'), v = ('value1', 'value2') - this function returns the corresponding dictionary without compound keys - """ - - if isinstance(k, tuple): - if not isinstance(v, tuple): - raise argparse.ArgumentTypeError( - """function _expand_compound_pair(k, v) requires that if type(k) is tuple, - type(v) must also be tuple.""") - else: - return dict(zip(k, v)) - else: - return {k: v} - - -def expand_compound_dict(m: Union[Dict, List]) -> Union[Dict, List]: - """ given a dictionary with some compound keys, aka tuples, - returns a dictionary which each compound key separated into primitives - - given a list of such dictionaries, will apply the transformation - described above to each dictionary and return the list, maintaining - structure - """ - - if isinstance(m, list): - return [expand_compound_dict(mi) for mi in m] - else: - expanded_dicts = [_expand_compound_pair(k, v) for k, v in m.items()] - return dict(ChainMap(*expanded_dicts)) - - -def tupleize_dict(m: Dict) -> Dict: - """ given a dictionary with compound keys, converts those keys to tuples, and - converts the corresponding values to a tuple or list of tuples - - Compound key: a string which uses square brackets to enclose - a comma-separated list, e.g. "[batch_size,learning_rate]" or "[a,b,c]" - """ - - formatted_items = [_tupleize_compound_item(k, v) for k, v in m.items()] - return dict(ChainMap(*formatted_items)) - - -def _tupleize_compound_item(k: Union[Tuple, str], v: Any) -> Dict: - """ converts a JSON-input compound key/value pair into a dictionary of tuples """ - if _is_compound_key(k): - return {_tupleize_compound_key(k): _tupleize_compound_value(v)} - else: - return {k: v} - - -def _tupleize_compound_key(k: str) -> List[str]: - """ converts a JSON-input compound key into a tuple """ - assert _is_compound_key(k), "{} must be a valid compound key".format(k) - return tuple([x.strip() for x in k.strip('][').split(',')]) - - -def _tupleize_compound_value( - v: Union[List, bool, str, int, float]) -> Union[List, Tuple]: - """ list of lists -> list of tuples - list of primitives -> tuple of primitives - single primitive -> length-1 tuple of that primitive - - E.g., [[0,1],[3,4]] -> [(0,1),(3,4)] - [0,1] -> (0,1) - 0 -> (0, ) - """ - if isinstance(v, list): - if isinstance(v[0], list): - # v is list of lists - return [tuple(vi) for vi in v] - else: - # v is list of primitives - return tuple(v) - else: - # v is a single primitive (bool, str, int, float) - return tuple([v]) - - -def _is_compound_key(s: Any) -> bool: - """ compound key is defined as a string which uses square brackets to enclose - a comma-separated list, e.g. "[batch_size,learning_rate]" or "[a,b,c]" - """ - - if type(s) is not str or len(s) <= 2: - return False - else: - return s[0] == '[' and s[-1] == ']' - - -def dict_product(m: Dict[Any, Any]) -> Iterable[Dict[Any, Any]]: - """Returns a dictionary generated by taking the cartesian product of each - list-typed value iterable with all others. - - The iterable of dictionaries returned represents every combination of values. - - If any value is NOT a list it will be treated as a singleton list. - - """ - - def wrap_v(v): - return v if isinstance(v, list) else [v] - - ks = m.keys() - vs = (wrap_v(v) for v in m.values()) - return (dict(zip(ks, x)) for x in it.product(*vs)) - - -def compose(l, r): - """Returns a function that's the composition of the two supplied functions. - - """ - - def inner(*args, **kwargs): - return l(r(*args, **kwargs)) - - return inner - - -def flipm(table: Dict[Any, Dict[Any, Any]]) -> Dict[Any, Dict[Any, Any]]: - """Handles shuffles for a particular kind of table.""" - ret = {} - for k, m in table.items(): - for k2, v in m.items(): - ret.setdefault(k2, {})[k] = v - - return ret - - -def invertm(table: Dict[Any, Iterable[Any]]) -> Dict[Any, Set[Any]]: - """Handles shuffles for a particular kind of table.""" - ret = {} - for k, vs in table.items(): - for v in vs: - ret.setdefault(v, set()).add(k) - - return ret - - -def reorderm(table: Dict[Any, Dict[Any, Iterable[Any]]], - order: Tuple[int, int, int]) -> Dict[Any, Dict[Any, Set[Any]]]: - """Handles shuffles for a particular kind of table.""" - ret = {} - for k, m in table.items(): - for k2, vs in m.items(): - for v in vs: - fields = [k, k2, v] - innerm = ret.setdefault(fields[order[0]], {}) - acc = innerm.setdefault(fields[order[1]], set()) - acc.add(fields[order[2]]) - - return ret - - -def merge(l: Dict[Any, Any], r: Dict[Any, Any]) -> Dict[Any, Any]: - """Returns a new dictionary by merging the two supplied dictionaries.""" - ret = l.copy() - ret.update(r) - return ret - - -def dict_by(keys: Set[str], f: Callable[[str], Any]) -> Dict[str, Any]: - """Returns a dictionary with keys equal to the supplied keyset. Each value is - the result of applying f to a key in keys. - - """ - return {k: f(k) for k in keys} - - -def expand_args(items: Dict[str, str]) -> List[str]: - """Converts the input map into a sequence of k, v pair strings. A None value is - interpreted to mean that the key is a solo flag; it's evicted from the - output. - - """ - pairs = [[k, v] if v is not None else [k] for k, v in items.items()] - return list(it.chain.from_iterable(pairs)) - - -def split_by(items: List[str], - separator: Optional[str] = None) -> Tuple[List[str], List[str]]: - """If the separator is present in the list, returns a 2-tuple of - - - the items before the separator, - - all items after the separator. - - If the separator isn't present, returns a tuple of - - - (the original list, []) - - """ - if separator is None: - separator = '--' - - try: - idx = items.index(separator) - return items[0:idx], items[idx + 1:] - except ValueError: - return (items, []) - - -class TempCopy(object): - """Inside its scope, this class: - - - generates a temporary file at tmp_name containing a copy of the file at - original_path, and - - deletes the new file at tmp_name when the scope exits. - - The temporary file will live inside the current directory where python's - being executed; it's a hidden file, but it will be live for the duration of - TempCopy's scope. - - We did NOT use a tmp directory here because the changing UUID name - invalidates the docker image each time a new temp path / directory is - generated. - - """ - - def __init__(self, original_path=None, tmp_name=None): - if tmp_name is None: - self.tmp_path = ".{}.json".format(str(uuid.uuid1())) - else: - self.tmp_path = tmp_name - - self.original_path = None - if original_path: - # handle tilde! - self.original_path = os.path.expanduser(original_path) - - self.path = None - - def __enter__(self): - if self.original_path is None: - return None - - current_dir = os.getcwd() - self.path = os.path.join(current_dir, self.tmp_path) - shutil.copy2(self.original_path, self.path) - return self.tmp_path - - def __exit__(self, exc_type, exc_val, exc_tb): - if self.path is not None: - os.remove(self.path) - self.path = None - - -def capture_stdout(cmd: List[str], - input_str: Optional[str] = None, - file=None) -> str: - """Executes the supplied command with the supplied string of std input, then - streams the output to stdout, and returns it as a string along with the - process's return code. - - Args: - cmd: list of strings to send in as the command - input_str: if supplied, this string will be passed as stdin to the supplied - command. if None, stdin will get closed immediately. - file: optional file-like object (stream): the output from the executed - process's stdout will get sent to this stream. Defaults to sys.stdout. - - Returns: - Pair of - - string of all stdout received during the command's execution - - return code of the process - - """ - if file is None: - file = sys.stdout - - buf = io.StringIO() - ret_code = None - - with subprocess.Popen(cmd, - stdin=subprocess.PIPE, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - universal_newlines=False, - bufsize=1) as p: - if input_str: - p.stdin.write(input_str.encode('utf-8')) - p.stdin.close() - - out = io.TextIOWrapper(p.stdout, newline='') - - for line in out: - buf.write(line) - file.write(line) - file.flush() - - # flush to force the contents to display. - file.flush() - - while p.poll() is None: - # Process hasn't exited yet, let's wait some - time.sleep(0.5) - - ret_code = p.returncode - p.stdout.close() - - return buf.getvalue(), ret_code - - -def path_to_module(path_str: str) -> str: - return path_str.replace(".py", "").replace(os.path.sep, ".") - - -def module_to_path(module_name: str) -> str: - """Converts the supplied python module (module names separated by dots) into - the python file represented by the module name. - - """ - return module_name.replace(".", os.path.sep) + ".py" - - -def file_exists_in_cwd(path: str) -> bool: - """Returns True if the current path references a valid file in the current - directory, False otherwise. - - """ - return os.path.isfile(os.path.join(os.getcwd(), path)) - - -def extract_root_directory(path: str) -> str: - """Returns the root directory of the supplied path.""" - items = path.split(os.path.sep) - return "." if len(items) == 1 else items[0] - - -def generate_package(path: str, - executable: Optional[List[str]] = None, - main_module: Optional[str] = None) -> Package: - """Takes in a string and generates a package instance that we can use for - imports. - """ - if executable is None: - _, ext = os.path.splitext(path) - executable = ["python"] if ext == ".py" else ["/bin/bash"] - - if main_module is None and not file_exists_in_cwd(path): - module_path = module_to_path(path) - - if file_exists_in_cwd(module_path): - return generate_package(module_path, - executable=["python", "-m"], - main_module=path_to_module(module_path)) - - root = extract_root_directory(path) - return Package(executable, root, path, main_module) - - -def validated_package(path: str) -> Package: - """similar to generate_package but runs argparse validation on packages that - don't actually exist in the filesystem. - - """ - p = generate_package(path) - - if not os.path.isdir(p.package_path): - raise argparse.ArgumentTypeError( - """Directory '{}' doesn't exist in directory. Code must be -nested in a folder that exists in the current directory.""".format( - p.package_path)) - - filename = p.script_path - if not file_exists_in_cwd(filename): - raise argparse.ArgumentTypeError( - """File '{}' doesn't exist locally as a script or python module; code -must live inside the current directory.""".format(filename)) - - return p - - -def parse_kv_pair(s: str) -> Tuple[str, str]: - """ - Parse a key, value pair, separated by '=' - - On the command line (argparse) a declaration will typically look like: - foo=hello - or - foo="hello world" - """ - items = s.split('=') - k = items[0].strip() # Remove whitespace around keys - - if len(items) <= 1: - raise argparse.ArgumentTypeError( - "Couldn't parse label '{}' into k=v format.".format(s)) - - v = '='.join(items[1:]) - return (k, v) - - -def _is_key(k: Optional[str]) -> bool: - """Returns True if the argument is a valid argparse optional arg input, False - otherwise. - - Strings that start with - or -- are considered valid for now. - - """ - return k is not None and len(k) > 0 and k[0] == "-" - - -def _truncate(s: str, max_length: int) -> str: - """Returns the input string s truncated to be at most max_length characters - long. - - """ - return s if len(s) <= max_length else s[0:max_length] - - -def _clean_label(s: Optional[str], is_key: bool) -> str: - """Processes the string into the sanitized format required by AI platform - labels. - - https://cloud.google.com/ml-engine/docs/resource-labels - - """ - if s is None: - return "" - - # periods are not allowed by AI Platform labels, but often occur in, - # e.g., learning rates - DECIMAL_REPLACEMENT = '_' - s = s.replace('.', DECIMAL_REPLACEMENT) - - # lowercase, letters, - and _ are valid, so strip the leading dashes, make - # everything lowercase and then kill any remaining unallowed characters. - cleaned = re.sub(r'[^a-z0-9_-]', '', s.lower()).lstrip("-") - - # Keys must start with a letter. If is_key is set and the cleaned version - # starts with something else, append `k`. - if is_key and cleaned != "" and not cleaned[0].isalpha(): - cleaned = "k" + cleaned - - return _truncate(cleaned, AI_PLATFORM_MAX_LABEL_LENGTH) - - -def key_label(k: Optional[str]) -> str: - """converts the argument into a valid label, suitable for submission as a label - key to Cloud. - - """ - return _clean_label(k, True) - - -def value_label(v: Optional[str]) -> str: - """converts the argument into a valid label, suitable for submission as a label - value to Cloud. - - """ - return _clean_label(v, False) - - -def n_chunks(items: List[Any], n_groups: int) -> List[List[Any]]: - """Returns a list of `n_groups` slices of the original list, guaranteed to - contain all of the original items. - - """ - return [items[i::n_groups] for i in range(n_groups)] - - -def chunks_below_limit(items: List[Any], limit: int) -> List[List[Any]]: - """Breaks the input list into a series of chunks guaranteed to be less than""" - quot, _ = divmod(len(items), limit) - return n_chunks(items, quot + 1) - - -def partition(seq: List[str], n: int) -> List[List[str]]: - """Generate groups of n items from seq by scanning across the sequence and - taking chunks of n, offset by 1. - - """ - for i in range(0, max(1, len(seq) - n + 1), 1): - yield seq[i:i + n] - - -def script_args_to_labels(script_args: Optional[List[str]]) -> Dict[str, str]: - """Converts the arguments supplied to our scripts into a dictionary usable as - labels valid for Cloud submission. - - """ - ret = {} - - def process_pair(k, v): - if _is_key(k): - clean_k = key_label(k) - if clean_k != "": - ret[clean_k] = "" if _is_key(v) else value_label(v) - - if script_args is None or len(script_args) == 0: - return ret - - elif len(script_args) == 1: - process_pair(script_args[0], None) - - # Handle the case where the final argument in the list is a boolean flag. - # This won't get picked up by partition. - elif len(script_args) > 1: - for k, v in partition(script_args, 2): - process_pair(k, v) - - process_pair(script_args[-1], None) - - return ret - - -def sanitize_labels( - pairs: Union[Dict[str, str], List[Tuple[str, str]]]) -> Dict[str, str]: - """Turns a dict, or a list of unsanitized key-value pairs (each represented by - a tuple) into a dictionary suitable to submit to Cloud as a label dict. - - """ - if isinstance(pairs, dict): - return sanitize_labels(pairs.items()) - - return {key_label(k): value_label(v) for (k, v) in pairs if key_label(k)} - - -def validated_directory(path: str) -> str: - """This validates that the supplied directory exists locally. - - """ - if not os.path.isdir(path): - raise argparse.ArgumentTypeError( - """Directory '{}' doesn't exist in this directory. Check yourself!""". - format(path)) - return path - - -def validated_file(path: str) -> str: - """This validates that the supplied file exists. Tilde expansion is supported. - - """ - expanded = os.path.expanduser(path) - if not os.path.isfile(expanded): - raise argparse.ArgumentTypeError( - """File '{}' isn't a valid file on your system. Try again!""".format( - path)) - return path - - -class TqdmFile(object): - """Dummy file-like that will write to tqdm""" - file = None - prefix = _term_move_up() + '\r' - - def __init__(self, file): - self.file = file - self._carriage_pending = False - - def write(self, line): - if self._carriage_pending: - line = self.prefix + line - self._carriage_pending = False - - if line.endswith('\r'): - self._carriage_pending = True - line = line[:-1] + '\n' - - tqdm.tqdm.write(line, file=self.file, end='') - - def flush(self): - return getattr(self.file, "flush", lambda: None)() - - def isatty(self): - return getattr(self.file, "isatty", lambda: False)() - - def close(self): - return getattr(self.file, "close", lambda: None)() - - -def config_logging(): - """Overrides logging to go through TQDM. - - TODO use this call to kill then restore: - https://github.com/tqdm/tqdm#redirecting-writing - - """ - h = logging.get_absl_handler() - old = h.python_handler - h._python_handler = logging.PythonHandler(stream=TqdmFile(sys.stderr)) - logging.use_python_logging() - - -@contextlib.contextmanager -def tqdm_logging(): - """Overrides logging to go through TQDM. - - https://github.com/tqdm/tqdm#redirecting-writing - - """ - handler = logging.get_absl_handler() - orig = handler.python_handler - - try: - handler._python_handler = logging.PythonHandler(stream=TqdmFile(sys.stderr)) - - # The changes won't take effect if this hasn't been called. Defensively - # call it again here. - logging.use_python_logging() - yield orig.stream - except Exception as exc: - raise exc - finally: - handler._python_handler = orig - - -def next_free_port(port: int, try_n: int = 1000, max_port=65535): - if try_n == 0 or port <= max_port: - sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - try: - sock.bind(('', port)) - sock.close() - return port - except OSError: - return next_free_port(port + 1, try_n - 1, max_port=max_port) - else: - raise IOError('no free ports') diff --git a/caliban/util/__init__.py b/caliban/util/__init__.py new file mode 100644 index 0000000..290200e --- /dev/null +++ b/caliban/util/__init__.py @@ -0,0 +1,180 @@ +#!/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. +""" +Utilities for our job runner. +""" +import getpass +import itertools as it +import platform +import sys +from enum import Enum +from typing import (Any, Callable, Dict, Iterable, List, NamedTuple, Optional, + Set, Tuple, Union) + +from blessings import Terminal + +t = Terminal() + +Package = NamedTuple("Package", [("executable", List[str]), + ("package_path", str), ("script_path", str), + ("main_module", Optional[str])]) + + +def err(s: str) -> None: + """Prints the supplied string to stderr in red text.""" + sys.stderr.write(t.red(s)) + + +def current_user() -> str: + return getpass.getuser() + + +def is_mac() -> bool: + """Returns True if the current code is executing on a Mac, False otherwise. + + """ + return platform.system() == "Darwin" + + +def is_linux() -> bool: + """Returns True if the current code is executing on a Linux system, False + otherwise. + + """ + return platform.system() == "Darwin" + + +def enum_vals(enum: Enum) -> List[str]: + """Returns the list of all values for a specific enum.""" + return [v.value for v in enum] + + +def any_of(value_s: str, union_type: Union) -> Any: + """Attempts to parse the supplied string into one of the components of the + supplied Union. Returns the value if possible, else raises a value error. + + union_type must be a union of enums! + + """ + + def attempt(s: str, enum_type: Enum) -> Optional[Any]: + try: + return enum_type(s) + except ValueError: + return None + + enums = union_type.__args__ + ret = None + + for enum_type in enums: + ret = attempt(value_s, enum_type) + if ret is not None: + break + + if ret is None: + raise ValueError("{} isn't a value of any of {}".format(value_s, enums)) + + return ret + + +def dict_product(m: Dict[Any, Any]) -> Iterable[Dict[Any, Any]]: + """Returns a dictionary generated by taking the cartesian product of each + list-typed value iterable with all others. + + The iterable of dictionaries returned represents every combination of values. + + If any value is NOT a list it will be treated as a singleton list. + + """ + + def wrap_v(v): + return v if isinstance(v, list) else [v] + + ks = m.keys() + vs = (wrap_v(v) for v in m.values()) + return (dict(zip(ks, x)) for x in it.product(*vs)) + + +def flipm(table: Dict[Any, Dict[Any, Any]]) -> Dict[Any, Dict[Any, Any]]: + """Handles shuffles for a particular kind of table.""" + ret = {} + for k, m in table.items(): + for k2, v in m.items(): + ret.setdefault(k2, {})[k] = v + + return ret + + +def invertm(table: Dict[Any, Iterable[Any]]) -> Dict[Any, Set[Any]]: + """Handles shuffles for a particular kind of table.""" + ret = {} + for k, vs in table.items(): + for v in vs: + ret.setdefault(v, set()).add(k) + + return ret + + +def reorderm(table: Dict[Any, Dict[Any, Iterable[Any]]], + order: Tuple[int, int, int]) -> Dict[Any, Dict[Any, Set[Any]]]: + """Handles shuffles for a particular kind of table.""" + ret = {} + for k, m in table.items(): + for k2, vs in m.items(): + for v in vs: + fields = [k, k2, v] + innerm = ret.setdefault(fields[order[0]], {}) + acc = innerm.setdefault(fields[order[1]], set()) + acc.add(fields[order[2]]) + + return ret + + +def merge(l: Dict[Any, Any], r: Dict[Any, Any]) -> Dict[Any, Any]: + """Returns a new dictionary by merging the two supplied dictionaries.""" + ret = l.copy() + ret.update(r) + return ret + + +def dict_by(keys: Set[str], f: Callable[[str], Any]) -> Dict[str, Any]: + """Returns a dictionary with keys equal to the supplied keyset. Each value is + the result of applying f to a key in keys. + + """ + return {k: f(k) for k in keys} + + +def split_by(items: List[str], + separator: Optional[str] = None) -> Tuple[List[str], List[str]]: + """If the separator is present in the list, returns a 2-tuple of + + - the items before the separator, + - all items after the separator. + + If the separator isn't present, returns a tuple of + + - (the original list, []) + + """ + if separator is None: + separator = '--' + + try: + idx = items.index(separator) + return items[0:idx], items[idx + 1:] + except ValueError: + return (items, []) diff --git a/caliban/util/argparse.py b/caliban/util/argparse.py new file mode 100644 index 0000000..7d36e3f --- /dev/null +++ b/caliban/util/argparse.py @@ -0,0 +1,129 @@ +#!/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. +""" +Utilities for our job runner. +""" +import caliban.util as u +import argparse +import contextlib +import getpass +import io +import itertools as it +import os +import platform +import re +import shutil +import socket +import subprocess +import sys +import time +import uuid +from collections import ChainMap +from enum import Enum +from typing import (Any, Callable, Dict, Iterable, List, NamedTuple, Optional, + Set, Tuple, Union) + +import tqdm +from absl import logging +from blessings import Terminal +from tqdm.utils import _term_move_up + +t = Terminal() + + +def expand_args(items: Dict[str, str]) -> List[str]: + """Converts the input map into a sequence of k, v pair strings. A None value is + interpreted to mean that the key is a solo flag; it's evicted from the + output. + + """ + pairs = [[k, v] if v is not None else [k] for k, v in items.items()] + return list(it.chain.from_iterable(pairs)) + + +def validated_package(path: str) -> u.Package: + """similar to generate_package but runs argparse validation on packages that + don't actually exist in the filesystem. + + """ + p = u.generate_package(path) + + if not os.path.isdir(p.package_path): + raise argparse.ArgumentTypeError( + """Directory '{}' doesn't exist in directory. Code must be +nested in a folder that exists in the current directory.""".format( + p.package_path)) + + filename = p.script_path + if not file_exists_in_cwd(filename): + raise argparse.ArgumentTypeError( + """File '{}' doesn't exist locally as a script or python module; code +must live inside the current directory.""".format(filename)) + + return p + + +def parse_kv_pair(s: str) -> Tuple[str, str]: + """ + Parse a key, value pair, separated by '=' + + On the command line (argparse) a declaration will typically look like: + foo=hello + or + foo="hello world" + """ + items = s.split('=') + k = items[0].strip() # Remove whitespace around keys + + if len(items) <= 1: + raise argparse.ArgumentTypeError( + "Couldn't parse label '{}' into k=v format.".format(s)) + + v = '='.join(items[1:]) + return (k, v) + + +def is_key(k: Optional[str]) -> bool: + """Returns True if the argument is a valid argparse optional arg input, False + otherwise. + + Strings that start with - or -- are considered valid for now. + + """ + return k is not None and len(k) > 0 and k[0] == "-" + + +def validated_directory(path: str) -> str: + """This validates that the supplied directory exists locally. + + """ + if not os.path.isdir(path): + raise argparse.ArgumentTypeError( + """Directory '{}' doesn't exist in this directory. Check yourself!""". + format(path)) + return path + + +def validated_file(path: str) -> str: + """This validates that the supplied file exists. Tilde expansion is supported. + + """ + expanded = os.path.expanduser(path) + if not os.path.isfile(expanded): + raise argparse.ArgumentTypeError( + """File '{}' isn't a valid file on your system. Try again!""".format( + path)) + return path diff --git a/caliban/util/fs.py b/caliban/util/fs.py new file mode 100644 index 0000000..27c80d0 --- /dev/null +++ b/caliban/util/fs.py @@ -0,0 +1,219 @@ +#!/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. +"""Utilities for interacting with the filesystem and packages. + +""" +import io +import os +import shutil +import socket +import subprocess +import sys +import time +import uuid +from typing import List, NamedTuple, Optional + +from blessings import Terminal + +t = Terminal() + +Package = NamedTuple("Package", [("executable", List[str]), + ("package_path", str), ("script_path", str), + ("main_module", Optional[str])]) + + +def file_exists_in_cwd(path: str) -> bool: + """Returns True if the current path references a valid file in the current + directory, False otherwise. + + """ + return os.path.isfile(os.path.join(os.getcwd(), path)) + + +def extract_root_directory(path: str) -> str: + """Returns the root directory of the supplied path.""" + items = path.split(os.path.sep) + return "." if len(items) == 1 else items[0] + + +def module_package(main_module: str) -> Package: + """Generates a Package instance for a python module executable that should be + executed with python -m. + + """ + script_path = module_to_path(main_module) + root = extract_root_directory(script_path) + return Package(["python", "-m"], + package_path=root, + script_path=script_path, + main_module=main_module) + + +def script_package(path: str, executable: str = "/bin/bash") -> Package: + """Generates a Package instance for a non-python-module executable.""" + root = extract_root_directory(path) + return Package([executable], + package_path=root, + script_path=path, + main_module=None) + + +def path_to_module(path_str: str) -> str: + return path_str.replace(".py", "").replace(os.path.sep, ".") + + +def module_to_path(module_name: str) -> str: + """Converts the supplied python module (module names separated by dots) into + the python file represented by the module name. + + """ + return module_name.replace(".", os.path.sep) + ".py" + + +def generate_package(path: str, + executable: Optional[List[str]] = None, + main_module: Optional[str] = None) -> Package: + """Takes in a string and generates a package instance that we can use for + imports. + """ + if executable is None: + _, ext = os.path.splitext(path) + executable = ["python"] if ext == ".py" else ["/bin/bash"] + + if main_module is None and not file_exists_in_cwd(path): + module_path = module_to_path(path) + + if file_exists_in_cwd(module_path): + return generate_package(module_path, + executable=["python", "-m"], + main_module=path_to_module(module_path)) + + root = extract_root_directory(path) + return Package(executable, root, path, main_module) + + +class TempCopy(object): + """Inside its scope, this class: + + - generates a temporary file at tmp_name containing a copy of the file at + original_path, and + - deletes the new file at tmp_name when the scope exits. + + The temporary file will live inside the current directory where python's + being executed; it's a hidden file, but it will be live for the duration of + TempCopy's scope. + + We did NOT use a tmp directory here because the changing UUID name + invalidates the docker image each time a new temp path / directory is + generated. + + """ + + def __init__(self, original_path=None, tmp_name=None): + if tmp_name is None: + self.tmp_path = ".{}.json".format(str(uuid.uuid1())) + else: + self.tmp_path = tmp_name + + self.original_path = None + if original_path: + # handle tilde! + self.original_path = os.path.expanduser(original_path) + + self.path = None + + def __enter__(self): + if self.original_path is None: + return None + + current_dir = os.getcwd() + self.path = os.path.join(current_dir, self.tmp_path) + shutil.copy2(self.original_path, self.path) + return self.tmp_path + + def __exit__(self, exc_type, exc_val, exc_tb): + if self.path is not None: + os.remove(self.path) + self.path = None + + +def capture_stdout(cmd: List[str], + input_str: Optional[str] = None, + file=None) -> str: + """Executes the supplied command with the supplied string of std input, then + streams the output to stdout, and returns it as a string along with the + process's return code. + + Args: + cmd: list of strings to send in as the command + input_str: if supplied, this string will be passed as stdin to the supplied + command. if None, stdin will get closed immediately. + file: optional file-like object (stream): the output from the executed + process's stdout will get sent to this stream. Defaults to sys.stdout. + + Returns: + Pair of + - string of all stdout received during the command's execution + - return code of the process + + """ + if file is None: + file = sys.stdout + + buf = io.StringIO() + ret_code = None + + with subprocess.Popen(cmd, + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + universal_newlines=False, + bufsize=1) as p: + if input_str: + p.stdin.write(input_str.encode('utf-8')) + p.stdin.close() + + out = io.TextIOWrapper(p.stdout, newline='') + + for line in out: + buf.write(line) + file.write(line) + file.flush() + + # flush to force the contents to display. + file.flush() + + while p.poll() is None: + # Process hasn't exited yet, let's wait some + time.sleep(0.5) + + ret_code = p.returncode + p.stdout.close() + + return buf.getvalue(), ret_code + + +def next_free_port(port: int, try_n: int = 1000, max_port=65535): + if try_n == 0 or port <= max_port: + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + try: + sock.bind(('', port)) + sock.close() + return port + except OSError: + return next_free_port(port + 1, try_n - 1, max_port=max_port) + else: + raise IOError('no free ports') diff --git a/caliban/util/tqdm.py b/caliban/util/tqdm.py new file mode 100644 index 0000000..13ba49d --- /dev/null +++ b/caliban/util/tqdm.py @@ -0,0 +1,95 @@ +#!/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. +""" +Progress bar utilities. +""" + +import contextlib +import sys + +from absl import logging +from blessings import Terminal + +import tqdm +from tqdm.utils import _term_move_up + +t = Terminal() + + +class TqdmFile(object): + """Dummy file-like that will write to tqdm""" + file = None + prefix = _term_move_up() + '\r' + + def __init__(self, file): + self.file = file + self._carriage_pending = False + + def write(self, line): + if self._carriage_pending: + line = self.prefix + line + self._carriage_pending = False + + if line.endswith('\r'): + self._carriage_pending = True + line = line[:-1] + '\n' + + tqdm.tqdm.write(line, file=self.file, end='') + + def flush(self): + return getattr(self.file, "flush", lambda: None)() + + def isatty(self): + return getattr(self.file, "isatty", lambda: False)() + + def close(self): + return getattr(self.file, "close", lambda: None)() + + +def config_logging(): + """Overrides logging to go through TQDM. + + TODO use this call to kill then restore: + https://github.com/tqdm/tqdm#redirecting-writing + + """ + h = logging.get_absl_handler() + old = h.python_handler + h._python_handler = logging.PythonHandler(stream=TqdmFile(sys.stderr)) + logging.use_python_logging() + + +@contextlib.contextmanager +def tqdm_logging(): + """Overrides logging to go through TQDM. + + https://github.com/tqdm/tqdm#redirecting-writing + + """ + handler = logging.get_absl_handler() + orig = handler.python_handler + + try: + handler._python_handler = logging.PythonHandler(stream=TqdmFile(sys.stderr)) + + # The changes won't take effect if this hasn't been called. Defensively + # call it again here. + logging.use_python_logging() + yield orig.stream + except Exception as exc: + raise exc + finally: + handler._python_handler = orig diff --git a/tests/caliban/config/test_config.py b/tests/caliban/config/test_config.py new file mode 100644 index 0000000..aa49498 --- /dev/null +++ b/tests/caliban/config/test_config.py @@ -0,0 +1,44 @@ +#!/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. + +import os +from argparse import ArgumentTypeError + +import caliban.cloud.types as ct +import caliban.config as c +import pytest + + +def test_extract_region(monkeypatch): + if os.environ.get('REGION'): + monkeypatch.delenv('REGION') + + assert c.extract_region({}) == c.DEFAULT_REGION + + # You have to provide a valid region. + with pytest.raises(ArgumentTypeError): + c.extract_region({"region": "face"}) + + # Same goes for the environment variable setting approach. + monkeypatch.setenv('REGION', "face") + with pytest.raises(ArgumentTypeError): + c.extract_region({}) + + # an empty string is fine, and ignored. + monkeypatch.setenv('REGION', "") + assert c.extract_region({}) == c.DEFAULT_REGION + + assert c.extract_region({"region": "us-west1"}) == ct.US.west1 diff --git a/tests/caliban/config/test_experiment.py b/tests/caliban/config/test_experiment.py new file mode 100644 index 0000000..30f134d --- /dev/null +++ b/tests/caliban/config/test_experiment.py @@ -0,0 +1,229 @@ +#!/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. + +from argparse import ArgumentTypeError + +import caliban.config.experiment as c +import pytest + + +def test_validate_experiment_config(): + """basic examples of validate experiment config.""" + invalid = {1: "face", "2": "3"} + with pytest.raises(ArgumentTypeError): + c.validate_experiment_config(invalid) + + # a dict value is invalid, even if it's hidden in a list. + with pytest.raises(ArgumentTypeError): + c.validate_experiment_config({"key": [{1: 2}, "face"]}) + + valid = {"a": [1.0, 2, 3], "b": True, "c": 1, "d": "e", "f": 1.2} + assert valid == c.validate_experiment_config(valid) + + # Lists are okay too... + items = [valid, valid] + assert items == c.validate_experiment_config(items) + + # As are lists of lists. + lol = [valid, [valid]] + assert lol == c.validate_experiment_config(lol) + + # Invalid types are caught even nested inside lists. + lol_invalid = [valid, valid, [invalid]] + + with pytest.raises(ArgumentTypeError): + c.validate_experiment_config(lol_invalid) + + # Compound keys which violate syntax rules are caught + invalid_compound = [{ + "[": 0 + }, { + "eh[": 0 + }, { + "[test,,fail]": 0 + }, { + "[,test,fail]": 0 + }, { + "[I,will,fail,]": 0 + }, { + "[I,,will,fail]": 0 + }, { + "]I,will,fail]": 0 + }] + + valid_compound = [{ + "[batch_size,learning_rate]": [0, 1] + }, { + "[batch_size,learning_rate,dataset_size]": [0.01, 0.02, 100] + }, { + "[batch_size,learning_rate,dataset_size]": [[0.01, 0.02, 100], + [0.03, 0.05, 200]] + }, { + "[batch_size,learning_rate]": [[0., 1.], [2., 3.]] + }, { + "[batch_size,learning_rate]": [[0., 1.], [2., 3.], [4., 5.]] + }, { + "[batch_size, learning_rate, dataset_size]": [0.01, 0.02, 100] + }, { + "[batch_size , learning_rate,dataset_size]": [[0.01, 0.02, 100], + [0.03, 0.05, 200]] + }, { + "[batch_size, learning_rate]": [[0., 1.], [2., 3.]] + }, { + "[batch_size ,learning_rate]": [[0., 1.], [2., 3.], [4., 5.]] + }] + + for i in invalid_compound: + with pytest.raises(Exception): + c.validate_experiment_config(i) + + for i in valid_compound: + assert i == c.validate_experiment_config(i) + + +def test_expand_experiment_config(): + # An empty config expands to a singleton list. This is important so that + # single job submission without a spec works. + assert [{}] == c.expand_experiment_config({}) + + +def test_compound_key_handling(): + tests = [{ + 'input': { + '[a,b]': [['c', 'd'], ['e', 'f']] + }, + 'after_tupleization': { + ('a', 'b'): [('c', 'd'), ('e', 'f')] + }, + 'after_dictproduct': [{ + ('a', 'b'): ('c', 'd') + }, { + ('a', 'b'): ('e', 'f') + }], + 'after_expansion': [{ + 'a': 'c', + 'b': 'd' + }, { + 'a': 'e', + 'b': 'f' + }] + }, { + 'input': { + '[a,b]': ['c', 'd'] + }, + 'after_tupleization': { + ('a', 'b'): ('c', 'd') + }, + 'after_dictproduct': [{ + ('a', 'b'): ('c', 'd') + }], + 'after_expansion': [{ + 'a': 'c', + 'b': 'd' + }] + }, { + 'input': { + 'hi': 'there', + '[k1,k2]': [['v1a', 'v2a'], ['v1b', 'v2b']] + }, + 'after_tupleization': { + 'hi': 'there', + ('k1', 'k2'): [('v1a', 'v2a'), ('v1b', 'v2b')] + }, + 'after_dictproduct': [{ + 'hi': 'there', + ('k1', 'k2'): ('v1a', 'v2a') + }, { + 'hi': 'there', + ('k1', 'k2'): ('v1b', 'v2b') + }], + 'after_expansion': [{ + 'hi': 'there', + 'k1': 'v1a', + 'k2': 'v2a' + }, { + 'hi': 'there', + 'k1': 'v1b', + 'k2': 'v2b' + }] + }, { + 'input': { + 'hi': 'there', + '[a,b]': ['c', 'd'] + }, + 'after_tupleization': { + 'hi': 'there', + ('a', 'b'): ('c', 'd') + }, + 'after_dictproduct': [{ + 'hi': 'there', + ('a', 'b'): ('c', 'd') + }], + 'after_expansion': [{ + 'hi': 'there', + 'a': 'c', + 'b': 'd' + }] + }, { + 'input': { + '[a,b]': [0, 1] + }, + 'after_tupleization': { + ('a', 'b'): (0, 1) + }, + 'after_dictproduct': [{ + ('a', 'b'): (0, 1) + }], + 'after_expansion': [{ + 'a': 0, + 'b': 1 + }] + }, { + 'input': { + '[a,b]': [[0, 1]] + }, + 'after_tupleization': { + ('a', 'b'): [(0, 1)] + }, + 'after_dictproduct': [{ + ('a', 'b'): (0, 1) + }], + 'after_expansion': [{ + 'a': 0, + 'b': 1 + }] + }, { + 'input': { + 'hi': 'blueshift', + '[a,b]': [[0, 1]] + }, + 'after_tupleization': { + 'hi': 'blueshift', + ('a', 'b'): [(0, 1)] + }, + 'after_dictproduct': [{ + 'hi': 'blueshift', + ('a', 'b'): (0, 1) + }], + 'after_expansion': [{ + 'hi': 'blueshift', + 'a': 0, + 'b': 1 + }] + }] + + for test in tests: + assert test['after_expansion'] == c.expand_experiment_config(test['input']) diff --git a/tests/caliban/test_config.py b/tests/caliban/test_config.py deleted file mode 100644 index df42969..0000000 --- a/tests/caliban/test_config.py +++ /dev/null @@ -1,257 +0,0 @@ -#!/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. - -import os -import unittest -from argparse import ArgumentTypeError - -import caliban.cloud.types as ct -import caliban.config as c -import pytest - - -def test_extract_region(monkeypatch): - if os.environ.get('REGION'): - monkeypatch.delenv('REGION') - - assert c.extract_region({}) == c.DEFAULT_REGION - - # You have to provide a valid region. - with pytest.raises(ArgumentTypeError): - c.extract_region({"region": "face"}) - - # Same goes for the environment variable setting approach. - monkeypatch.setenv('REGION', "face") - with pytest.raises(ArgumentTypeError): - c.extract_region({}) - - # an empty string is fine, and ignored. - monkeypatch.setenv('REGION', "") - assert c.extract_region({}) == c.DEFAULT_REGION - - assert c.extract_region({"region": "us-west1"}) == ct.US.west1 - - -class ConfigTestSuite(unittest.TestCase): - """Tests for the config package.""" - - def test_validate_experiment_config(self): - """basic examples of validate experiment config.""" - invalid = {1: "face", "2": "3"} - with self.assertRaises(ArgumentTypeError): - c.validate_experiment_config(invalid) - - # a dict value is invalid, even if it's hidden in a list. - with self.assertRaises(ArgumentTypeError): - c.validate_experiment_config({"key": [{1: 2}, "face"]}) - - valid = {"a": [1.0, 2, 3], "b": True, "c": 1, "d": "e", "f": 1.2} - self.assertDictEqual(valid, c.validate_experiment_config(valid)) - - # Lists are okay too... - items = [valid, valid] - self.assertListEqual(items, c.validate_experiment_config(items)) - - # As are lists of lists. - lol = [valid, [valid]] - self.assertListEqual(lol, c.validate_experiment_config(lol)) - - # Invalid types are caught even nested inside lists. - lol_invalid = [valid, valid, [invalid]] - with self.assertRaises(ArgumentTypeError): - c.validate_experiment_config(lol_invalid) - - # Compound keys which violate syntax rules are caught - invalid_compound = [{ - "[": 0 - }, { - "eh[": 0 - }, { - "[test,,fail]": 0 - }, { - "[,test,fail]": 0 - }, { - "[I,will,fail,]": 0 - }, { - "[I,,will,fail]": 0 - }, { - "]I,will,fail]": 0 - }] - valid_compound = [{ - "[batch_size,learning_rate]": [0, 1] - }, { - "[batch_size,learning_rate,dataset_size]": [0.01, 0.02, 100] - }, { - "[batch_size,learning_rate,dataset_size]": [[0.01, 0.02, 100], - [0.03, 0.05, 200]] - }, { - "[batch_size,learning_rate]": [[0., 1.], [2., 3.]] - }, { - "[batch_size,learning_rate]": [[0., 1.], [2., 3.], [4., 5.]] - }, { - "[batch_size, learning_rate, dataset_size]": [0.01, 0.02, 100] - }, { - "[batch_size , learning_rate,dataset_size]": [[0.01, 0.02, 100], - [0.03, 0.05, 200]] - }, { - "[batch_size, learning_rate]": [[0., 1.], [2., 3.]] - }, { - "[batch_size ,learning_rate]": [[0., 1.], [2., 3.], [4., 5.]] - }] - - for i in invalid_compound: - with self.assertRaises(Exception): - c.validate_experiment_config(i) - for i in valid_compound: - self.assertDictEqual(i, c.validate_experiment_config(i)) - - def test_expand_experiment_config(self): - # An empty config expands to a singleton list. This is important so that - # single job submission without a spec works. - self.assertListEqual([{}], c.expand_experiment_config({})) - - def test_compound_key_handling(self): - tests = [{ - 'input': { - '[a,b]': [['c', 'd'], ['e', 'f']] - }, - 'after_tupleization': { - ('a', 'b'): [('c', 'd'), ('e', 'f')] - }, - 'after_dictproduct': [{ - ('a', 'b'): ('c', 'd') - }, { - ('a', 'b'): ('e', 'f') - }], - 'after_expansion': [{ - 'a': 'c', - 'b': 'd' - }, { - 'a': 'e', - 'b': 'f' - }] - }, { - 'input': { - '[a,b]': ['c', 'd'] - }, - 'after_tupleization': { - ('a', 'b'): ('c', 'd') - }, - 'after_dictproduct': [{ - ('a', 'b'): ('c', 'd') - }], - 'after_expansion': [{ - 'a': 'c', - 'b': 'd' - }] - }, { - 'input': { - 'hi': 'there', - '[k1,k2]': [['v1a', 'v2a'], ['v1b', 'v2b']] - }, - 'after_tupleization': { - 'hi': 'there', - ('k1', 'k2'): [('v1a', 'v2a'), ('v1b', 'v2b')] - }, - 'after_dictproduct': [{ - 'hi': 'there', - ('k1', 'k2'): ('v1a', 'v2a') - }, { - 'hi': 'there', - ('k1', 'k2'): ('v1b', 'v2b') - }], - 'after_expansion': [{ - 'hi': 'there', - 'k1': 'v1a', - 'k2': 'v2a' - }, { - 'hi': 'there', - 'k1': 'v1b', - 'k2': 'v2b' - }] - }, { - 'input': { - 'hi': 'there', - '[a,b]': ['c', 'd'] - }, - 'after_tupleization': { - 'hi': 'there', - ('a', 'b'): ('c', 'd') - }, - 'after_dictproduct': [{ - 'hi': 'there', - ('a', 'b'): ('c', 'd') - }], - 'after_expansion': [{ - 'hi': 'there', - 'a': 'c', - 'b': 'd' - }] - }, { - 'input': { - '[a,b]': [0, 1] - }, - 'after_tupleization': { - ('a', 'b'): (0, 1) - }, - 'after_dictproduct': [{ - ('a', 'b'): (0, 1) - }], - 'after_expansion': [{ - 'a': 0, - 'b': 1 - }] - }, { - 'input': { - '[a,b]': [[0, 1]] - }, - 'after_tupleization': { - ('a', 'b'): [(0, 1)] - }, - 'after_dictproduct': [{ - ('a', 'b'): (0, 1) - }], - 'after_expansion': [{ - 'a': 0, - 'b': 1 - }] - }, { - 'input': { - 'hi': 'blueshift', - '[a,b]': [[0, 1]] - }, - 'after_tupleization': { - 'hi': 'blueshift', - ('a', 'b'): [(0, 1)] - }, - 'after_dictproduct': [{ - 'hi': 'blueshift', - ('a', 'b'): (0, 1) - }], - 'after_expansion': [{ - 'hi': 'blueshift', - 'a': 0, - 'b': 1 - }] - }] - - for test in tests: - self.assertListEqual(test['after_expansion'], - c.expand_experiment_config(test['input'])) - - -if __name__ == '__main__': - unittest.main() diff --git a/tests/caliban/test_util.py b/tests/caliban/test_util.py index c02ea03..d46719b 100644 --- a/tests/caliban/test_util.py +++ b/tests/caliban/test_util.py @@ -302,22 +302,6 @@ def check_expansion(test_dict): check_dictproduct(test) check_expansion(test) - @given(st.integers()) - def test_compose(self, x): - """Functions should compose; the composed function accepts any arguments that the rightmost function accepts.""" - - def plus1(x): - return x + 1 - - def square(x): - return x * x - - square_plus_one = u.compose(plus1, square) - times_plus_one = u.compose(plus1, lambda l, r: l * r) - - self.assertEqual(square_plus_one(x), x * x + 1) - self.assertEqual(square_plus_one(x), times_plus_one(l=x, r=x)) - @given(st.dictionaries(st.text(), st.text()), st.dictionaries(st.text(), st.text())) def test_merge(self, m1, m2): From 8170450aba98a63be60591ab03d46a01311149b7 Mon Sep 17 00:00:00 2001 From: Sam Ritchie Date: Fri, 10 Jul 2020 09:15:21 -0600 Subject: [PATCH 02/15] another pass --- caliban/builder.py | 5 +- caliban/cli.py | 2 +- caliban/docker/push.py | 5 +- caliban/history/cli.py | 8 ++-- caliban/history/{utils.py => util.py} | 13 +++--- caliban/platform/cloud/core.py | 5 +- caliban/platform/gke/cli.py | 27 ++++++----- caliban/platform/gke/cluster.py | 37 ++++++++------- caliban/platform/gke/{utils.py => util.py} | 0 caliban/platform/notebook.py | 5 +- caliban/platform/run.py | 5 +- caliban/platform/shell.py | 5 +- caliban/types.py | 46 ------------------- tests/caliban/cloud/__init__.py | 31 ------------- tests/caliban/config/__init__.py | 15 ++++++ tests/caliban/docker/__init__.py | 15 ++++++ tests/caliban/{ => docker}/test_docker.py | 0 tests/caliban/gke/__init__.py | 31 ------------- tests/caliban/history/__init__.py | 16 ------- tests/caliban/history/test_history.py | 31 +++---------- tests/caliban/platform/__init__.py | 15 ++++++ tests/caliban/platform/cloud/__init__.py | 15 ++++++ .../{ => platform}/cloud/test_types.py | 18 -------- tests/caliban/platform/gke/__init__.py | 15 ++++++ .../caliban/{ => platform}/gke/test_types.py | 22 +-------- .../caliban/{ => platform}/gke/test_utils.py | 14 ++---- tests/caliban/util/__init__.py | 15 ++++++ tests/caliban/{ => util}/test_util.py | 0 28 files changed, 156 insertions(+), 260 deletions(-) rename caliban/history/{utils.py => util.py} (98%) rename caliban/platform/gke/{utils.py => util.py} (100%) delete mode 100644 caliban/types.py delete mode 100644 tests/caliban/cloud/__init__.py create mode 100644 tests/caliban/config/__init__.py create mode 100644 tests/caliban/docker/__init__.py rename tests/caliban/{ => docker}/test_docker.py (100%) delete mode 100644 tests/caliban/gke/__init__.py create mode 100644 tests/caliban/platform/__init__.py create mode 100644 tests/caliban/platform/cloud/__init__.py rename tests/caliban/{ => platform}/cloud/test_types.py (81%) create mode 100644 tests/caliban/platform/gke/__init__.py rename tests/caliban/{ => platform}/gke/test_types.py (65%) rename tests/caliban/{ => platform}/gke/test_utils.py (98%) create mode 100644 tests/caliban/util/__init__.py rename tests/caliban/{ => util}/test_util.py (100%) diff --git a/caliban/builder.py b/caliban/builder.py index 298aa1f..893438c 100644 --- a/caliban/builder.py +++ b/caliban/builder.py @@ -37,9 +37,8 @@ import caliban.config as c import caliban.util as u from caliban.history.types import Experiment, Job, JobSpec, JobStatus, Platform -from caliban.history.utils import (create_experiments, generate_container_spec, - get_mem_engine, get_sql_engine, - session_scope) +from caliban.history.util import (create_experiments, generate_container_spec, + get_mem_engine, get_sql_engine, session_scope) t = Terminal() diff --git a/caliban/cli.py b/caliban/cli.py index 1a2eeee..ae270aa 100644 --- a/caliban/cli.py +++ b/caliban/cli.py @@ -31,7 +31,7 @@ import caliban.gke as gke import caliban.gke.constants as gke_k import caliban.gke.types as gke_t -import caliban.gke.utils as gke_u +import caliban.gke.util as gke_u import caliban.util as u from caliban import __version__ diff --git a/caliban/docker/push.py b/caliban/docker/push.py index 298aa1f..893438c 100644 --- a/caliban/docker/push.py +++ b/caliban/docker/push.py @@ -37,9 +37,8 @@ import caliban.config as c import caliban.util as u from caliban.history.types import Experiment, Job, JobSpec, JobStatus, Platform -from caliban.history.utils import (create_experiments, generate_container_spec, - get_mem_engine, get_sql_engine, - session_scope) +from caliban.history.util import (create_experiments, generate_container_spec, + get_mem_engine, get_sql_engine, session_scope) t = Terminal() diff --git a/caliban/history/cli.py b/caliban/history/cli.py index 5cf3679..4eb8e14 100644 --- a/caliban/history/cli.py +++ b/caliban/history/cli.py @@ -24,13 +24,13 @@ from sqlalchemy.orm import Session from caliban.util import current_user, Package -from caliban.history.utils import (get_sql_engine, session_scope, - update_job_status, get_gke_job_name, - stop_job, replace_job_spec_image) +from caliban.history.util import (get_sql_engine, session_scope, + update_job_status, get_gke_job_name, stop_job, + replace_job_spec_image) from caliban.history.submit import submit_job_specs from caliban.history.types import (ContainerSpec, ExperimentGroup, Experiment, JobSpec, Job, Platform, JobStatus, Platform) -from caliban.gke.utils import user_verify, credentials +from caliban.gke.util import user_verify, credentials from caliban.docker import build_image, push_uuid_tag, execute_jobs diff --git a/caliban/history/utils.py b/caliban/history/util.py similarity index 98% rename from caliban/history/utils.py rename to caliban/history/util.py index 5b7737c..917d9de 100644 --- a/caliban/history/utils.py +++ b/caliban/history/util.py @@ -17,25 +17,24 @@ import os import sys from contextlib import contextmanager -from typing import Optional, Dict, Any, List from copy import deepcopy +from typing import Any, Dict, List, Optional + from absl import logging from blessings import Terminal from googleapiclient import discovery - from sqlalchemy import create_engine from sqlalchemy.engine.base import Engine from sqlalchemy.exc import OperationalError from sqlalchemy.orm import Session, sessionmaker import caliban.config as conf +from caliban.cloud.types import JobStatus as CloudStatus from caliban.gke.cluster import Cluster -from caliban.gke.utils import default_credentials from caliban.gke.types import JobStatus as GkeStatus -from caliban.cloud.types import JobStatus as CloudStatus -from caliban.history.types import (init_db, Job, JobStatus, Platform, - ContainerSpec, Experiment, ExperimentGroup, - JobSpec) +from caliban.gke.util import default_credentials +from caliban.history.types import (ContainerSpec, Experiment, ExperimentGroup, + Job, JobSpec, JobStatus, Platform, init_db) DB_URL_ENV = 'CALIBAN_DB_URL' MEMORY_DB_URL = 'sqlite:///:memory:' diff --git a/caliban/platform/cloud/core.py b/caliban/platform/cloud/core.py index 543c115..b6c4829 100644 --- a/caliban/platform/cloud/core.py +++ b/caliban/platform/cloud/core.py @@ -36,9 +36,8 @@ import caliban.docker as d import caliban.history.types as ht import caliban.util as u -from caliban.history.utils import (create_experiments, generate_container_spec, - get_mem_engine, get_sql_engine, - session_scope) +from caliban.history.util import (create_experiments, generate_container_spec, + get_mem_engine, get_sql_engine, session_scope) t = Terminal() diff --git a/caliban/platform/gke/cli.py b/caliban/platform/gke/cli.py index e42422d..a71f1d7 100644 --- a/caliban/platform/gke/cli.py +++ b/caliban/platform/gke/cli.py @@ -29,13 +29,12 @@ import caliban.cli as cli import caliban.config as conf import caliban.gke.constants as k -import caliban.gke.utils as utils +import caliban.gke.util as util import caliban.util as u from caliban.cloud.core import generate_image_tag from caliban.gke.cluster import Cluster -from caliban.history.utils import (create_experiments, generate_container_spec, - get_mem_engine, get_sql_engine, - session_scope) +from caliban.history.util import (create_experiments, generate_container_spec, + get_mem_engine, get_sql_engine, session_scope) # ---------------------------------------------------------------------------- @@ -46,7 +45,7 @@ def wrapper(args: dict): project_id = args.get('project_id') creds_file = args.get('cloud_key') - creds_data = utils.credentials(creds_file) + creds_data = util.credentials(creds_file) creds = creds_data.credentials if project_id is None: @@ -105,8 +104,8 @@ def _check_for_existing_cluster(cluster_name: str, project_id: str, for c in clusters: logging.info(c) - return utils.user_verify('Do you really want to create a new cluster?', - default=False) + return util.user_verify('Do you really want to create a new cluster?', + default=False) # ---------------------------------------------------------------------------- @@ -126,11 +125,11 @@ def _export_jobs(export: str, jobs: List[V1Job]) -> bool: """ if len(jobs) == 1: - return utils.export_job(jobs[0], export) + return util.export_job(jobs[0], export) else: base, ext = os.path.splitext(export) for i, j in enumerate(jobs): - if not utils.export_job(j, f'{base}_{i}{ext}'): + if not util.export_job(j, f'{base}_{i}{ext}'): return False return True @@ -149,7 +148,7 @@ def _cluster_create(args: dict, project_id: str, creds: Credentials) -> None: dry_run = args['dry_run'] cluster_name = args['cluster_name'] or k.DEFAULT_CLUSTER_NAME zone = args['zone'] - dashboard_url = utils.dashboard_cluster_url(cluster_name, zone, project_id) + dashboard_url = util.dashboard_cluster_url(cluster_name, zone, project_id) release_channel = args['release_channel'] single_zone = args['single_zone'] @@ -208,9 +207,9 @@ def _cluster_delete(args: dict, cluster: Cluster) -> None: None """ - if utils.user_verify('Are you sure you want to delete {}?'.format( + if util.user_verify('Are you sure you want to delete {}?'.format( cluster.name), - default=False): + default=False): cluster.delete() return @@ -433,7 +432,7 @@ def _job_submit(args: dict, cluster: Cluster) -> None: specs = list( cluster.create_simple_experiment_job_specs( - name=utils.sanitize_job_name(job_name), + name=util.sanitize_job_name(job_name), image=image_tag, min_cpu=min_cpu, min_mem=min_mem, @@ -483,7 +482,7 @@ def _job_submit_file(args: dict, cluster: Cluster) -> None: job_file = args['job_file'] - job_spec = utils.parse_job_file(job_file) + job_spec = util.parse_job_file(job_file) if job_spec is None: logging.error('error parsing job file {}'.format(job_file)) return diff --git a/caliban/platform/gke/cluster.py b/caliban/platform/gke/cluster.py index ed256a2..aaa987e 100644 --- a/caliban/platform/gke/cluster.py +++ b/caliban/platform/gke/cluster.py @@ -39,11 +39,11 @@ import caliban.config as conf import caliban.gke.constants as k -import caliban.gke.utils as utils +import caliban.gke.util as util from caliban.cloud.types import (GPU, TPU, Accelerator, GPUSpec, MachineType, TPUSpec) from caliban.gke.types import NodeImage, OpStatus, ReleaseChannel -from caliban.gke.utils import trap +from caliban.gke.util import trap from caliban.history.types import Experiment, Job, JobSpec, JobStatus, Platform urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) @@ -265,8 +265,8 @@ def _set_gke_cluster(self) -> bool: logging.error('error getting cluster management client') return False - cluster_list = utils.get_gke_clusters(self._cluster_client, self.project_id, - self.zone) + cluster_list = util.get_gke_clusters(self._cluster_client, self.project_id, + self.zone) if cluster_list is None: return False if len(cluster_list) < 1: @@ -313,7 +313,7 @@ def list(project_id: str, logging.error('error getting cluster management client') return False - clusters = utils.get_gke_clusters(client, project_id, zone) + clusters = util.get_gke_clusters(client, project_id, zone) return [c.name for c in clusters] if clusters is not None else None # -------------------------------------------------------------------------- @@ -603,7 +603,7 @@ def create_v1job( ) -> V1Job: '''creates a V1Job from a JobSpec, a job name, and an optional set of labels''' - name = utils.sanitize_job_name(name) + name = util.sanitize_job_name(name) job_metadata = V1ObjectMeta(generate_name=name + '-', labels=labels) return V1Job(api_version=k.BATCH_V1_VERSION, @@ -861,7 +861,7 @@ def convert_accel_spec( @connected(None) def dashboard_url(self) -> str: """returns gke dashboard url for this cluster""" - return utils.dashboard_cluster_url(self.name, self.zone, self.project_id) + return util.dashboard_cluster_url(self.name, self.zone, self.project_id) # -------------------------------------------------------------------------- @connected(None) @@ -891,7 +891,7 @@ def get_tpu_types(self) -> Optional[List[TPUSpec]]: list of supported tpu types on success, None otherwise """ - return utils.get_zone_tpu_types(self._tpu_api, self.project_id, self.zone) + return util.get_zone_tpu_types(self._tpu_api, self.project_id, self.zone) # -------------------------------------------------------------------------- @connected(None) @@ -958,14 +958,13 @@ def validate_gpu_spec(self, gpu_spec: Optional[GPUSpec]) -> bool: credentials=self.credentials, cache_discovery=False) - zone_gpus = utils.get_zone_gpu_types(compute_api, self.project_id, - self.zone) + zone_gpus = util.get_zone_gpu_types(compute_api, self.project_id, self.zone) if zone_gpus is None: return False gpu_limits = dict([(x.gpu, x.count) for x in zone_gpus]) - if not utils.validate_gpu_spec_against_limits(gpu_spec, gpu_limits, 'zone'): + if not util.validate_gpu_spec_against_limits(gpu_spec, gpu_limits, 'zone'): return False # ------------------------------------------------------------------------ @@ -975,8 +974,8 @@ def validate_gpu_spec(self, gpu_spec: Optional[GPUSpec]) -> bool: return False gpu_limits = dict([(x.gpu, x.count) for x in available_gpu]) - if not utils.validate_gpu_spec_against_limits(gpu_spec, gpu_limits, - 'cluster'): + if not util.validate_gpu_spec_against_limits(gpu_spec, gpu_limits, + 'cluster'): return False return True @@ -1061,7 +1060,7 @@ def get_tpu_drivers(self) -> Optional[List[str]]: list of supported tpu drivers on success, None otherwise """ - return utils.get_tpu_drivers(self._tpu_api, self.project_id, self.zone) + return util.get_tpu_drivers(self._tpu_api, self.project_id, self.zone) # -------------------------------------------------------------------------- @connected(None) @@ -1125,8 +1124,8 @@ def create_request(cluster_api: discovery.Resource, creds: Credentials, credentials=creds, cache_discovery=False) - resource_limits = utils.generate_resource_limits(compute_api, project_id, - region) + resource_limits = util.generate_resource_limits(compute_api, project_id, + region) if resource_limits is None: logging.error('error generating resource limits') @@ -1135,7 +1134,7 @@ def create_request(cluster_api: discovery.Resource, creds: Credentials, if single_zone: node_zones = [zone] else: - node_zones = utils.get_zones_in_region(compute_api, project_id, region) + node_zones = util.get_zones_in_region(compute_api, project_id, region) if node_zones is None: logging.error('error getting zones for region {}'.format(region)) @@ -1170,7 +1169,7 @@ def create(cluster_api: discovery.Resource, creds: Credentials, Cluster instance on success, None otherwise ''' - daemonset_url = utils.nvidia_daemonset_url(NodeImage.COS) + daemonset_url = util.nvidia_daemonset_url(NodeImage.COS) body = json.loads(request.body) zone = body['cluster']['zone'] cluster_name = body['cluster']['name'] @@ -1184,7 +1183,7 @@ def create(cluster_api: discovery.Resource, creds: Credentials, # wait for creation operation to complete operation_name = rsp['name'] - rsp = utils.wait_for_operation( + rsp = util.wait_for_operation( cluster_api, 'projects/{}/locations/{}/operations/{}'.format(project_id, zone, operation_name)) diff --git a/caliban/platform/gke/utils.py b/caliban/platform/gke/util.py similarity index 100% rename from caliban/platform/gke/utils.py rename to caliban/platform/gke/util.py diff --git a/caliban/platform/notebook.py b/caliban/platform/notebook.py index 88c4d02..9da052d 100644 --- a/caliban/platform/notebook.py +++ b/caliban/platform/notebook.py @@ -37,9 +37,8 @@ import caliban.config as c import caliban.util as u from caliban.history.types import Experiment, Job, JobSpec, JobStatus, Platform -from caliban.history.utils import (create_experiments, generate_container_spec, - get_mem_engine, get_sql_engine, - session_scope) +from caliban.history.util import (create_experiments, generate_container_spec, + get_mem_engine, get_sql_engine, session_scope) t = Terminal() diff --git a/caliban/platform/run.py b/caliban/platform/run.py index 88c4d02..9da052d 100644 --- a/caliban/platform/run.py +++ b/caliban/platform/run.py @@ -37,9 +37,8 @@ import caliban.config as c import caliban.util as u from caliban.history.types import Experiment, Job, JobSpec, JobStatus, Platform -from caliban.history.utils import (create_experiments, generate_container_spec, - get_mem_engine, get_sql_engine, - session_scope) +from caliban.history.util import (create_experiments, generate_container_spec, + get_mem_engine, get_sql_engine, session_scope) t = Terminal() diff --git a/caliban/platform/shell.py b/caliban/platform/shell.py index 88c4d02..9da052d 100644 --- a/caliban/platform/shell.py +++ b/caliban/platform/shell.py @@ -37,9 +37,8 @@ import caliban.config as c import caliban.util as u from caliban.history.types import Experiment, Job, JobSpec, JobStatus, Platform -from caliban.history.utils import (create_experiments, generate_container_spec, - get_mem_engine, get_sql_engine, - session_scope) +from caliban.history.util import (create_experiments, generate_container_spec, + get_mem_engine, get_sql_engine, session_scope) t = Terminal() diff --git a/caliban/types.py b/caliban/types.py deleted file mode 100644 index a6d75f5..0000000 --- a/caliban/types.py +++ /dev/null @@ -1,46 +0,0 @@ -#!/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. -'''general caliban types''' - -from typing import Union, Optional, TypeVar, Generic, Any, Type, NewType - - -class Ignored(): - '''ignored''' - - -Ignore = Ignored() -''' This is a simple TypeVar that is useful when you want to use - typing.Optional, but None is a valid type that you need to handle. - For example: - def foo(x: Union[Optional[str], Ignored] = Ignore): - if x == Ignore: - print('ignored') - else: - print(x) - - foo() - foo(None) - foo('bar') - - returns: - ignored - None - bar - -Note that mypy sometimes will not be able to handle the above conditional, -so you may have to use isinstance(x, Ignored) instead of (x == Ignore). -''' diff --git a/tests/caliban/cloud/__init__.py b/tests/caliban/cloud/__init__.py deleted file mode 100644 index 2aa960d..0000000 --- a/tests/caliban/cloud/__init__.py +++ /dev/null @@ -1,31 +0,0 @@ -#!/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. - -#!/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/tests/caliban/config/__init__.py b/tests/caliban/config/__init__.py new file mode 100644 index 0000000..79c6a2f --- /dev/null +++ b/tests/caliban/config/__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/tests/caliban/docker/__init__.py b/tests/caliban/docker/__init__.py new file mode 100644 index 0000000..79c6a2f --- /dev/null +++ b/tests/caliban/docker/__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/tests/caliban/test_docker.py b/tests/caliban/docker/test_docker.py similarity index 100% rename from tests/caliban/test_docker.py rename to tests/caliban/docker/test_docker.py diff --git a/tests/caliban/gke/__init__.py b/tests/caliban/gke/__init__.py deleted file mode 100644 index 2aa960d..0000000 --- a/tests/caliban/gke/__init__.py +++ /dev/null @@ -1,31 +0,0 @@ -#!/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. - -#!/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/tests/caliban/history/__init__.py b/tests/caliban/history/__init__.py index 2aa960d..79c6a2f 100644 --- a/tests/caliban/history/__init__.py +++ b/tests/caliban/history/__init__.py @@ -13,19 +13,3 @@ # 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. - -#!/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/tests/caliban/history/test_history.py b/tests/caliban/history/test_history.py index be81b67..9f24be2 100644 --- a/tests/caliban/history/test_history.py +++ b/tests/caliban/history/test_history.py @@ -1,19 +1,3 @@ -#!/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. - #!/usr/bin/python # # Copyright 2020 Google LLC @@ -31,18 +15,17 @@ # limitations under the License. """unit tests for caliban history""" -import pytest #type: ignore -# https://mypy.readthedocs.io/en/latest/jobning_mypy.html#missing-imports +from datetime import datetime from sqlalchemy.engine.base import Engine -from caliban.history.utils import session_scope, get_mem_engine, get_sql_engine -from caliban.history.types import (JobStatus, Platform, ExperimentGroup, - Experiment, JobSpec, Job, ContainerSpec) +import pytest # type: ignore +from caliban.history.types import (ContainerSpec, Experiment, ExperimentGroup, + Job, JobSpec, Platform) +from caliban.history.util import get_mem_engine, session_scope from caliban.util import current_user -import random -from datetime import datetime -from time import sleep + +# https://mypy.readthedocs.io/en/latest/jobning_mypy.html#missing-imports # we create and exist session scopes here to test persistence diff --git a/tests/caliban/platform/__init__.py b/tests/caliban/platform/__init__.py new file mode 100644 index 0000000..79c6a2f --- /dev/null +++ b/tests/caliban/platform/__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/tests/caliban/platform/cloud/__init__.py b/tests/caliban/platform/cloud/__init__.py new file mode 100644 index 0000000..79c6a2f --- /dev/null +++ b/tests/caliban/platform/cloud/__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/tests/caliban/cloud/test_types.py b/tests/caliban/platform/cloud/test_types.py similarity index 81% rename from tests/caliban/cloud/test_types.py rename to tests/caliban/platform/cloud/test_types.py index 2d6d46b..690cc6e 100644 --- a/tests/caliban/cloud/test_types.py +++ b/tests/caliban/platform/cloud/test_types.py @@ -14,26 +14,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -#!/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. - import unittest from argparse import ArgumentTypeError -from typing import (Any, Callable, Dict, Iterable, List, NamedTuple, Optional, - Set, Tuple, Union) import hypothesis.strategies as st from hypothesis import given diff --git a/tests/caliban/platform/gke/__init__.py b/tests/caliban/platform/gke/__init__.py new file mode 100644 index 0000000..79c6a2f --- /dev/null +++ b/tests/caliban/platform/gke/__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/tests/caliban/gke/test_types.py b/tests/caliban/platform/gke/test_types.py similarity index 65% rename from tests/caliban/gke/test_types.py rename to tests/caliban/platform/gke/test_types.py index d5eb20c..f727502 100644 --- a/tests/caliban/gke/test_types.py +++ b/tests/caliban/platform/gke/test_types.py @@ -1,19 +1,3 @@ -#!/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. - #!/usr/bin/python # # Copyright 2020 Google LLC @@ -33,9 +17,7 @@ import unittest import hypothesis.strategies as st -from hypothesis import given, settings -from typing import Dict, List, Any -import re +from hypothesis import given from caliban.gke.types import ReleaseChannel @@ -54,5 +36,3 @@ def test_release_channel(self, invalid: str, valid: ReleaseChannel): x = ReleaseChannel(invalid) self.assertEqual(valid, ReleaseChannel(valid.value)) - - return diff --git a/tests/caliban/gke/test_utils.py b/tests/caliban/platform/gke/test_utils.py similarity index 98% rename from tests/caliban/gke/test_utils.py rename to tests/caliban/platform/gke/test_utils.py index cc21633..f6b307d 100644 --- a/tests/caliban/gke/test_utils.py +++ b/tests/caliban/platform/gke/test_utils.py @@ -14,23 +14,19 @@ # See the License for the specific language governing permissions and # limitations under the License. """unit tests for gke utilities""" +import random import unittest +from typing import List from unittest import mock -import random import hypothesis.strategies as st from hypothesis import given, settings -from typing import Dict, List, Any -import re -import random -import pprint as pp import caliban.cloud.types as ct -import caliban.gke -import caliban.gke.utils as utils import caliban.gke.constants as k +import caliban.gke.utils as utils +from caliban.gke.types import NodeImage, OpStatus from caliban.gke.utils import trap -from caliban.gke.types import NodeImage, OpStatus, ReleaseChannel # ---------------------------------------------------------------------------- @@ -622,5 +618,3 @@ def _invalid(): # normal execution api.execute = _normal self.assertEqual(zones, utils.get_zones_in_region(api, 'p', region)) - - return diff --git a/tests/caliban/util/__init__.py b/tests/caliban/util/__init__.py new file mode 100644 index 0000000..79c6a2f --- /dev/null +++ b/tests/caliban/util/__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/tests/caliban/test_util.py b/tests/caliban/util/test_util.py similarity index 100% rename from tests/caliban/test_util.py rename to tests/caliban/util/test_util.py From 309d380400e54bbf810eac04838fae90450c6f81 Mon Sep 17 00:00:00 2001 From: Sam Ritchie Date: Fri, 10 Jul 2020 09:24:46 -0600 Subject: [PATCH 03/15] get tests restructured --- caliban/config/__init__.py | 2 +- caliban/config/experiment.py | 4 +- tests/caliban/config/test_config.py | 2 +- tests/caliban/config/test_experiment.py | 11 ++ tests/caliban/util/test_util.py | 148 ------------------------ 5 files changed, 15 insertions(+), 152 deletions(-) diff --git a/caliban/config/__init__.py b/caliban/config/__init__.py index 28e00a4..ac60130 100644 --- a/caliban/config/__init__.py +++ b/caliban/config/__init__.py @@ -28,7 +28,7 @@ import commentjson import yaml -import caliban.cloud.types as ct +import caliban.platform.cloud.types as ct class JobMode(str, Enum): diff --git a/caliban/config/experiment.py b/caliban/config/experiment.py index 0fcc1bc..c0fecd5 100644 --- a/caliban/config/experiment.py +++ b/caliban/config/experiment.py @@ -148,8 +148,8 @@ def expand_experiment_config(items: ExpConf) -> List[Experiment]: itertools.chain.from_iterable( [expand_experiment_config(m) for m in items])) - tupleized_items = u.tupleize_dict(items) - return [u.expand_compound_dict(d) for d in u.dict_product(tupleized_items)] + tupleized_items = tupleize_dict(items) + return [expand_compound_dict(d) for d in u.dict_product(tupleized_items)] def validate_compound_keys(m: ExpConf) -> ExpConf: diff --git a/tests/caliban/config/test_config.py b/tests/caliban/config/test_config.py index aa49498..126c21d 100644 --- a/tests/caliban/config/test_config.py +++ b/tests/caliban/config/test_config.py @@ -17,7 +17,7 @@ import os from argparse import ArgumentTypeError -import caliban.cloud.types as ct +import caliban.platform.cloud.types as ct import caliban.config as c import pytest diff --git a/tests/caliban/config/test_experiment.py b/tests/caliban/config/test_experiment.py index 30f134d..5ca1e3d 100644 --- a/tests/caliban/config/test_experiment.py +++ b/tests/caliban/config/test_experiment.py @@ -17,6 +17,7 @@ from argparse import ArgumentTypeError import caliban.config.experiment as c +import caliban.util as u import pytest @@ -101,6 +102,11 @@ def test_expand_experiment_config(): def test_compound_key_handling(): + """tests the full assembly line transforming a configuration dictionary + including compound keys into a list of dictionaries for passing to the + script + + """ tests = [{ 'input': { '[a,b]': [['c', 'd'], ['e', 'f']] @@ -226,4 +232,9 @@ def test_compound_key_handling(): }] for test in tests: + assert test['after_tupleization'] == c.tupleize_dict(test['input']) + assert test['after_expansion'] == list( + c.expand_compound_dict(test['after_dictproduct'])) assert test['after_expansion'] == c.expand_experiment_config(test['input']) + assert test['after_dictproduct'] == list( + u.dict_product(test['after_tupleization'])) diff --git a/tests/caliban/util/test_util.py b/tests/caliban/util/test_util.py index d46719b..26bf24b 100644 --- a/tests/caliban/util/test_util.py +++ b/tests/caliban/util/test_util.py @@ -154,154 +154,6 @@ def test_dict_product(self): self.assertListEqual(result, expected) - def test_compound_key_handling(self): - """ tests the full assembly line transforming a configuration dictionary - including compound keys into a list of dictionaries for passing to the script - """ - - tests = [{ - 'input': { - '[a,b]': [['c', 'd'], ['e', 'f']] - }, - 'after_tupleization': { - ('a', 'b'): [('c', 'd'), ('e', 'f')] - }, - 'after_dictproduct': [{ - ('a', 'b'): ('c', 'd') - }, { - ('a', 'b'): ('e', 'f') - }], - 'after_expansion': [{ - 'a': 'c', - 'b': 'd' - }, { - 'a': 'e', - 'b': 'f' - }] - }, { - 'input': { - '[a,b]': ['c', 'd'] - }, - 'after_tupleization': { - ('a', 'b'): ('c', 'd') - }, - 'after_dictproduct': [{ - ('a', 'b'): ('c', 'd') - }], - 'after_expansion': [{ - 'a': 'c', - 'b': 'd' - }] - }, { - 'input': { - 'hi': 'there', - '[k1,k2]': [['v1a', 'v2a'], ['v1b', 'v2b']] - }, - 'after_tupleization': { - 'hi': 'there', - ('k1', 'k2'): [('v1a', 'v2a'), ('v1b', 'v2b')] - }, - 'after_dictproduct': [{ - 'hi': 'there', - ('k1', 'k2'): ('v1a', 'v2a') - }, { - 'hi': 'there', - ('k1', 'k2'): ('v1b', 'v2b') - }], - 'after_expansion': [{ - 'hi': 'there', - 'k1': 'v1a', - 'k2': 'v2a' - }, { - 'hi': 'there', - 'k1': 'v1b', - 'k2': 'v2b' - }] - }, { - 'input': { - 'hi': 'there', - '[a,b]': ['c', 'd'] - }, - 'after_tupleization': { - 'hi': 'there', - ('a', 'b'): ('c', 'd') - }, - 'after_dictproduct': [{ - 'hi': 'there', - ('a', 'b'): ('c', 'd') - }], - 'after_expansion': [{ - 'hi': 'there', - 'a': 'c', - 'b': 'd' - }] - }, { - 'input': { - '[a,b]': [0, 1] - }, - 'after_tupleization': { - ('a', 'b'): (0, 1) - }, - 'after_dictproduct': [{ - ('a', 'b'): (0, 1) - }], - 'after_expansion': [{ - 'a': 0, - 'b': 1 - }] - }, { - 'input': { - '[a,b]': [[0, 1]] - }, - 'after_tupleization': { - ('a', 'b'): [(0, 1)] - }, - 'after_dictproduct': [{ - ('a', 'b'): (0, 1) - }], - 'after_expansion': [{ - 'a': 0, - 'b': 1 - }] - }, { - 'input': { - 'hi': 'blueshift', - '[a,b]': [[0, 1]] - }, - 'after_tupleization': { - 'hi': 'blueshift', - ('a', 'b'): [(0, 1)] - }, - 'after_dictproduct': [{ - 'hi': 'blueshift', - ('a', 'b'): (0, 1) - }], - 'after_expansion': [{ - 'hi': 'blueshift', - 'a': 0, - 'b': 1 - }] - }] - - def check_tupleization(test_dict): - self.assertDictEqual(test_dict['after_tupleization'], - u.tupleize_dict(test_dict['input'])) - - def check_dictproduct(test_dict): - self.assertListEqual( - test_dict['after_dictproduct'], - list(u.dict_product(test_dict['after_tupleization']))) - - def check_expansion(test_dict): - self.assertListEqual( - test_dict['after_expansion'], - list(u.expand_compound_dict(test_dict['after_dictproduct']))) - - for test in tests: - check_tupleization(test) - check_dictproduct(test) - check_expansion(test) - @given(st.dictionaries(st.text(), st.text()), st.dictionaries(st.text(), st.text())) def test_merge(self, m1, m2): From a07ef6b4fb11a5d0a5871439da0b59449a80b3e1 Mon Sep 17 00:00:00 2001 From: Sam Ritchie Date: Fri, 10 Jul 2020 09:37:39 -0600 Subject: [PATCH 04/15] more conversion --- caliban/builder.py | 10 +-- caliban/docker/build.py | 3 +- caliban/docker/push.py | 16 ++--- caliban/platform/notebook.py | 82 +++------------------- caliban/platform/run.py | 119 +------------------------------- caliban/platform/shell.py | 57 +-------------- tests/caliban/util/test_fs.py | 36 ++++++++++ tests/caliban/util/test_tqdm.py | 47 +++++++++++++ tests/caliban/util/test_util.py | 17 ----- 9 files changed, 111 insertions(+), 276 deletions(-) create mode 100644 tests/caliban/util/test_fs.py create mode 100644 tests/caliban/util/test_tqdm.py diff --git a/caliban/builder.py b/caliban/builder.py index 893438c..ec333a8 100644 --- a/caliban/builder.py +++ b/caliban/builder.py @@ -23,22 +23,18 @@ import json import os import subprocess -import sys from enum import Enum from pathlib import Path from typing import (Any, Callable, Dict, Iterable, List, NamedTuple, NewType, Optional, Union) -import tqdm from absl import logging from blessings import Terminal -from tqdm.utils import _screen_shape_wrapper import caliban.config as c import caliban.util as u -from caliban.history.types import Experiment, Job, JobSpec, JobStatus, Platform -from caliban.history.util import (create_experiments, generate_container_spec, - get_mem_engine, get_sql_engine, session_scope) +import caliban.util.fs as ufs +from caliban.history.types import JobSpec t = Terminal() @@ -613,7 +609,7 @@ def build_image(job_mode: c.JobMode, logging.info("Running command: {}".format(joined_cmd)) try: - output, ret_code = u.capture_stdout(cmd, input_str=dockerfile) + output, ret_code = ufs.capture_stdout(cmd, input_str=dockerfile) if ret_code == 0: return docker_image_id(output) else: diff --git a/caliban/docker/build.py b/caliban/docker/build.py index 1d16aa5..3185eae 100644 --- a/caliban/docker/build.py +++ b/caliban/docker/build.py @@ -33,6 +33,7 @@ import caliban.config as c import caliban.util as u +import caliban.util.fs as ufs t = Terminal() @@ -607,7 +608,7 @@ def build_image(job_mode: c.JobMode, logging.info("Running command: {}".format(joined_cmd)) try: - output, ret_code = u.capture_stdout(cmd, input_str=dockerfile) + output, ret_code = ufs.capture_stdout(cmd, input_str=dockerfile) if ret_code == 0: return docker_image_id(output) else: diff --git a/caliban/docker/push.py b/caliban/docker/push.py index 893438c..e4acd5a 100644 --- a/caliban/docker/push.py +++ b/caliban/docker/push.py @@ -23,22 +23,18 @@ import json import os import subprocess -import sys from enum import Enum from pathlib import Path from typing import (Any, Callable, Dict, Iterable, List, NamedTuple, NewType, Optional, Union) -import tqdm from absl import logging from blessings import Terminal -from tqdm.utils import _screen_shape_wrapper import caliban.config as c import caliban.util as u -from caliban.history.types import Experiment, Job, JobSpec, JobStatus, Platform -from caliban.history.util import (create_experiments, generate_container_spec, - get_mem_engine, get_sql_engine, session_scope) +import caliban.util.fs as ufs +from caliban.history.types import JobSpec t = Terminal() @@ -598,9 +594,9 @@ def build_image(job_mode: c.JobMode, the problem. """ - with u.TempCopy(credentials_path, - tmp_name=".caliban_default_creds.json") as creds: - with u.TempCopy(adc_path, tmp_name=".caliban_adc_creds.json") as adc: + with ufs.TempCopy(credentials_path, + tmp_name=".caliban_default_creds.json") as creds: + with ufs.TempCopy(adc_path, tmp_name=".caliban_adc_creds.json") as adc: cache_args = ["--no-cache"] if no_cache else [] cmd = ["docker", "build"] + cache_args + ["--rm", "-f-", build_path] @@ -613,7 +609,7 @@ def build_image(job_mode: c.JobMode, logging.info("Running command: {}".format(joined_cmd)) try: - output, ret_code = u.capture_stdout(cmd, input_str=dockerfile) + output, ret_code = ufs.capture_stdout(cmd, input_str=dockerfile) if ret_code == 0: return docker_image_id(output) else: diff --git a/caliban/platform/notebook.py b/caliban/platform/notebook.py index 9da052d..cfdcf46 100644 --- a/caliban/platform/notebook.py +++ b/caliban/platform/notebook.py @@ -36,6 +36,8 @@ import caliban.config as c import caliban.util as u +import caliban.util.fs as ufs +import caliban.platform.shell as ps from caliban.history.types import Experiment, Job, JobSpec, JobStatus, Platform from caliban.history.util import (create_experiments, generate_container_spec, get_mem_engine, get_sql_engine, session_scope) @@ -613,7 +615,7 @@ def build_image(job_mode: c.JobMode, logging.info("Running command: {}".format(joined_cmd)) try: - output, ret_code = u.capture_stdout(cmd, input_str=dockerfile) + output, ret_code = ufs.capture_stdout(cmd, input_str=dockerfile) if ret_code == 0: return docker_image_id(output) else: @@ -808,7 +810,7 @@ def execute_jobs( command = job_spec.spec['command'] logging.info(f'Running command: {" ".join(command)}') if not dry_run: - _, ret_code = u.capture_stdout(command, "", u.TqdmFile(sys.stderr)) + _, ret_code = ufs.capture_stdout(command, "", u.TqdmFile(sys.stderr)) else: ret_code = 0 j = Job(spec=job_spec, @@ -941,68 +943,6 @@ def run(job_mode: c.JobMode, return None -def run_interactive(job_mode: c.JobMode, - workdir: Optional[str] = None, - image_id: Optional[str] = None, - run_args: Optional[List[str]] = None, - mount_home: Optional[bool] = None, - shell: Optional[Shell] = None, - entrypoint: Optional[str] = None, - entrypoint_args: Optional[List[str]] = None, - **build_image_kwargs) -> None: - """Start a live shell in the terminal, with all dependencies installed and the - current working directory (and optionally the user's home directory) mounted. - - Keyword args: - - - job_mode: c.JobMode. - - image_id: ID of the image to run. Supplying this will skip an image build. - - run_args: extra arguments to supply to `docker run`. - - mount_home: if true, mounts the user's $HOME directory into the container - to `/home/$USERNAME`. If False, nothing. - - shell: name of the shell to install into the container. Also configures the - entrypoint if that's not supplied. - - entrypoint: command to run. Defaults to the executable command for the - supplied shell. - - entrypoint_args: extra arguments to supply to the entrypoint. - - any extra kwargs supplied are passed through to build_image. - - """ - if workdir is None: - workdir = DEFAULT_WORKDIR - - if run_args is None: - run_args = [] - - if entrypoint_args is None: - entrypoint_args = [] - - if mount_home is None: - mount_home = True - - if shell is None: - # Only set a default shell if we're also mounting the home volume. - # Otherwise a custom shell won't have access to the user's profile. - shell = default_shell() if mount_home else Shell.bash - - if entrypoint is None: - entrypoint = SHELL_DICT[shell].executable - - interactive_run_args = _interactive_opts(workdir) + [ - "-it", \ - "--entrypoint", entrypoint - ] + _home_mount_cmds(mount_home) + run_args - - run(job_mode=job_mode, - run_args=interactive_run_args, - script_args=entrypoint_args, - image_id=image_id, - shell=shell, - workdir=workdir, - **build_image_kwargs) - - def run_notebook(job_mode: c.JobMode, port: Optional[int] = None, lab: Optional[bool] = None, @@ -1046,10 +986,10 @@ def run_notebook(job_mode: c.JobMode, ] docker_args = ["-p", "{}:{}".format(port, port)] + run_args - run_interactive(job_mode, - entrypoint="/opt/conda/envs/caliban/bin/python", - entrypoint_args=jupyter_args, - run_args=docker_args, - inject_notebook=inject_arg, - jupyter_version=version, - **run_interactive_kwargs) + ps.run_interactive(job_mode, + entrypoint="/opt/conda/envs/caliban/bin/python", + entrypoint_args=jupyter_args, + run_args=docker_args, + inject_notebook=inject_arg, + jupyter_version=version, + **run_interactive_kwargs) diff --git a/caliban/platform/run.py b/caliban/platform/run.py index 9da052d..0dc0068 100644 --- a/caliban/platform/run.py +++ b/caliban/platform/run.py @@ -36,6 +36,7 @@ import caliban.config as c import caliban.util as u +import caliban.util.fs as ufs from caliban.history.types import Experiment, Job, JobSpec, JobStatus, Platform from caliban.history.util import (create_experiments, generate_container_spec, get_mem_engine, get_sql_engine, session_scope) @@ -613,7 +614,7 @@ def build_image(job_mode: c.JobMode, logging.info("Running command: {}".format(joined_cmd)) try: - output, ret_code = u.capture_stdout(cmd, input_str=dockerfile) + output, ret_code = ufs.capture_stdout(cmd, input_str=dockerfile) if ret_code == 0: return docker_image_id(output) else: @@ -808,7 +809,7 @@ def execute_jobs( command = job_spec.spec['command'] logging.info(f'Running command: {" ".join(command)}') if not dry_run: - _, ret_code = u.capture_stdout(command, "", u.TqdmFile(sys.stderr)) + _, ret_code = ufs.capture_stdout(command, "", u.TqdmFile(sys.stderr)) else: ret_code = 0 j = Job(spec=job_spec, @@ -939,117 +940,3 @@ def run(job_mode: c.JobMode, logging.info("Running command: {}".format(' '.join(command))) subprocess.call(command) return None - - -def run_interactive(job_mode: c.JobMode, - workdir: Optional[str] = None, - image_id: Optional[str] = None, - run_args: Optional[List[str]] = None, - mount_home: Optional[bool] = None, - shell: Optional[Shell] = None, - entrypoint: Optional[str] = None, - entrypoint_args: Optional[List[str]] = None, - **build_image_kwargs) -> None: - """Start a live shell in the terminal, with all dependencies installed and the - current working directory (and optionally the user's home directory) mounted. - - Keyword args: - - - job_mode: c.JobMode. - - image_id: ID of the image to run. Supplying this will skip an image build. - - run_args: extra arguments to supply to `docker run`. - - mount_home: if true, mounts the user's $HOME directory into the container - to `/home/$USERNAME`. If False, nothing. - - shell: name of the shell to install into the container. Also configures the - entrypoint if that's not supplied. - - entrypoint: command to run. Defaults to the executable command for the - supplied shell. - - entrypoint_args: extra arguments to supply to the entrypoint. - - any extra kwargs supplied are passed through to build_image. - - """ - if workdir is None: - workdir = DEFAULT_WORKDIR - - if run_args is None: - run_args = [] - - if entrypoint_args is None: - entrypoint_args = [] - - if mount_home is None: - mount_home = True - - if shell is None: - # Only set a default shell if we're also mounting the home volume. - # Otherwise a custom shell won't have access to the user's profile. - shell = default_shell() if mount_home else Shell.bash - - if entrypoint is None: - entrypoint = SHELL_DICT[shell].executable - - interactive_run_args = _interactive_opts(workdir) + [ - "-it", \ - "--entrypoint", entrypoint - ] + _home_mount_cmds(mount_home) + run_args - - run(job_mode=job_mode, - run_args=interactive_run_args, - script_args=entrypoint_args, - image_id=image_id, - shell=shell, - workdir=workdir, - **build_image_kwargs) - - -def run_notebook(job_mode: c.JobMode, - port: Optional[int] = None, - lab: Optional[bool] = None, - version: Optional[bool] = None, - run_args: Optional[List[str]] = None, - **run_interactive_kwargs) -> None: - """Start a notebook in the current working directory; the process will run - inside of a Docker container that's identical to the environment available to - Cloud jobs that are submitted by `caliban cloud`, or local jobs run with - `caliban run.` - - if you pass mount_home=True your jupyter settings will persist across calls. - - Keyword args: - - - port: the port to pass to Jupyter when it boots, useful if you have - multiple instances running on one machine. - - lab: if True, starts jupyter lab, else jupyter notebook. - - version: explicit Jupyter version to install. - - run_interactive_kwargs are all extra arguments taken by run_interactive. - - """ - - if port is None: - port = u.next_free_port(8888) - - if lab is None: - lab = False - - if run_args is None: - run_args = [] - - inject_arg = NotebookInstall.lab if lab else NotebookInstall.jupyter - jupyter_cmd = "lab" if lab else "notebook" - jupyter_args = [ - "-m", "jupyter", jupyter_cmd, \ - "--ip=0.0.0.0", \ - "--port={}".format(port), \ - "--no-browser" - ] - docker_args = ["-p", "{}:{}".format(port, port)] + run_args - - run_interactive(job_mode, - entrypoint="/opt/conda/envs/caliban/bin/python", - entrypoint_args=jupyter_args, - run_args=docker_args, - inject_notebook=inject_arg, - jupyter_version=version, - **run_interactive_kwargs) diff --git a/caliban/platform/shell.py b/caliban/platform/shell.py index 9da052d..8655dfe 100644 --- a/caliban/platform/shell.py +++ b/caliban/platform/shell.py @@ -36,6 +36,7 @@ import caliban.config as c import caliban.util as u +import caliban.util.fs as ufs from caliban.history.types import Experiment, Job, JobSpec, JobStatus, Platform from caliban.history.util import (create_experiments, generate_container_spec, get_mem_engine, get_sql_engine, session_scope) @@ -613,7 +614,7 @@ def build_image(job_mode: c.JobMode, logging.info("Running command: {}".format(joined_cmd)) try: - output, ret_code = u.capture_stdout(cmd, input_str=dockerfile) + output, ret_code = ufs.capture_stdout(cmd, input_str=dockerfile) if ret_code == 0: return docker_image_id(output) else: @@ -808,7 +809,7 @@ def execute_jobs( command = job_spec.spec['command'] logging.info(f'Running command: {" ".join(command)}') if not dry_run: - _, ret_code = u.capture_stdout(command, "", u.TqdmFile(sys.stderr)) + _, ret_code = ufs.capture_stdout(command, "", u.TqdmFile(sys.stderr)) else: ret_code = 0 j = Job(spec=job_spec, @@ -1001,55 +1002,3 @@ def run_interactive(job_mode: c.JobMode, shell=shell, workdir=workdir, **build_image_kwargs) - - -def run_notebook(job_mode: c.JobMode, - port: Optional[int] = None, - lab: Optional[bool] = None, - version: Optional[bool] = None, - run_args: Optional[List[str]] = None, - **run_interactive_kwargs) -> None: - """Start a notebook in the current working directory; the process will run - inside of a Docker container that's identical to the environment available to - Cloud jobs that are submitted by `caliban cloud`, or local jobs run with - `caliban run.` - - if you pass mount_home=True your jupyter settings will persist across calls. - - Keyword args: - - - port: the port to pass to Jupyter when it boots, useful if you have - multiple instances running on one machine. - - lab: if True, starts jupyter lab, else jupyter notebook. - - version: explicit Jupyter version to install. - - run_interactive_kwargs are all extra arguments taken by run_interactive. - - """ - - if port is None: - port = u.next_free_port(8888) - - if lab is None: - lab = False - - if run_args is None: - run_args = [] - - inject_arg = NotebookInstall.lab if lab else NotebookInstall.jupyter - jupyter_cmd = "lab" if lab else "notebook" - jupyter_args = [ - "-m", "jupyter", jupyter_cmd, \ - "--ip=0.0.0.0", \ - "--port={}".format(port), \ - "--no-browser" - ] - docker_args = ["-p", "{}:{}".format(port, port)] + run_args - - run_interactive(job_mode, - entrypoint="/opt/conda/envs/caliban/bin/python", - entrypoint_args=jupyter_args, - run_args=docker_args, - inject_notebook=inject_arg, - jupyter_version=version, - **run_interactive_kwargs) diff --git a/tests/caliban/util/test_fs.py b/tests/caliban/util/test_fs.py new file mode 100644 index 0000000..b9ab88f --- /dev/null +++ b/tests/caliban/util/test_fs.py @@ -0,0 +1,36 @@ +#!/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. + +import io + +import caliban.util.fs as ufs + + +def test_capture_stdout(): + buf = io.StringIO() + ret_string, code = ufs.capture_stdout(["echo", "hello!"], file=buf) + assert code == 0 + + # Verify that the stdout is reported to the supplied file, and that it's + # captured by the function and returned correctly. + assert ret_string == "hello!\n" + assert buf.getvalue() == ret_string + + +def test_capture_stdout_input(): + ret_string, code = ufs.capture_stdout(["cat"], input_str="hello!") + assert code == 0 + assert ret_string.rstrip() == "hello!" diff --git a/tests/caliban/util/test_tqdm.py b/tests/caliban/util/test_tqdm.py new file mode 100644 index 0000000..ebd6e37 --- /dev/null +++ b/tests/caliban/util/test_tqdm.py @@ -0,0 +1,47 @@ +#!/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. + +import io + +from tqdm._utils import _term_move_up + +import caliban.util.tqdm as ut + + +def test_carriage_return(): + + def through(xs): + buf = io.StringIO() + f = ut.TqdmFile(file=buf) + + for x in xs: + f.write(x) + f.flush() + + return buf.getvalue() + + # Strings pass through tqdmfile with no newline attached. + assert through(["Yo!"]) == "Yo!" + + # Empty lines do nothing. + assert through(["", "", ""]) == "" + + # A carriage return is converted to a newline, but the next line, if it's + # written, will have the proper prefix to trigger a carriage return. + assert through(["Yo!\r"]) == "Yo!\n" + + # Boom, triggered. + assert through(["Yo!\r", "continue"]) == f"Yo!\n{_term_move_up()}\rcontinue" diff --git a/tests/caliban/util/test_util.py b/tests/caliban/util/test_util.py index 26bf24b..2c99318 100644 --- a/tests/caliban/util/test_util.py +++ b/tests/caliban/util/test_util.py @@ -36,17 +36,6 @@ def non_empty_dict(vgen): return st.dictionaries(st.text(), vgen, min_size=1) -def test_capture_stdout(): - buf = io.StringIO() - ret_string, code = u.capture_stdout(["echo", "hello!"], file=buf) - assert code == 0 - - # Verify that the stdout is reported to the supplied file, and that it's - # captured by the function and returned correctly. - assert ret_string == "hello!\n" - assert buf.getvalue() == ret_string - - def test_carriage_return(): def through(xs): @@ -73,12 +62,6 @@ def through(xs): assert through(["Yo!\r", "continue"]) == f"Yo!\n{_term_move_up()}\rcontinue" -def test_capture_stdout_input(): - ret_string, code = u.capture_stdout(["cat"], input_str="hello!") - assert code == 0 - assert ret_string.rstrip() == "hello!" - - class UtilTestSuite(unittest.TestCase): """Tests for the util package.""" From 76fe3e0212f987a50157bd756938c89bc067f697 Mon Sep 17 00:00:00 2001 From: Sam Ritchie Date: Fri, 10 Jul 2020 13:00:33 -0600 Subject: [PATCH 05/15] rename push --- caliban/builder.py | 729 --------------------- caliban/cli.py | 10 +- caliban/config/__init__.py | 2 +- caliban/docker/build.py | 6 +- caliban/docker/push.py | 678 ------------------- caliban/history/cli.py | 2 +- caliban/history/submit.py | 4 +- caliban/history/util.py | 8 +- caliban/main.py | 6 +- caliban/platform/cloud/core.py | 2 +- caliban/platform/gke/cli.py | 8 +- caliban/platform/gke/cluster.py | 12 +- caliban/platform/gke/constants.py | 4 +- caliban/platform/gke/util.py | 6 +- tests/caliban/docker/test_build.py | 23 + tests/caliban/docker/test_docker.py | 39 -- tests/caliban/docker/test_push.py | 29 + tests/caliban/platform/cloud/test_types.py | 4 +- tests/caliban/platform/gke/test_types.py | 4 +- tests/caliban/platform/gke/test_utils.py | 14 +- tests/caliban/test_builder.py | 0 tests/caliban/test_cli.py | 2 +- 22 files changed, 99 insertions(+), 1493 deletions(-) delete mode 100644 caliban/builder.py create mode 100644 tests/caliban/docker/test_build.py delete mode 100644 tests/caliban/docker/test_docker.py create mode 100644 tests/caliban/docker/test_push.py create mode 100644 tests/caliban/test_builder.py diff --git a/caliban/builder.py b/caliban/builder.py deleted file mode 100644 index ec333a8..0000000 --- a/caliban/builder.py +++ /dev/null @@ -1,729 +0,0 @@ -#!/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. -"""Functions required to interact with Docker to build and run images, shells -and notebooks in a Docker environment. - -""" - -from __future__ import absolute_import, division, print_function - -import json -import os -import subprocess -from enum import Enum -from pathlib import Path -from typing import (Any, Callable, Dict, Iterable, List, NamedTuple, NewType, - Optional, Union) - -from absl import logging -from blessings import Terminal - -import caliban.config as c -import caliban.util as u -import caliban.util.fs as ufs -from caliban.history.types import JobSpec - -t = Terminal() - -DEV_CONTAINER_ROOT = "gcr.io/blueshift-playground/blueshift" -TF_VERSIONS = {"2.2.0", "1.12.3", "1.14.0", "1.15.0"} -DEFAULT_WORKDIR = "/usr/app" -CREDS_DIR = "/.creds" -CONDA_BIN = "/opt/conda/bin/conda" - -ImageId = NewType('ImageId', str) -ArgSeq = NewType('ArgSeq', List[str]) - - -class DockerError(Exception): - """Exception that passes info on a failed Docker command.""" - - def __init__(self, message, cmd, ret_code): - super().__init__(message) - self.message = message - self.cmd = cmd - self.ret_code = ret_code - - @property - def command(self): - return " ".join(self.cmd) - - -class NotebookInstall(Enum): - """Flag to decide what to do .""" - none = 'none' - lab = 'lab' - jupyter = 'jupyter' - - def __str__(self) -> str: - return self.value - - -class Shell(Enum): - """Add new shells here and below, in SHELL_DICT.""" - bash = 'bash' - zsh = 'zsh' - - def __str__(self) -> str: - return self.value - - -# Tuple to track the information required to install and execute some custom -# shell into a container. -ShellData = NamedTuple("ShellData", [("executable", str), - ("packages", List[str])]) - - -def apt_install(*packages: str) -> str: - """Returns a command that will install the supplied list of packages without - requiring confirmation or any user interaction. - """ - package_str = ' '.join(packages) - no_prompt = "DEBIAN_FRONTEND=noninteractive" - return f"{no_prompt} apt-get install --yes --no-install-recommends {package_str}" - - -def apt_command(commands: List[str]) -> List[str]: - """Pre-and-ap-pends the supplied commands with the appropriate in-container and - cleanup command for aptitude. - - """ - update = ["apt-get update"] - cleanup = ["apt-get clean", "rm -rf /var/lib/apt/lists/*"] - return update + commands + cleanup - - -# Dict linking a particular supported shell to the data required to run and -# install the shell inside a container. -# -# : Dict[Shell, ShellData] -SHELL_DICT = { - Shell.bash: ShellData("/bin/bash", []), - Shell.zsh: ShellData("/bin/zsh", ["zsh"]) -} - - -def default_shell() -> Shell: - """Returns the shell to load into the container. Defaults to Shell.bash, but if - the user's SHELL variable refers to a supported sub-shell, returns that - instead. - - """ - ret = Shell.bash - - if "zsh" in os.environ.get("SHELL"): - ret = Shell.zsh - - return ret - - -def adc_location(home_dir: Optional[str] = None) -> str: - """Returns the location for application default credentials, INSIDE the - container (so, hardcoded unix separators), given the supplied home directory. - - """ - if home_dir is None: - home_dir = Path.home() - - return "{}/.config/gcloud/application_default_credentials.json".format( - home_dir) - - -def container_home(): - """Returns the location of the home directory inside the generated - container. - - """ - return "/home/{}".format(u.current_user()) - - -def tf_base_image(job_mode: c.JobMode, tensorflow_version: str) -> str: - """Returns the base image to use, depending on whether or not we're using a - GPU. This is JUST for building our base images for Blueshift; not for - actually using in a job. - - List of available tags: https://hub.docker.com/r/tensorflow/tensorflow/tags - - """ - if tensorflow_version not in TF_VERSIONS: - raise Exception("""{} is not a valid tensorflow version. - Try one of: {}""".format(tensorflow_version, TF_VERSIONS)) - - gpu = "-gpu" if c.gpu(job_mode) else "" - return "tensorflow/tensorflow:{}{}-py3".format(tensorflow_version, gpu) - - -def base_image_suffix(job_mode: c.JobMode) -> str: - return "gpu" if c.gpu(job_mode) else "cpu" - - -def base_image_id(job_mode: c.JobMode) -> str: - """Returns the default base image for all caliban Dockerfiles.""" - base_suffix = base_image_suffix(job_mode) - return "{}:{}".format(DEV_CONTAINER_ROOT, base_suffix) - - -def extras_string(extras: List[str]) -> str: - """Returns the argument passed to `pip install` to install a project from its - setup.py and target a specific set of extras_require dependencies. - - Args: - extras: (potentially empty) list of extra_requires deps. - """ - ret = "." - if len(extras) > 0: - ret += "[{}]".format(','.join(extras)) - return ret - - -def base_extras(job_mode: c.JobMode, path: str, - extras: Optional[List[str]]) -> Optional[List[str]]: - """Returns None if the supplied path doesn't exist (it's assumed it points to a - setup.py file). - - If the path DOES exist, generates a list of extras to install. gpu or cpu are - always added to the beginning of the list, depending on the mode. - - """ - ret = None - - if os.path.exists(path): - base = extras or [] - extra = 'gpu' if c.gpu(job_mode) else 'cpu' - ret = base if extra in base else [extra] + base - - return ret - - -def _dependency_entries(workdir: str, - user_id: int, - user_group: int, - requirements_path: Optional[str] = None, - conda_env_path: Optional[str] = None, - setup_extras: Optional[List[str]] = None) -> str: - """Returns the Dockerfile entries required to install dependencies from either: - - - a requirements.txt file, path supplied by requirements_path - - a conda environment.yml file, path supplied by conda_env_path. - - a setup.py file, if some sequence of dependencies is supplied. - - An empty list for setup_extras means, run `pip install -c .` with no extras. - None for this argument means do nothing. If a list of strings is supplied, - they'll be treated as extras dependency sets. - """ - ret = "" - - if setup_extras is not None: - ret += f""" -COPY --chown={user_id}:{user_group} setup.py {workdir} -RUN /bin/bash -c "pip install --no-cache-dir {extras_string(setup_extras)}" -""" - - if conda_env_path is not None: - ret += f""" -COPY --chown={user_id}:{user_group} {conda_env_path} {workdir} -RUN /bin/bash -c "{CONDA_BIN} env update \ - --quiet --name caliban \ - --file {conda_env_path} && \ - {CONDA_BIN} clean -y -q --all" -""" - - if requirements_path is not None: - ret += f""" -COPY --chown={user_id}:{user_group} {requirements_path} {workdir} -RUN /bin/bash -c "pip install --no-cache-dir -r {requirements_path}" -""" - - return ret - - -def _package_entries(workdir: str, user_id: int, user_group: int, - package: u.Package) -> str: - """Returns the Dockerfile entries required to: - - - copy a directory of code into a docker container - - inject an entrypoint that executes a python module inside that directory. - - Python code runs as modules vs scripts so that we can enforce import hygiene - between files inside a project. - - """ - owner = "{}:{}".format(user_id, user_group) - - arg = package.main_module or package.script_path - - # This needs to use json so that quotes print as double quotes, not single - # quotes. - entrypoint_s = json.dumps(package.executable + [arg]) - - return """ -# Copy project code into the docker container. -COPY --chown={owner} {package_path} {workdir}/{package_path} - -# Declare an entrypoint that actually runs the container. -ENTRYPOINT {entrypoint_s} - """.format_map({ - "owner": owner, - "package_path": package.package_path, - "workdir": workdir, - "entrypoint_s": entrypoint_s - }) - - -def _service_account_entry(user_id: int, user_group: int, credentials_path: str, - docker_credentials_dir: str, - write_adc_placeholder: bool): - """Generates the Dockerfile entries required to transfer a set of Cloud service - account credentials into the Docker container. - - NOTE the write_adc_placeholder variable is here because the "ctpu" script - that we use to interact with TPUs has a bug in it, as of 1/21/2020, where the - script will fail if the application_default_credentials.json file isn't - present, EVEN THOUGH it properly uses the service account credentials - registered with gcloud instead of ADC creds. - - If a service account is present, we write a placeholder string to get past - this problem. This shouldn't matter for anyone else since adc isn't used if a - service account is present. - - """ - container_creds = "{}/credentials.json".format(docker_credentials_dir) - ret = """ -COPY --chown={user_id}:{user_group} {credentials_path} {container_creds} - -# Use the credentials file to activate gcloud, gsutil inside the container. -RUN gcloud auth activate-service-account --key-file={container_creds} && \ - git config --global credential.'https://source.developers.google.com'.helper gcloud.sh - -ENV GOOGLE_APPLICATION_CREDENTIALS={container_creds} -""".format_map({ - "user_id": user_id, - "user_group": user_group, - "credentials_path": credentials_path, - "container_creds": container_creds - }) - - if write_adc_placeholder: - ret += """ -RUN echo "placeholder" >> {} -""".format(adc_location(container_home())) - - return ret - - -def _adc_entry(user_id: int, user_group: int, adc_path: str): - """Returns the Dockerfile line required to transfer the - application_default_credentials.json file into the container's home - directory. - - """ - return """ -COPY --chown={user_id}:{user_group} {adc_path} {adc_loc} - """.format_map({ - "user_id": user_id, - "user_group": user_group, - "adc_path": adc_path, - "adc_loc": adc_location(container_home()) - }) - - -def _credentials_entries(user_id: int, - user_group: int, - adc_path: Optional[str], - credentials_path: Optional[str], - docker_credentials_dir: Optional[str] = None) -> str: - """Returns the Dockerfile entries necessary to copy a user's Cloud credentials - into the Docker container. - - - adc_path is the relative path inside the current directory to an - application_default_credentials.json file containing... well, you get it. - - credentials_path is the relative path inside the current directory to a - JSON credentials file. - - docker_credentials_dir is the relative path inside the docker container - where the JSON file will be copied on build. - - """ - if docker_credentials_dir is None: - docker_credentials_dir = CREDS_DIR - - ret = "" - if credentials_path is not None: - ret += _service_account_entry(user_id, - user_group, - credentials_path, - docker_credentials_dir, - write_adc_placeholder=adc_path is None) - - if adc_path is not None: - ret += _adc_entry(user_id, user_group, adc_path) - - return ret - - -def _notebook_entries(lab: bool = False, version: Optional[str] = None) -> str: - """Returns the Dockerfile entries necessary to install Jupyter{lab}. - - Optionally takes a version string. - - """ - version_suffix = "" - - if version is not None: - version_suffix = "=={}".format(version) - - library = "jupyterlab" if lab else "jupyter" - - return """ -RUN pip install {}{} -""".format(library, version_suffix) - - -def _custom_packages( - user_id: int, - user_group: int, - packages: Optional[List[str]] = None, - shell: Optional[Shell] = None, -) -> str: - """Returns the Dockerfile entries necessary to install custom dependencies for - the supplied shell and sequence of aptitude packages. - - """ - if packages is None: - packages = [] - - if shell is None: - shell = Shell.bash - - ret = "" - - to_install = sorted(packages + SHELL_DICT[shell].packages) - - if len(to_install) != 0: - commands = apt_command([apt_install(*to_install)]) - ret = """ -USER root - -RUN {commands} - -USER {user_id}:{user_group} -""".format_map({ - "commands": " && ".join(commands), - "user_id": user_id, - "user_group": user_group - }) - - return ret - - -def _copy_dir_entry(workdir: str, user_id: int, user_group: int, - dirname: str) -> str: - """Returns the Dockerfile entry necessary to copy a single extra subdirectory - from the current directory into a docker container during build. - - """ - owner = "{}:{}".format(user_id, user_group) - return """# Copy {dirname} into the Docker container. -COPY --chown={owner} {dirname} {workdir}/{dirname} -""".format_map({ - "owner": owner, - "workdir": workdir, - "dirname": dirname - }) - - -def _extra_dir_entries(workdir: str, user_id: int, user_group: int, - extra_dirs: List[str]) -> str: - """Returns the Dockerfile entries necessary to copy all directories in the - extra_dirs list into a docker container during build. - - """ - ret = "" - for d in extra_dirs: - ret += "\n{}".format(_copy_dir_entry(workdir, user_id, user_group, d)) - return ret - - -def _dockerfile_template( - job_mode: c.JobMode, - workdir: Optional[str] = None, - base_image_fn: Optional[Callable[[c.JobMode], str]] = None, - package: Optional[Union[List, u.Package]] = None, - requirements_path: Optional[str] = None, - conda_env_path: Optional[str] = None, - setup_extras: Optional[List[str]] = None, - adc_path: Optional[str] = None, - credentials_path: Optional[str] = None, - jupyter_version: Optional[str] = None, - inject_notebook: NotebookInstall = NotebookInstall.none, - shell: Optional[Shell] = None, - extra_dirs: Optional[List[str]] = None, - caliban_config: Optional[Dict[str, Any]] = None) -> str: - """Returns a Dockerfile that builds on a local CPU or GPU base image (depending - on the value of job_mode) to create a container that: - - - installs any dependency specified in a requirements.txt file living at - requirements_path, a conda environment at conda_env_path, or any - dependencies in a setup.py file, including extra dependencies, if - setup_extras isn't None - - injects gcloud credentials into the container, so Cloud interaction works - just like it does locally - - potentially installs a custom shell, or jupyterlab for notebook support - - copies all source needed by the main module specified by package, and - potentially injects an entrypoint that, on run, will run that main module - - Most functions that call _dockerfile_template pass along any kwargs that they - receive. It should be enough to add kwargs here, then rely on that mechanism - to pass them along, vs adding kwargs all the way down the call chain. - - Supply a custom base_image_fn (function from job_mode -> image ID) to inject - more complex Docker commands into the Caliban environments by, for example, - building your own image on top of the TF base images, then using that. - - """ - uid = os.getuid() - gid = os.getgid() - username = u.current_user() - - if isinstance(package, list): - package = u.Package(*package) - - if workdir is None: - workdir = DEFAULT_WORKDIR - - if base_image_fn is None: - base_image_fn = base_image_id - - base_image = base_image_fn(job_mode) - - dockerfile = """ -FROM {base_image} - -# Create the same group we're using on the host machine. -RUN [ $(getent group {gid}) ] || groupadd --gid {gid} {gid} - -# Create the user by name. --no-log-init guards against a crash with large user -# IDs. -RUN useradd --no-log-init --no-create-home -u {uid} -g {gid} --shell /bin/bash {username} - -# The directory is created by root. This sets permissions so that any user can -# access the folder. -RUN mkdir -m 777 {workdir} {creds_dir} {c_home} - -ENV HOME={c_home} - -WORKDIR {workdir} - -USER {uid}:{gid} -""".format_map({ - "base_image": base_image, - "username": username, - "uid": uid, - "gid": gid, - "workdir": workdir, - "c_home": container_home(), - "creds_dir": CREDS_DIR - }) - dockerfile += _credentials_entries(uid, - gid, - adc_path=adc_path, - credentials_path=credentials_path) - - dockerfile += _dependency_entries(workdir, - uid, - gid, - requirements_path=requirements_path, - conda_env_path=conda_env_path, - setup_extras=setup_extras) - - if inject_notebook.value != 'none': - install_lab = inject_notebook == NotebookInstall.lab - dockerfile += _notebook_entries(lab=install_lab, version=jupyter_version) - - if extra_dirs is not None: - dockerfile += _extra_dir_entries(workdir, uid, gid, extra_dirs) - - dockerfile += _custom_packages(uid, - gid, - packages=c.apt_packages( - caliban_config, job_mode), - shell=shell) - - if package is not None: - # The actual entrypoint and final copied code. - dockerfile += _package_entries(workdir, uid, gid, package) - - return dockerfile - - -def docker_image_id(output: str) -> ImageId: - """Accepts a string containing the output of a successful `docker build` - command and parses the Docker image ID from the stream. - - NOTE this is probably quite brittle! I can imagine this breaking quite easily - on a Docker upgrade. - - """ - return ImageId(output.splitlines()[-1].split()[-1]) - - -def build_image(job_mode: c.JobMode, - build_path: str, - credentials_path: Optional[str] = None, - adc_path: Optional[str] = None, - no_cache: bool = False, - **kwargs) -> str: - """Builds a Docker image by generating a Dockerfile and passing it to `docker - build` via stdin. All output from the `docker build` process prints to - stdout. - - Returns the image ID of the new docker container; if the command fails, - throws on error with information about the command and any issues that caused - the problem. - - """ - with u.TempCopy(credentials_path, - tmp_name=".caliban_default_creds.json") as creds: - with u.TempCopy(adc_path, tmp_name=".caliban_adc_creds.json") as adc: - cache_args = ["--no-cache"] if no_cache else [] - cmd = ["docker", "build"] + cache_args + ["--rm", "-f-", build_path] - - dockerfile = _dockerfile_template(job_mode, - credentials_path=creds, - adc_path=adc, - **kwargs) - - joined_cmd = " ".join(cmd) - logging.info("Running command: {}".format(joined_cmd)) - - try: - output, ret_code = ufs.capture_stdout(cmd, input_str=dockerfile) - if ret_code == 0: - return docker_image_id(output) - else: - error_msg = "Docker failed with error code {}.".format(ret_code) - raise DockerError(error_msg, cmd, ret_code) - - except subprocess.CalledProcessError as e: - logging.error(e.output) - logging.error(e.stderr) - - -def _image_tag_for_project(project_id: str, image_id: str) -> str: - """Generate the GCR Docker image tag for the supplied pair of project_id and - image_id. - - This function properly handles "domain scoped projects", where the project ID - contains a domain name and project ID separated by : - https://cloud.google.com/container-registry/docs/overview#domain-scoped_projects. - - """ - project_s = project_id.replace(":", "/") - return "gcr.io/{}/{}:latest".format(project_s, image_id) - - -def push_uuid_tag(project_id: str, image_id: str) -> str: - """Takes a base image and tags it for upload, then pushes it to a remote Google - Container Registry. - - Returns the tag on a successful push. - - TODO should this just check first before attempting to push if the image - exists? Immutable names means that if the tag is up there, we're done. - Potentially use docker-py for this. - - """ - image_tag = _image_tag_for_project(project_id, image_id) - subprocess.run(["docker", "tag", image_id, image_tag], check=True) - subprocess.run(["docker", "push", image_tag], check=True) - return image_tag - - -def _run_cmd(job_mode: c.JobMode, - run_args: Optional[List[str]] = None) -> List[str]: - """Returns the sequence of commands for the subprocess run functions required - to execute `docker run`. in CPU or GPU mode, depending on the value of - job_mode. - - Keyword args: - - run_args: list of args to pass to docker run. - - """ - if run_args is None: - run_args = [] - - runtime = ["--runtime", "nvidia"] if c.gpu(job_mode) else [] - return ["docker", "run"] + runtime + ["--ipc", "host"] + run_args - - -def _home_mount_cmds(enable_home_mount: bool) -> List[str]: - """Returns the argument needed by Docker to mount a user's local home directory - into the home directory location inside their container. - - If enable_home_mount is false returns an empty list. - - """ - ret = [] - if enable_home_mount: - ret = ["-v", "{}:{}".format(Path.home(), container_home())] - return ret - - -def _interactive_opts(workdir: str) -> List[str]: - """Returns the basic arguments we want to run a docker process locally. - - """ - return [ - "-w", workdir, \ - "-u", "{}:{}".format(os.getuid(), os.getgid()), \ - "-v", "{}:{}".format(os.getcwd(), workdir) \ - ] - - -def log_job_spec_instance(job_spec: JobSpec, i: int) -> JobSpec: - """Prints logging as a side effect for the supplied sequence of job specs - generated from an experiment definition; returns the input job spec. - - """ - args = c.experiment_to_args(job_spec.experiment.kwargs, - job_spec.experiment.args) - logging.info("") - logging.info("Job {} - Experiment args: {}".format(i, t.yellow(str(args)))) - return job_spec - - -def logged_job_specs(job_specs: Iterable[JobSpec]) -> Iterable[JobSpec]: - """Generates an iterable of job specs that should be passed to `docker run` to - execute the experiments defined by the supplied iterable. - - """ - for i, s in enumerate(job_specs, 1): - yield log_job_spec_instance(s, i) - - -def execute_dry_run(job_specs: Iterable[JobSpec]) -> None: - """Expands the supplied sequence of experiments into sequences of args and logs - the jobs that WOULD have been executed, had the dry run flag not been - applied. - - """ - list(logged_job_specs(job_specs)) - - logging.info('') - logging.info( - t.yellow("To build your image and execute these jobs, \ -run your command again without {}.".format(c.DRY_RUN_FLAG))) - logging.info('') - return None diff --git a/caliban/cli.py b/caliban/cli.py index ae270aa..427c743 100644 --- a/caliban/cli.py +++ b/caliban/cli.py @@ -24,14 +24,14 @@ from absl.flags import argparse_flags from blessings import Terminal -import caliban.cloud.types as ct +import caliban.platform.cloud.types as ct import caliban.config as conf import caliban.config.experiment as ce import caliban.docker as docker -import caliban.gke as gke -import caliban.gke.constants as gke_k -import caliban.gke.types as gke_t -import caliban.gke.util as gke_u +import caliban.platform.gke as gke +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.util as u from caliban import __version__ diff --git a/caliban/config/__init__.py b/caliban/config/__init__.py index ac60130..52f0b04 100644 --- a/caliban/config/__init__.py +++ b/caliban/config/__init__.py @@ -137,7 +137,7 @@ def extract_project_id(m: Dict[str, Any]) -> str: def extract_region(m: Dict[str, Any]) -> ct.Region: """Returns the region specified in the args; defaults to an environment variable. If that's not supplied defaults to the default cloud provider from - caliban.cloud. + caliban.platform.cloud. """ region = m.get("region") or os.environ.get("REGION") diff --git a/caliban/docker/build.py b/caliban/docker/build.py index 3185eae..0740853 100644 --- a/caliban/docker/build.py +++ b/caliban/docker/build.py @@ -593,9 +593,9 @@ def build_image(job_mode: c.JobMode, the problem. """ - with u.TempCopy(credentials_path, - tmp_name=".caliban_default_creds.json") as creds: - with u.TempCopy(adc_path, tmp_name=".caliban_adc_creds.json") as adc: + with ufs.TempCopy(credentials_path, + tmp_name=".caliban_default_creds.json") as creds: + with ufs.TempCopy(adc_path, tmp_name=".caliban_adc_creds.json") as adc: cache_args = ["--no-cache"] if no_cache else [] cmd = ["docker", "build"] + cache_args + ["--rm", "-f-", build_path] diff --git a/caliban/docker/push.py b/caliban/docker/push.py index e4acd5a..3ec7928 100644 --- a/caliban/docker/push.py +++ b/caliban/docker/push.py @@ -18,607 +18,7 @@ """ -from __future__ import absolute_import, division, print_function - -import json -import os import subprocess -from enum import Enum -from pathlib import Path -from typing import (Any, Callable, Dict, Iterable, List, NamedTuple, NewType, - Optional, Union) - -from absl import logging -from blessings import Terminal - -import caliban.config as c -import caliban.util as u -import caliban.util.fs as ufs -from caliban.history.types import JobSpec - -t = Terminal() - -DEV_CONTAINER_ROOT = "gcr.io/blueshift-playground/blueshift" -TF_VERSIONS = {"2.2.0", "1.12.3", "1.14.0", "1.15.0"} -DEFAULT_WORKDIR = "/usr/app" -CREDS_DIR = "/.creds" -CONDA_BIN = "/opt/conda/bin/conda" - -ImageId = NewType('ImageId', str) -ArgSeq = NewType('ArgSeq', List[str]) - - -class DockerError(Exception): - """Exception that passes info on a failed Docker command.""" - - def __init__(self, message, cmd, ret_code): - super().__init__(message) - self.message = message - self.cmd = cmd - self.ret_code = ret_code - - @property - def command(self): - return " ".join(self.cmd) - - -class NotebookInstall(Enum): - """Flag to decide what to do .""" - none = 'none' - lab = 'lab' - jupyter = 'jupyter' - - def __str__(self) -> str: - return self.value - - -class Shell(Enum): - """Add new shells here and below, in SHELL_DICT.""" - bash = 'bash' - zsh = 'zsh' - - def __str__(self) -> str: - return self.value - - -# Tuple to track the information required to install and execute some custom -# shell into a container. -ShellData = NamedTuple("ShellData", [("executable", str), - ("packages", List[str])]) - - -def apt_install(*packages: str) -> str: - """Returns a command that will install the supplied list of packages without - requiring confirmation or any user interaction. - """ - package_str = ' '.join(packages) - no_prompt = "DEBIAN_FRONTEND=noninteractive" - return f"{no_prompt} apt-get install --yes --no-install-recommends {package_str}" - - -def apt_command(commands: List[str]) -> List[str]: - """Pre-and-ap-pends the supplied commands with the appropriate in-container and - cleanup command for aptitude. - - """ - update = ["apt-get update"] - cleanup = ["apt-get clean", "rm -rf /var/lib/apt/lists/*"] - return update + commands + cleanup - - -# Dict linking a particular supported shell to the data required to run and -# install the shell inside a container. -# -# : Dict[Shell, ShellData] -SHELL_DICT = { - Shell.bash: ShellData("/bin/bash", []), - Shell.zsh: ShellData("/bin/zsh", ["zsh"]) -} - - -def default_shell() -> Shell: - """Returns the shell to load into the container. Defaults to Shell.bash, but if - the user's SHELL variable refers to a supported sub-shell, returns that - instead. - - """ - ret = Shell.bash - - if "zsh" in os.environ.get("SHELL"): - ret = Shell.zsh - - return ret - - -def adc_location(home_dir: Optional[str] = None) -> str: - """Returns the location for application default credentials, INSIDE the - container (so, hardcoded unix separators), given the supplied home directory. - - """ - if home_dir is None: - home_dir = Path.home() - - return "{}/.config/gcloud/application_default_credentials.json".format( - home_dir) - - -def container_home(): - """Returns the location of the home directory inside the generated - container. - - """ - return "/home/{}".format(u.current_user()) - - -def tf_base_image(job_mode: c.JobMode, tensorflow_version: str) -> str: - """Returns the base image to use, depending on whether or not we're using a - GPU. This is JUST for building our base images for Blueshift; not for - actually using in a job. - - List of available tags: https://hub.docker.com/r/tensorflow/tensorflow/tags - - """ - if tensorflow_version not in TF_VERSIONS: - raise Exception("""{} is not a valid tensorflow version. - Try one of: {}""".format(tensorflow_version, TF_VERSIONS)) - - gpu = "-gpu" if c.gpu(job_mode) else "" - return "tensorflow/tensorflow:{}{}-py3".format(tensorflow_version, gpu) - - -def base_image_suffix(job_mode: c.JobMode) -> str: - return "gpu" if c.gpu(job_mode) else "cpu" - - -def base_image_id(job_mode: c.JobMode) -> str: - """Returns the default base image for all caliban Dockerfiles.""" - base_suffix = base_image_suffix(job_mode) - return "{}:{}".format(DEV_CONTAINER_ROOT, base_suffix) - - -def extras_string(extras: List[str]) -> str: - """Returns the argument passed to `pip install` to install a project from its - setup.py and target a specific set of extras_require dependencies. - - Args: - extras: (potentially empty) list of extra_requires deps. - """ - ret = "." - if len(extras) > 0: - ret += "[{}]".format(','.join(extras)) - return ret - - -def base_extras(job_mode: c.JobMode, path: str, - extras: Optional[List[str]]) -> Optional[List[str]]: - """Returns None if the supplied path doesn't exist (it's assumed it points to a - setup.py file). - - If the path DOES exist, generates a list of extras to install. gpu or cpu are - always added to the beginning of the list, depending on the mode. - - """ - ret = None - - if os.path.exists(path): - base = extras or [] - extra = 'gpu' if c.gpu(job_mode) else 'cpu' - ret = base if extra in base else [extra] + base - - return ret - - -def _dependency_entries(workdir: str, - user_id: int, - user_group: int, - requirements_path: Optional[str] = None, - conda_env_path: Optional[str] = None, - setup_extras: Optional[List[str]] = None) -> str: - """Returns the Dockerfile entries required to install dependencies from either: - - - a requirements.txt file, path supplied by requirements_path - - a conda environment.yml file, path supplied by conda_env_path. - - a setup.py file, if some sequence of dependencies is supplied. - - An empty list for setup_extras means, run `pip install -c .` with no extras. - None for this argument means do nothing. If a list of strings is supplied, - they'll be treated as extras dependency sets. - """ - ret = "" - - if setup_extras is not None: - ret += f""" -COPY --chown={user_id}:{user_group} setup.py {workdir} -RUN /bin/bash -c "pip install --no-cache-dir {extras_string(setup_extras)}" -""" - - if conda_env_path is not None: - ret += f""" -COPY --chown={user_id}:{user_group} {conda_env_path} {workdir} -RUN /bin/bash -c "{CONDA_BIN} env update \ - --quiet --name caliban \ - --file {conda_env_path} && \ - {CONDA_BIN} clean -y -q --all" -""" - - if requirements_path is not None: - ret += f""" -COPY --chown={user_id}:{user_group} {requirements_path} {workdir} -RUN /bin/bash -c "pip install --no-cache-dir -r {requirements_path}" -""" - - return ret - - -def _package_entries(workdir: str, user_id: int, user_group: int, - package: u.Package) -> str: - """Returns the Dockerfile entries required to: - - - copy a directory of code into a docker container - - inject an entrypoint that executes a python module inside that directory. - - Python code runs as modules vs scripts so that we can enforce import hygiene - between files inside a project. - - """ - owner = "{}:{}".format(user_id, user_group) - - arg = package.main_module or package.script_path - - # This needs to use json so that quotes print as double quotes, not single - # quotes. - entrypoint_s = json.dumps(package.executable + [arg]) - - return """ -# Copy project code into the docker container. -COPY --chown={owner} {package_path} {workdir}/{package_path} - -# Declare an entrypoint that actually runs the container. -ENTRYPOINT {entrypoint_s} - """.format_map({ - "owner": owner, - "package_path": package.package_path, - "workdir": workdir, - "entrypoint_s": entrypoint_s - }) - - -def _service_account_entry(user_id: int, user_group: int, credentials_path: str, - docker_credentials_dir: str, - write_adc_placeholder: bool): - """Generates the Dockerfile entries required to transfer a set of Cloud service - account credentials into the Docker container. - - NOTE the write_adc_placeholder variable is here because the "ctpu" script - that we use to interact with TPUs has a bug in it, as of 1/21/2020, where the - script will fail if the application_default_credentials.json file isn't - present, EVEN THOUGH it properly uses the service account credentials - registered with gcloud instead of ADC creds. - - If a service account is present, we write a placeholder string to get past - this problem. This shouldn't matter for anyone else since adc isn't used if a - service account is present. - - """ - container_creds = "{}/credentials.json".format(docker_credentials_dir) - ret = """ -COPY --chown={user_id}:{user_group} {credentials_path} {container_creds} - -# Use the credentials file to activate gcloud, gsutil inside the container. -RUN gcloud auth activate-service-account --key-file={container_creds} && \ - git config --global credential.'https://source.developers.google.com'.helper gcloud.sh - -ENV GOOGLE_APPLICATION_CREDENTIALS={container_creds} -""".format_map({ - "user_id": user_id, - "user_group": user_group, - "credentials_path": credentials_path, - "container_creds": container_creds - }) - - if write_adc_placeholder: - ret += """ -RUN echo "placeholder" >> {} -""".format(adc_location(container_home())) - - return ret - - -def _adc_entry(user_id: int, user_group: int, adc_path: str): - """Returns the Dockerfile line required to transfer the - application_default_credentials.json file into the container's home - directory. - - """ - return """ -COPY --chown={user_id}:{user_group} {adc_path} {adc_loc} - """.format_map({ - "user_id": user_id, - "user_group": user_group, - "adc_path": adc_path, - "adc_loc": adc_location(container_home()) - }) - - -def _credentials_entries(user_id: int, - user_group: int, - adc_path: Optional[str], - credentials_path: Optional[str], - docker_credentials_dir: Optional[str] = None) -> str: - """Returns the Dockerfile entries necessary to copy a user's Cloud credentials - into the Docker container. - - - adc_path is the relative path inside the current directory to an - application_default_credentials.json file containing... well, you get it. - - credentials_path is the relative path inside the current directory to a - JSON credentials file. - - docker_credentials_dir is the relative path inside the docker container - where the JSON file will be copied on build. - - """ - if docker_credentials_dir is None: - docker_credentials_dir = CREDS_DIR - - ret = "" - if credentials_path is not None: - ret += _service_account_entry(user_id, - user_group, - credentials_path, - docker_credentials_dir, - write_adc_placeholder=adc_path is None) - - if adc_path is not None: - ret += _adc_entry(user_id, user_group, adc_path) - - return ret - - -def _notebook_entries(lab: bool = False, version: Optional[str] = None) -> str: - """Returns the Dockerfile entries necessary to install Jupyter{lab}. - - Optionally takes a version string. - - """ - version_suffix = "" - - if version is not None: - version_suffix = "=={}".format(version) - - library = "jupyterlab" if lab else "jupyter" - - return """ -RUN pip install {}{} -""".format(library, version_suffix) - - -def _custom_packages( - user_id: int, - user_group: int, - packages: Optional[List[str]] = None, - shell: Optional[Shell] = None, -) -> str: - """Returns the Dockerfile entries necessary to install custom dependencies for - the supplied shell and sequence of aptitude packages. - - """ - if packages is None: - packages = [] - - if shell is None: - shell = Shell.bash - - ret = "" - - to_install = sorted(packages + SHELL_DICT[shell].packages) - - if len(to_install) != 0: - commands = apt_command([apt_install(*to_install)]) - ret = """ -USER root - -RUN {commands} - -USER {user_id}:{user_group} -""".format_map({ - "commands": " && ".join(commands), - "user_id": user_id, - "user_group": user_group - }) - - return ret - - -def _copy_dir_entry(workdir: str, user_id: int, user_group: int, - dirname: str) -> str: - """Returns the Dockerfile entry necessary to copy a single extra subdirectory - from the current directory into a docker container during build. - - """ - owner = "{}:{}".format(user_id, user_group) - return """# Copy {dirname} into the Docker container. -COPY --chown={owner} {dirname} {workdir}/{dirname} -""".format_map({ - "owner": owner, - "workdir": workdir, - "dirname": dirname - }) - - -def _extra_dir_entries(workdir: str, user_id: int, user_group: int, - extra_dirs: List[str]) -> str: - """Returns the Dockerfile entries necessary to copy all directories in the - extra_dirs list into a docker container during build. - - """ - ret = "" - for d in extra_dirs: - ret += "\n{}".format(_copy_dir_entry(workdir, user_id, user_group, d)) - return ret - - -def _dockerfile_template( - job_mode: c.JobMode, - workdir: Optional[str] = None, - base_image_fn: Optional[Callable[[c.JobMode], str]] = None, - package: Optional[Union[List, u.Package]] = None, - requirements_path: Optional[str] = None, - conda_env_path: Optional[str] = None, - setup_extras: Optional[List[str]] = None, - adc_path: Optional[str] = None, - credentials_path: Optional[str] = None, - jupyter_version: Optional[str] = None, - inject_notebook: NotebookInstall = NotebookInstall.none, - shell: Optional[Shell] = None, - extra_dirs: Optional[List[str]] = None, - caliban_config: Optional[Dict[str, Any]] = None) -> str: - """Returns a Dockerfile that builds on a local CPU or GPU base image (depending - on the value of job_mode) to create a container that: - - - installs any dependency specified in a requirements.txt file living at - requirements_path, a conda environment at conda_env_path, or any - dependencies in a setup.py file, including extra dependencies, if - setup_extras isn't None - - injects gcloud credentials into the container, so Cloud interaction works - just like it does locally - - potentially installs a custom shell, or jupyterlab for notebook support - - copies all source needed by the main module specified by package, and - potentially injects an entrypoint that, on run, will run that main module - - Most functions that call _dockerfile_template pass along any kwargs that they - receive. It should be enough to add kwargs here, then rely on that mechanism - to pass them along, vs adding kwargs all the way down the call chain. - - Supply a custom base_image_fn (function from job_mode -> image ID) to inject - more complex Docker commands into the Caliban environments by, for example, - building your own image on top of the TF base images, then using that. - - """ - uid = os.getuid() - gid = os.getgid() - username = u.current_user() - - if isinstance(package, list): - package = u.Package(*package) - - if workdir is None: - workdir = DEFAULT_WORKDIR - - if base_image_fn is None: - base_image_fn = base_image_id - - base_image = base_image_fn(job_mode) - - dockerfile = """ -FROM {base_image} - -# Create the same group we're using on the host machine. -RUN [ $(getent group {gid}) ] || groupadd --gid {gid} {gid} - -# Create the user by name. --no-log-init guards against a crash with large user -# IDs. -RUN useradd --no-log-init --no-create-home -u {uid} -g {gid} --shell /bin/bash {username} - -# The directory is created by root. This sets permissions so that any user can -# access the folder. -RUN mkdir -m 777 {workdir} {creds_dir} {c_home} - -ENV HOME={c_home} - -WORKDIR {workdir} - -USER {uid}:{gid} -""".format_map({ - "base_image": base_image, - "username": username, - "uid": uid, - "gid": gid, - "workdir": workdir, - "c_home": container_home(), - "creds_dir": CREDS_DIR - }) - dockerfile += _credentials_entries(uid, - gid, - adc_path=adc_path, - credentials_path=credentials_path) - - dockerfile += _dependency_entries(workdir, - uid, - gid, - requirements_path=requirements_path, - conda_env_path=conda_env_path, - setup_extras=setup_extras) - - if inject_notebook.value != 'none': - install_lab = inject_notebook == NotebookInstall.lab - dockerfile += _notebook_entries(lab=install_lab, version=jupyter_version) - - if extra_dirs is not None: - dockerfile += _extra_dir_entries(workdir, uid, gid, extra_dirs) - - dockerfile += _custom_packages(uid, - gid, - packages=c.apt_packages( - caliban_config, job_mode), - shell=shell) - - if package is not None: - # The actual entrypoint and final copied code. - dockerfile += _package_entries(workdir, uid, gid, package) - - return dockerfile - - -def docker_image_id(output: str) -> ImageId: - """Accepts a string containing the output of a successful `docker build` - command and parses the Docker image ID from the stream. - - NOTE this is probably quite brittle! I can imagine this breaking quite easily - on a Docker upgrade. - - """ - return ImageId(output.splitlines()[-1].split()[-1]) - - -def build_image(job_mode: c.JobMode, - build_path: str, - credentials_path: Optional[str] = None, - adc_path: Optional[str] = None, - no_cache: bool = False, - **kwargs) -> str: - """Builds a Docker image by generating a Dockerfile and passing it to `docker - build` via stdin. All output from the `docker build` process prints to - stdout. - - Returns the image ID of the new docker container; if the command fails, - throws on error with information about the command and any issues that caused - the problem. - - """ - with ufs.TempCopy(credentials_path, - tmp_name=".caliban_default_creds.json") as creds: - with ufs.TempCopy(adc_path, tmp_name=".caliban_adc_creds.json") as adc: - cache_args = ["--no-cache"] if no_cache else [] - cmd = ["docker", "build"] + cache_args + ["--rm", "-f-", build_path] - - dockerfile = _dockerfile_template(job_mode, - credentials_path=creds, - adc_path=adc, - **kwargs) - - joined_cmd = " ".join(cmd) - logging.info("Running command: {}".format(joined_cmd)) - - try: - output, ret_code = ufs.capture_stdout(cmd, input_str=dockerfile) - if ret_code == 0: - return docker_image_id(output) - else: - error_msg = "Docker failed with error code {}.".format(ret_code) - raise DockerError(error_msg, cmd, ret_code) - - except subprocess.CalledProcessError as e: - logging.error(e.output) - logging.error(e.stderr) def _image_tag_for_project(project_id: str, image_id: str) -> str: @@ -649,81 +49,3 @@ def push_uuid_tag(project_id: str, image_id: str) -> str: subprocess.run(["docker", "tag", image_id, image_tag], check=True) subprocess.run(["docker", "push", image_tag], check=True) return image_tag - - -def _run_cmd(job_mode: c.JobMode, - run_args: Optional[List[str]] = None) -> List[str]: - """Returns the sequence of commands for the subprocess run functions required - to execute `docker run`. in CPU or GPU mode, depending on the value of - job_mode. - - Keyword args: - - run_args: list of args to pass to docker run. - - """ - if run_args is None: - run_args = [] - - runtime = ["--runtime", "nvidia"] if c.gpu(job_mode) else [] - return ["docker", "run"] + runtime + ["--ipc", "host"] + run_args - - -def _home_mount_cmds(enable_home_mount: bool) -> List[str]: - """Returns the argument needed by Docker to mount a user's local home directory - into the home directory location inside their container. - - If enable_home_mount is false returns an empty list. - - """ - ret = [] - if enable_home_mount: - ret = ["-v", "{}:{}".format(Path.home(), container_home())] - return ret - - -def _interactive_opts(workdir: str) -> List[str]: - """Returns the basic arguments we want to run a docker process locally. - - """ - return [ - "-w", workdir, \ - "-u", "{}:{}".format(os.getuid(), os.getgid()), \ - "-v", "{}:{}".format(os.getcwd(), workdir) \ - ] - - -def log_job_spec_instance(job_spec: JobSpec, i: int) -> JobSpec: - """Prints logging as a side effect for the supplied sequence of job specs - generated from an experiment definition; returns the input job spec. - - """ - args = c.experiment_to_args(job_spec.experiment.kwargs, - job_spec.experiment.args) - logging.info("") - logging.info("Job {} - Experiment args: {}".format(i, t.yellow(str(args)))) - return job_spec - - -def logged_job_specs(job_specs: Iterable[JobSpec]) -> Iterable[JobSpec]: - """Generates an iterable of job specs that should be passed to `docker run` to - execute the experiments defined by the supplied iterable. - - """ - for i, s in enumerate(job_specs, 1): - yield log_job_spec_instance(s, i) - - -def execute_dry_run(job_specs: Iterable[JobSpec]) -> None: - """Expands the supplied sequence of experiments into sequences of args and logs - the jobs that WOULD have been executed, had the dry run flag not been - applied. - - """ - list(logged_job_specs(job_specs)) - - logging.info('') - logging.info( - t.yellow("To build your image and execute these jobs, \ -run your command again without {}.".format(c.DRY_RUN_FLAG))) - logging.info('') - return None diff --git a/caliban/history/cli.py b/caliban/history/cli.py index 4eb8e14..1b08bc0 100644 --- a/caliban/history/cli.py +++ b/caliban/history/cli.py @@ -30,7 +30,7 @@ from caliban.history.submit import submit_job_specs from caliban.history.types import (ContainerSpec, ExperimentGroup, Experiment, JobSpec, Job, Platform, JobStatus, Platform) -from caliban.gke.util import user_verify, credentials +from caliban.platform.gke.util import user_verify, credentials from caliban.docker import build_image, push_uuid_tag, execute_jobs diff --git a/caliban/history/submit.py b/caliban/history/submit.py index 3b61bf2..f827ee7 100644 --- a/caliban/history/submit.py +++ b/caliban/history/submit.py @@ -20,8 +20,8 @@ from caliban.history.types import JobSpec, Job, Platform import caliban.docker as docker -import caliban.cloud.core as cloud -import caliban.gke.cli as gke_cli +import caliban.platform.cloud.core as cloud +import caliban.platform.gke.cli as gke_cli # ---------------------------------------------------------------------------- diff --git a/caliban/history/util.py b/caliban/history/util.py index 917d9de..e653aad 100644 --- a/caliban/history/util.py +++ b/caliban/history/util.py @@ -29,10 +29,10 @@ from sqlalchemy.orm import Session, sessionmaker import caliban.config as conf -from caliban.cloud.types import JobStatus as CloudStatus -from caliban.gke.cluster import Cluster -from caliban.gke.types import JobStatus as GkeStatus -from caliban.gke.util import default_credentials +from caliban.platform.cloud.types import JobStatus as CloudStatus +from caliban.platform.gke.cluster import Cluster +from caliban.platform.gke.types import JobStatus as GkeStatus +from caliban.platform.gke.util import default_credentials from caliban.history.types import (ContainerSpec, Experiment, ExperimentGroup, Job, JobSpec, JobStatus, Platform, init_db) diff --git a/caliban/main.py b/caliban/main.py index d1120ff..b701523 100644 --- a/caliban/main.py +++ b/caliban/main.py @@ -25,11 +25,11 @@ from blessings import Terminal import caliban.cli as cli -import caliban.cloud.core as cloud +import caliban.platform.cloud.core as cloud import caliban.config as c import caliban.docker as docker -import caliban.gke as gke -import caliban.gke.cli +import caliban.platform.gke as gke +import caliban.platform.gke.cli import caliban.util as u import caliban.history.cli diff --git a/caliban/platform/cloud/core.py b/caliban/platform/cloud/core.py index b6c4829..5556266 100644 --- a/caliban/platform/cloud/core.py +++ b/caliban/platform/cloud/core.py @@ -31,7 +31,7 @@ from googleapiclient import discovery from googleapiclient.errors import HttpError -import caliban.cloud.types as ct +import caliban.platform.cloud.types as ct import caliban.config as conf import caliban.docker as d import caliban.history.types as ht diff --git a/caliban/platform/gke/cli.py b/caliban/platform/gke/cli.py index a71f1d7..a1c253d 100644 --- a/caliban/platform/gke/cli.py +++ b/caliban/platform/gke/cli.py @@ -28,11 +28,11 @@ import caliban.cli as cli import caliban.config as conf -import caliban.gke.constants as k -import caliban.gke.util as util +import caliban.platform.gke.constants as k +import caliban.platform.gke.util as util import caliban.util as u -from caliban.cloud.core import generate_image_tag -from caliban.gke.cluster import Cluster +from caliban.platform.cloud.core import generate_image_tag +from caliban.platform.gke.cluster import Cluster from caliban.history.util import (create_experiments, generate_container_spec, get_mem_engine, get_sql_engine, session_scope) diff --git a/caliban/platform/gke/cluster.py b/caliban/platform/gke/cluster.py index aaa987e..512f558 100644 --- a/caliban/platform/gke/cluster.py +++ b/caliban/platform/gke/cluster.py @@ -38,12 +38,12 @@ from kubernetes.client.api_client import ApiClient import caliban.config as conf -import caliban.gke.constants as k -import caliban.gke.util as util -from caliban.cloud.types import (GPU, TPU, Accelerator, GPUSpec, MachineType, - TPUSpec) -from caliban.gke.types import NodeImage, OpStatus, ReleaseChannel -from caliban.gke.util import trap +import caliban.platform.gke.constants as k +import caliban.platform.gke.util as util +from caliban.platform.cloud.types import (GPU, TPU, Accelerator, GPUSpec, + MachineType, TPUSpec) +from caliban.platform.gke.types import NodeImage, OpStatus, ReleaseChannel +from caliban.platform.gke.util import trap from caliban.history.types import Experiment, Job, JobSpec, JobStatus, Platform urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) diff --git a/caliban/platform/gke/constants.py b/caliban/platform/gke/constants.py index dcd6b38..d78c4b6 100644 --- a/caliban/platform/gke/constants.py +++ b/caliban/platform/gke/constants.py @@ -17,9 +17,9 @@ import re -from caliban.cloud.types import GPU, GPUSpec +from caliban.platform.cloud.types import GPU, GPUSpec from caliban.config import DEFAULT_MACHINE_TYPE, JobMode -from caliban.gke.types import ReleaseChannel +from caliban.platform.gke.types import ReleaseChannel COMPUTE_SCOPE_URL = 'https://www.googleapis.com/auth/compute' COMPUTE_READONLY_SCOPE_URL = 'https://www.googleapis.com/auth/compute.readonly' diff --git a/caliban/platform/gke/util.py b/caliban/platform/gke/util.py index 22926a3..dea441d 100644 --- a/caliban/platform/gke/util.py +++ b/caliban/platform/gke/util.py @@ -41,9 +41,9 @@ from yaspin import yaspin from yaspin.spinners import Spinners -import caliban.gke.constants as k -from caliban.cloud.types import GPU, TPU, GPUSpec, TPUSpec -from caliban.gke.types import (CredentialsData, NodeImage, OpStatus) +import caliban.platform.gke.constants as k +from caliban.platform.cloud.types import GPU, TPU, GPUSpec, TPUSpec +from caliban.platform.gke.types import (CredentialsData, NodeImage, OpStatus) # ---------------------------------------------------------------------------- diff --git a/tests/caliban/docker/test_build.py b/tests/caliban/docker/test_build.py new file mode 100644 index 0000000..16a9afd --- /dev/null +++ b/tests/caliban/docker/test_build.py @@ -0,0 +1,23 @@ +#!/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. + +import caliban.docker.build as b + + +def test_shell_dict(): + """Tests that the shell dict has an entry for all possible Shell values.""" + + assert set(b.Shell) == set(b.SHELL_DICT.keys()) diff --git a/tests/caliban/docker/test_docker.py b/tests/caliban/docker/test_docker.py deleted file mode 100644 index 04600c0..0000000 --- a/tests/caliban/docker/test_docker.py +++ /dev/null @@ -1,39 +0,0 @@ -#!/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. - -import unittest - -import caliban.docker as d - - -class DockerTestSuite(unittest.TestCase): - """Tests for the docker package.""" - - def test_shell_dict(self): - """Tests that the shell dict has an entry for all possible Shell values.""" - - self.assertSetEqual(set(d.Shell), set(d.SHELL_DICT.keys())) - - def test_image_tag_for_project(self): - """Tests that we generate a valid image tag for domain-scoped and modern - project IDs. - - """ - self.assertEqual(d._image_tag_for_project("face", "imageid"), - "gcr.io/face/imageid:latest") - - self.assertEqual(d._image_tag_for_project("google.com:face", "imageid"), - "gcr.io/google.com/face/imageid:latest") diff --git a/tests/caliban/docker/test_push.py b/tests/caliban/docker/test_push.py new file mode 100644 index 0000000..be0a02d --- /dev/null +++ b/tests/caliban/docker/test_push.py @@ -0,0 +1,29 @@ +#!/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. + +import caliban.docker.push as p + + +def test_image_tag_for_project(): + """Tests that we generate a valid image tag for domain-scoped and modern + project IDs. + + """ + assert p._image_tag_for_project("face", + "imageid") == "gcr.io/face/imageid:latest" + + assert p._image_tag_for_project( + "google.com:face", "imageid") == "gcr.io/google.com/face/imageid:latest" diff --git a/tests/caliban/platform/cloud/test_types.py b/tests/caliban/platform/cloud/test_types.py index 690cc6e..f8db40a 100644 --- a/tests/caliban/platform/cloud/test_types.py +++ b/tests/caliban/platform/cloud/test_types.py @@ -20,11 +20,11 @@ import hypothesis.strategies as st from hypothesis import given -import caliban.cloud.types as ct +import caliban.platform.cloud.types as ct class TypesTestSuite(unittest.TestCase): - """Tests for caliban.cloud.types.""" + """Tests for caliban.platform.cloud.types.""" @given(st.integers(min_value=0, max_value=40), st.sampled_from(list(ct.GPU) + list(ct.TPU))) diff --git a/tests/caliban/platform/gke/test_types.py b/tests/caliban/platform/gke/test_types.py index f727502..04b876d 100644 --- a/tests/caliban/platform/gke/test_types.py +++ b/tests/caliban/platform/gke/test_types.py @@ -19,12 +19,12 @@ import hypothesis.strategies as st from hypothesis import given -from caliban.gke.types import ReleaseChannel +from caliban.platform.gke.types import ReleaseChannel # ---------------------------------------------------------------------------- class TypesTestSuite(unittest.TestCase): - """tests for caliban.gke.types""" + """tests for caliban.platform.gke.types""" # -------------------------------------------------------------------------- @given(st.from_regex('\A(?!UNSPECIFIED\Z|RAPID\Z|REGULAR\Z|STABLE\Z).*\Z'), diff --git a/tests/caliban/platform/gke/test_utils.py b/tests/caliban/platform/gke/test_utils.py index f6b307d..59d3eba 100644 --- a/tests/caliban/platform/gke/test_utils.py +++ b/tests/caliban/platform/gke/test_utils.py @@ -22,11 +22,11 @@ import hypothesis.strategies as st from hypothesis import given, settings -import caliban.cloud.types as ct -import caliban.gke.constants as k -import caliban.gke.utils as utils -from caliban.gke.types import NodeImage, OpStatus -from caliban.gke.utils import trap +import caliban.platform.cloud.types as ct +import caliban.platform.gke.constants as k +import caliban.platform.gke.utils as utils +from caliban.platform.gke.types import NodeImage, OpStatus +from caliban.platform.gke.utils import trap # ---------------------------------------------------------------------------- @@ -43,7 +43,7 @@ def everything_except(excluded_types): # ---------------------------------------------------------------------------- class UtilsTestSuite(unittest.TestCase): - """tests for caliban.gke.utils""" + """tests for caliban.platform.gke.utils""" # -------------------------------------------------------------------------- @given( @@ -89,7 +89,7 @@ def test_nvidia_daemonset_url(self): return # -------------------------------------------------------------------------- - @mock.patch('caliban.gke.utils.input', create=True) + @mock.patch('caliban.platform.gke.utils.input', create=True) @given(st.lists(st.from_regex('^[^yYnN]+$'), min_size=0, max_size=8)) def test_user_verify( self, diff --git a/tests/caliban/test_builder.py b/tests/caliban/test_builder.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/caliban/test_cli.py b/tests/caliban/test_cli.py index 8074c23..ab20033 100644 --- a/tests/caliban/test_cli.py +++ b/tests/caliban/test_cli.py @@ -17,7 +17,7 @@ import unittest import caliban.cli as c -import caliban.cloud.types as ct +import caliban.platform.cloud.types as ct from caliban.config import JobMode From bd2e4141c830010bb9f428c01df82f670122a16f Mon Sep 17 00:00:00 2001 From: Sam Ritchie Date: Sat, 11 Jul 2020 13:17:59 -0600 Subject: [PATCH 06/15] getting closer --- caliban/cli.py | 9 +- caliban/config/__init__.py | 12 + caliban/expansion.py | 12 +- caliban/history/cli.py | 25 +- caliban/history/submit.py | 9 +- caliban/history/util.py | 10 +- caliban/main.py | 58 +- caliban/platform/cloud/core.py | 16 +- caliban/platform/gke/cli.py | 4 +- caliban/platform/gke/cluster.py | 6 +- caliban/platform/gke/constants.py | 2 +- caliban/platform/gke/util.py | 2 +- caliban/platform/notebook.py | 15 +- caliban/platform/run.py | 3 +- caliban/platform/shell.py | 13 +- caliban/util/__init__.py | 21 + caliban/util/argparse.py | 24 +- tests/caliban/platform/cloud/test_util.py | 155 ++++ .../gke/{test_utils.py => test_util.py} | 81 +- tests/caliban/test_builder.py | 0 tests/caliban/util/test_argparse.py | 42 + tests/caliban/util/test_fs.py | 65 ++ tests/caliban/util/test_tqdm.py | 2 +- tests/caliban/util/test_util.py | 729 ++++++------------ 24 files changed, 687 insertions(+), 628 deletions(-) create mode 100644 tests/caliban/platform/cloud/test_util.py rename tests/caliban/platform/gke/{test_utils.py => test_util.py} (84%) delete mode 100644 tests/caliban/test_builder.py create mode 100644 tests/caliban/util/test_argparse.py diff --git a/caliban/cli.py b/caliban/cli.py index 427c743..7d18ad6 100644 --- a/caliban/cli.py +++ b/caliban/cli.py @@ -24,10 +24,10 @@ from absl.flags import argparse_flags from blessings import Terminal -import caliban.platform.cloud.types as ct import caliban.config as conf import caliban.config.experiment as ce -import caliban.docker as docker +import caliban.docker.build as docker +import caliban.platform.cloud.types as ct import caliban.platform.gke as gke import caliban.platform.gke.constants as gke_k import caliban.platform.gke.types as gke_t @@ -540,7 +540,8 @@ def generate_docker_args(job_mode: conf.JobMode, # Get extra dependencies in case you want to install your requirements via a # setup.py file. - setup_extras = docker.base_extras(job_mode, "setup.py", args.get("extras")) + setup_extras = docker.build.base_extras(job_mode, "setup.py", + args.get("extras")) # Google application credentials, from the CLI or from an env variable. creds_path = conf.extract_cloud_key(args) @@ -553,7 +554,7 @@ def generate_docker_args(job_mode: conf.JobMode, reqs = "requirements.txt" conda_env = "environment.yml" - # Arguments that make their way down to caliban.docker.build_image. + # Arguments that make their way down to caliban.docker.build.build_image. docker_args = { "extra_dirs": args.get("dir"), "requirements_path": reqs if os.path.exists(reqs) else None, diff --git a/caliban/config/__init__.py b/caliban/config/__init__.py index 52f0b04..ed3a05e 100644 --- a/caliban/config/__init__.py +++ b/caliban/config/__init__.py @@ -179,3 +179,15 @@ def apt_packages(conf: CalibanConfig, mode: JobMode) -> List[str]: raise argparse.ArgumentTypeError( """{}'s "apt_packages" entry must be a dictionary or list, not '{}'""". format(CALIBAN_CONFIG, packages)) + + +def caliban_config() -> CalibanConfig: + """Returns a dict that represents a `.calibanconfig.json` file if present, + empty dictionary otherwise. + """ + if not os.path.isfile(CALIBAN_CONFIG): + return {} + + with open(CALIBAN_CONFIG) as f: + conf = commentjson.load(f) + return conf diff --git a/caliban/expansion.py b/caliban/expansion.py index 7659aba..6add41a 100644 --- a/caliban/expansion.py +++ b/caliban/expansion.py @@ -24,7 +24,7 @@ from absl import app, logging from absl.flags import argparse_flags -import caliban.config.experiment as c +import caliban.config.experiment as ce from caliban import __version__ ll.getLogger('caliban.expansion').setLevel(logging.ERROR) @@ -54,7 +54,7 @@ def expansion_parser(): one per line.") parser.add_argument( "experiment_config", - type=c.load_experiment_config, + type=ce.load_experiment_config, help="Path to an experiment config, or 'stdin' to read from stdin.") return parser @@ -69,17 +69,17 @@ def parse_flags(argv): return expansion_parser().parse_args(args) -def _print_flags(expanded: List[c.Experiment]) -> None: +def _print_flags(expanded: List[ce.Experiment]) -> None: """Print the flags associated with each experiment in the supplied expansion list. """ for m in expanded: - flags = c.experiment_to_args(m) + flags = ce.experiment_to_args(m) print(' '.join(flags)) -def _print_json(expanded: List[c.Experiment], pprint: bool = False) -> None: +def _print_json(expanded: List[ce.Experiment], pprint: bool = False) -> None: """Print the list of expanded experiments to stdout; if pprint is true, pretty-prints each JSON dict using an indent of 2, else prints the list with no newlines. @@ -95,7 +95,7 @@ def run_app(args): """ conf = args.experiment_config - expanded = c.expand_experiment_config(conf) + expanded = ce.expand_experiment_config(conf) if args.print_flags: _print_flags(expanded) diff --git a/caliban/history/cli.py b/caliban/history/cli.py index 1b08bc0..6c4be77 100644 --- a/caliban/history/cli.py +++ b/caliban/history/cli.py @@ -15,24 +15,23 @@ # limitations under the License. '''caliban history cli support''' -import os import logging -import pprint as pp -from typing import Optional, Iterable, Dict, Any, List +import os +from typing import Any, Dict, Iterable, List, Optional -from sqlalchemy import or_, and_ +from sqlalchemy import or_ from sqlalchemy.orm import Session -from caliban.util import current_user, Package -from caliban.history.util import (get_sql_engine, session_scope, - update_job_status, get_gke_job_name, stop_job, - replace_job_spec_image) +from caliban.docker.build import build_image +from caliban.docker.push import push_uuid_tag from caliban.history.submit import submit_job_specs -from caliban.history.types import (ContainerSpec, ExperimentGroup, Experiment, - JobSpec, Job, Platform, JobStatus, Platform) -from caliban.platform.gke.util import user_verify, credentials - -from caliban.docker import build_image, push_uuid_tag, execute_jobs +from caliban.history.types import (ContainerSpec, Experiment, ExperimentGroup, + Job, JobStatus, Platform) +from caliban.history.util import (get_gke_job_name, get_sql_engine, + replace_job_spec_image, session_scope, + stop_job, update_job_status) +from caliban.platform.gke.util import credentials, user_verify +from caliban.util import Package, current_user # default max jobs to return for status command _DEFAULT_STATUS_MAX_JOBS = 8 diff --git a/caliban/history/submit.py b/caliban/history/submit.py index f827ee7..64d0323 100644 --- a/caliban/history/submit.py +++ b/caliban/history/submit.py @@ -15,13 +15,12 @@ # limitations under the License. '''caliban utilities for job re-submission''' -from typing import Optional, Iterable, List +from typing import List, Optional -from caliban.history.types import JobSpec, Job, Platform - -import caliban.docker as docker import caliban.platform.cloud.core as cloud import caliban.platform.gke.cli as gke_cli +import caliban.platform.run as r +from caliban.history.types import JobSpec, Platform # ---------------------------------------------------------------------------- @@ -37,7 +36,7 @@ def submit_job_specs( return if platform == Platform.LOCAL: - return docker.execute_jobs(job_specs=specs) + return r.execute_jobs(job_specs=specs) if platform == Platform.CAIP: return cloud.submit_job_specs( diff --git a/caliban/history/util.py b/caliban/history/util.py index e653aad..c701390 100644 --- a/caliban/history/util.py +++ b/caliban/history/util.py @@ -28,13 +28,13 @@ from sqlalchemy.exc import OperationalError from sqlalchemy.orm import Session, sessionmaker -import caliban.config as conf +import caliban.config.experiment as ce +from caliban.history.types import (ContainerSpec, Experiment, ExperimentGroup, + Job, JobSpec, JobStatus, Platform, init_db) from caliban.platform.cloud.types import JobStatus as CloudStatus from caliban.platform.gke.cluster import Cluster from caliban.platform.gke.types import JobStatus as GkeStatus from caliban.platform.gke.util import default_credentials -from caliban.history.types import (ContainerSpec, Experiment, ExperimentGroup, - Job, JobSpec, JobStatus, Platform, init_db) DB_URL_ENV = 'CALIBAN_DB_URL' MEMORY_DB_URL = 'sqlite:///:memory:' @@ -203,7 +203,7 @@ def create_experiments( session: Session, container_spec: ContainerSpec, script_args: List[str], - experiment_config: conf.ExpConf, + experiment_config: ce.ExpConf, xgroup: Optional[str] = None, ) -> List[Experiment]: '''create experiment instances @@ -230,7 +230,7 @@ def create_experiments( container_spec=container_spec, args=script_args, kwargs=kwargs, - ) for kwargs in conf.expand_experiment_config(experiment_config) + ) for kwargs in ce.expand_experiment_config(experiment_config) ] diff --git a/caliban/main.py b/caliban/main.py index b701523..1b34e0d 100644 --- a/caliban/main.py +++ b/caliban/main.py @@ -18,20 +18,22 @@ from __future__ import absolute_import, division, print_function import logging as ll -import os import sys from absl import app, logging from blessings import Terminal import caliban.cli as cli -import caliban.platform.cloud.core as cloud import caliban.config as c -import caliban.docker as docker +import caliban.docker.build as b +import caliban.history.cli +import caliban.platform.cloud.core as cloud import caliban.platform.gke as gke import caliban.platform.gke.cli +import caliban.platform.notebook as pn +import caliban.platform.run as pr +import caliban.platform.shell as ps import caliban.util as u -import caliban.history.cli ll.getLogger('caliban.main').setLevel(logging.ERROR) t = Terminal() @@ -58,29 +60,29 @@ def run_app(arg_input): mount_home = not args['bare'] image_id = args.get("image_id") shell = args['shell'] - docker.run_interactive(job_mode, - image_id=image_id, - run_args=docker_run_args, - mount_home=mount_home, - shell=shell, - **docker_args) + ps.run_interactive(job_mode, + image_id=image_id, + run_args=docker_run_args, + mount_home=mount_home, + shell=shell, + **docker_args) elif command == "notebook": port = args.get("port") lab = args.get("lab") version = args.get("jupyter_version") mount_home = not args['bare'] - docker.run_notebook(job_mode, - port=port, - lab=lab, - version=version, - run_args=docker_run_args, - mount_home=mount_home, - **docker_args) + pn.run_notebook(job_mode, + port=port, + lab=lab, + version=version, + run_args=docker_run_args, + mount_home=mount_home, + **docker_args) elif command == "build": package = args["module"] - docker.build_image(job_mode, package=package, **docker_args) + b.build_image(job_mode, package=package, **docker_args) elif command == 'status': caliban.history.cli.get_status(args) @@ -98,15 +100,15 @@ def run_app(arg_input): exp_config = args.get("experiment_config") xgroup = args.get('xgroup') - docker.run_experiments(job_mode, - run_args=docker_run_args, - script_args=script_args, - image_id=image_id, - experiment_config=exp_config, - dry_run=dry_run, - package=package, - xgroup=xgroup, - **docker_args) + pr.run_experiments(job_mode, + run_args=docker_run_args, + script_args=script_args, + image_id=image_id, + experiment_config=exp_config, + dry_run=dry_run, + package=package, + xgroup=xgroup, + **docker_args) elif command == "cloud": project_id = c.extract_project_id(args) @@ -156,7 +158,7 @@ def main(): except KeyboardInterrupt: logging.info('Shutting down.') sys.exit(0) - except docker.DockerError as e: + except b.DockerError as e: # Handle a failed Docker command. logging.error(t.red(e.message)) logging.error(t.red("Original command: {}".format(e.command))) diff --git a/caliban/platform/cloud/core.py b/caliban/platform/cloud/core.py index 5556266..056c2be 100644 --- a/caliban/platform/cloud/core.py +++ b/caliban/platform/cloud/core.py @@ -31,10 +31,12 @@ from googleapiclient import discovery from googleapiclient.errors import HttpError -import caliban.platform.cloud.types as ct import caliban.config as conf -import caliban.docker as d +import caliban.config.experiment as ce +import caliban.docker.build as db +import caliban.docker.push as dp import caliban.history.types as ht +import caliban.platform.cloud.types as ct import caliban.util as u from caliban.history.util import (create_experiments, generate_container_spec, get_mem_engine, get_sql_engine, session_scope) @@ -459,7 +461,7 @@ def _job_specs( """ for idx, m in enumerate(experiments, 1): - args = conf.experiment_to_args(m.kwargs, m.args) + args = ce.experiment_to_args(m.kwargs, m.args) yield _job_spec(job_name=job_name, idx=idx, training_input={ @@ -527,8 +529,8 @@ def generate_image_tag(project_id, docker_args, dry_run: bool = False): logging.info("Dry run - skipping actual 'docker build' and 'docker push'.") image_tag = "dry_run_tag" else: - image_id = d.build_image(**docker_args) - image_tag = d.push_uuid_tag(project_id, image_id) + image_id = db.build_image(**docker_args) + image_tag = dp.push_uuid_tag(project_id, image_id) return image_tag @@ -582,7 +584,7 @@ def submit_ml_job(job_mode: conf.JobMode, tpu_spec: Optional[ct.TPUSpec] = None, image_tag: Optional[str] = None, labels: Optional[Dict[str, str]] = None, - experiment_config: Optional[conf.ExpConf] = None, + experiment_config: Optional[ce.ExpConf] = None, script_args: Optional[List[str]] = None, request_retries: Optional[int] = None, xgroup: Optional[str] = None) -> None: @@ -599,7 +601,7 @@ def submit_ml_job(job_mode: conf.JobMode, - job_mode: caliban.config.JobMode. - docker_args: these arguments are passed through to - caliban.docker.build_image. + caliban.docker.build.build_image. - region: the region to use for AI Platform job submission. Different regions support different GPUs. - project_id: GCloud project ID for container storage and job submission. diff --git a/caliban/platform/gke/cli.py b/caliban/platform/gke/cli.py index a1c253d..625c4b2 100644 --- a/caliban/platform/gke/cli.py +++ b/caliban/platform/gke/cli.py @@ -31,10 +31,10 @@ import caliban.platform.gke.constants as k import caliban.platform.gke.util as util import caliban.util as u -from caliban.platform.cloud.core import generate_image_tag -from caliban.platform.gke.cluster import Cluster 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 +from caliban.platform.gke.cluster import Cluster # ---------------------------------------------------------------------------- diff --git a/caliban/platform/gke/cluster.py b/caliban/platform/gke/cluster.py index 512f558..014540c 100644 --- a/caliban/platform/gke/cluster.py +++ b/caliban/platform/gke/cluster.py @@ -37,14 +37,14 @@ V1Toleration) from kubernetes.client.api_client import ApiClient -import caliban.config as conf +import caliban.config.experiment as ce import caliban.platform.gke.constants as k import caliban.platform.gke.util as util +from caliban.history.types import Experiment, Job, JobSpec, JobStatus, Platform from caliban.platform.cloud.types import (GPU, TPU, Accelerator, GPUSpec, MachineType, TPUSpec) from caliban.platform.gke.types import NodeImage, OpStatus, ReleaseChannel from caliban.platform.gke.util import trap -from caliban.history.types import Experiment, Job, JobSpec, JobStatus, Platform urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) @@ -697,7 +697,7 @@ def create_simple_job_spec( JobSpec on success, None otherwise """ - args = conf.experiment_to_args(experiment.kwargs, experiment.args) + args = ce.experiment_to_args(experiment.kwargs, experiment.args) # ------------------------------------------------------------------------ # container diff --git a/caliban/platform/gke/constants.py b/caliban/platform/gke/constants.py index d78c4b6..4aadeb7 100644 --- a/caliban/platform/gke/constants.py +++ b/caliban/platform/gke/constants.py @@ -17,8 +17,8 @@ import re -from caliban.platform.cloud.types import GPU, GPUSpec from caliban.config import DEFAULT_MACHINE_TYPE, JobMode +from caliban.platform.cloud.types import GPU, GPUSpec from caliban.platform.gke.types import ReleaseChannel COMPUTE_SCOPE_URL = 'https://www.googleapis.com/auth/compute' diff --git a/caliban/platform/gke/util.py b/caliban/platform/gke/util.py index dea441d..9a754ea 100644 --- a/caliban/platform/gke/util.py +++ b/caliban/platform/gke/util.py @@ -43,7 +43,7 @@ import caliban.platform.gke.constants as k from caliban.platform.cloud.types import GPU, TPU, GPUSpec, TPUSpec -from caliban.platform.gke.types import (CredentialsData, NodeImage, OpStatus) +from caliban.platform.gke.types import CredentialsData, NodeImage, OpStatus # ---------------------------------------------------------------------------- diff --git a/caliban/platform/notebook.py b/caliban/platform/notebook.py index cfdcf46..19041f2 100644 --- a/caliban/platform/notebook.py +++ b/caliban/platform/notebook.py @@ -35,9 +35,10 @@ from tqdm.utils import _screen_shape_wrapper import caliban.config as c +import caliban.config.experiment as ce +import caliban.platform.shell as ps import caliban.util as u import caliban.util.fs as ufs -import caliban.platform.shell as ps from caliban.history.types import Experiment, Job, JobSpec, JobStatus, Platform from caliban.history.util import (create_experiments, generate_container_spec, get_mem_engine, get_sql_engine, session_scope) @@ -703,8 +704,8 @@ def log_job_spec_instance(job_spec: JobSpec, i: int) -> JobSpec: generated from an experiment definition; returns the input job spec. """ - args = c.experiment_to_args(job_spec.experiment.kwargs, - job_spec.experiment.args) + args = ce.experiment_to_args(job_spec.experiment.kwargs, + job_spec.experiment.args) logging.info("") logging.info("Job {} - Experiment args: {}".format(i, t.yellow(str(args)))) return job_spec @@ -746,8 +747,8 @@ def local_callback(idx: int, job: Job) -> None: else: logging.error( t.red(f'Job {idx} failed with return code {job.details["ret_code"]}.')) - args = c.experiment_to_args(job.spec.experiment.kwargs, - job.spec.experiment.args) + args = ce.experiment_to_args(job.spec.experiment.kwargs, + job.spec.experiment.args) logging.error(t.red(f'Failing args for job {idx}: {args}')) @@ -783,7 +784,7 @@ def _create_job_spec_dict( terminal_cmds = ["-e" "PYTHONUNBUFFERED=1"] + window_size_env_cmds() base_cmd = _run_cmd(job_mode, run_args) + terminal_cmds + [image_id] - command = base_cmd + c.experiment_to_args(experiment.kwargs, experiment.args) + command = base_cmd + ce.experiment_to_args(experiment.kwargs, experiment.args) return {'command': command, 'container': image_id} @@ -832,7 +833,7 @@ def run_experiments(job_mode: c.JobMode, script_args: Optional[List[str]] = None, image_id: Optional[str] = None, dry_run: bool = False, - experiment_config: Optional[c.ExpConf] = None, + experiment_config: Optional[ce.ExpConf] = None, xgroup: Optional[str] = None, **build_image_kwargs) -> None: """Builds an image using the supplied **build_image_kwargs and calls `docker diff --git a/caliban/platform/run.py b/caliban/platform/run.py index 0dc0068..9304ee4 100644 --- a/caliban/platform/run.py +++ b/caliban/platform/run.py @@ -35,6 +35,7 @@ from tqdm.utils import _screen_shape_wrapper import caliban.config as c +import caliban.config.experiment as ce import caliban.util as u import caliban.util.fs as ufs from caliban.history.types import Experiment, Job, JobSpec, JobStatus, Platform @@ -831,7 +832,7 @@ def run_experiments(job_mode: c.JobMode, script_args: Optional[List[str]] = None, image_id: Optional[str] = None, dry_run: bool = False, - experiment_config: Optional[c.ExpConf] = None, + experiment_config: Optional[ce.ExpConf] = None, xgroup: Optional[str] = None, **build_image_kwargs) -> None: """Builds an image using the supplied **build_image_kwargs and calls `docker diff --git a/caliban/platform/shell.py b/caliban/platform/shell.py index 8655dfe..e15ff83 100644 --- a/caliban/platform/shell.py +++ b/caliban/platform/shell.py @@ -35,6 +35,7 @@ from tqdm.utils import _screen_shape_wrapper import caliban.config as c +import caliban.config.experiment as ce import caliban.util as u import caliban.util.fs as ufs from caliban.history.types import Experiment, Job, JobSpec, JobStatus, Platform @@ -702,8 +703,8 @@ def log_job_spec_instance(job_spec: JobSpec, i: int) -> JobSpec: generated from an experiment definition; returns the input job spec. """ - args = c.experiment_to_args(job_spec.experiment.kwargs, - job_spec.experiment.args) + args = ce.experiment_to_args(job_spec.experiment.kwargs, + job_spec.experiment.args) logging.info("") logging.info("Job {} - Experiment args: {}".format(i, t.yellow(str(args)))) return job_spec @@ -745,8 +746,8 @@ def local_callback(idx: int, job: Job) -> None: else: logging.error( t.red(f'Job {idx} failed with return code {job.details["ret_code"]}.')) - args = c.experiment_to_args(job.spec.experiment.kwargs, - job.spec.experiment.args) + args = ce.experiment_to_args(job.spec.experiment.kwargs, + job.spec.experiment.args) logging.error(t.red(f'Failing args for job {idx}: {args}')) @@ -782,7 +783,7 @@ def _create_job_spec_dict( terminal_cmds = ["-e" "PYTHONUNBUFFERED=1"] + window_size_env_cmds() base_cmd = _run_cmd(job_mode, run_args) + terminal_cmds + [image_id] - command = base_cmd + c.experiment_to_args(experiment.kwargs, experiment.args) + command = base_cmd + ce.experiment_to_args(experiment.kwargs, experiment.args) return {'command': command, 'container': image_id} @@ -831,7 +832,7 @@ def run_experiments(job_mode: c.JobMode, script_args: Optional[List[str]] = None, image_id: Optional[str] = None, dry_run: bool = False, - experiment_config: Optional[c.ExpConf] = None, + experiment_config: Optional[ce.ExpConf] = None, xgroup: Optional[str] = None, **build_image_kwargs) -> None: """Builds an image using the supplied **build_image_kwargs and calls `docker diff --git a/caliban/util/__init__.py b/caliban/util/__init__.py index 290200e..f027e91 100644 --- a/caliban/util/__init__.py +++ b/caliban/util/__init__.py @@ -178,3 +178,24 @@ def split_by(items: List[str], return items[0:idx], items[idx + 1:] except ValueError: return (items, []) + + +def n_chunks(items: List[Any], n_groups: int) -> List[List[Any]]: + """Returns a list of `n_groups` slices of the original list, guaranteed to + contain all of the original items. + """ + return [items[i::n_groups] for i in range(n_groups)] + + +def chunks_below_limit(items: List[Any], limit: int) -> List[List[Any]]: + """Breaks the input list into a series of chunks guaranteed to be less than""" + quot, _ = divmod(len(items), limit) + return n_chunks(items, quot + 1) + + +def partition(seq: List[str], n: int) -> List[List[str]]: + """Generate groups of n items from seq by scanning across the sequence and + taking chunks of n, offset by 1. + """ + for i in range(0, max(1, len(seq) - n + 1), 1): + yield seq[i:i + n] diff --git a/caliban/util/argparse.py b/caliban/util/argparse.py index 7d36e3f..53f0b2f 100644 --- a/caliban/util/argparse.py +++ b/caliban/util/argparse.py @@ -16,30 +16,14 @@ """ Utilities for our job runner. """ -import caliban.util as u import argparse -import contextlib -import getpass -import io import itertools as it import os -import platform -import re -import shutil -import socket -import subprocess -import sys -import time -import uuid -from collections import ChainMap -from enum import Enum -from typing import (Any, Callable, Dict, Iterable, List, NamedTuple, Optional, - Set, Tuple, Union) - -import tqdm -from absl import logging +from typing import Dict, List, Optional, Tuple + from blessings import Terminal -from tqdm.utils import _term_move_up + +import caliban.util as u t = Terminal() diff --git a/tests/caliban/platform/cloud/test_util.py b/tests/caliban/platform/cloud/test_util.py new file mode 100644 index 0000000..1788c2c --- /dev/null +++ b/tests/caliban/platform/cloud/test_util.py @@ -0,0 +1,155 @@ +#!/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. + +import re + +import hypothesis.strategies as st +from hypothesis import given + +import caliban.platform.cloud.util as u + + +def non_empty_dict(vgen): + return st.dictionaries(st.text(), vgen, min_size=1) + + +def test_key_value_label(): + """unit tests for specific cases of key and value label conversion.""" + assert "face" == u.key_label("--face") + assert "fa_ce" == u.key_label("------fA.!! ce") + + # Empty string roundtrips. + assert "" == u.key_label("") + assert "" == u.value_label("") + + # keys can't have leading digits, just letters, so we append a k. + assert "k0helper" == u.key_label("--0helper") + + # values CAN have leading digits and underscores. + assert "0helper" == u.value_label("--0helper") + assert "_helper" == u.value_label("--_helper") + + +def assert_valid_label(label): + """Assertion that passes if the supplied string is a valid label by all the + rules of Cloud. + + """ + assert len(label) <= 63 + + if label != "": + # check that the output has only lowercase, letters, dashes or + # underscores. + assert re.match('^[a-z0-9_-]+$', label) + + +def assert_valid_key_label(k): + """Assertion that passes if the input is a valid key by all the rules of + Cloud. + + """ + assert_valid_label(k) + if k != "": + assert k[0].isalpha() + + +@given(st.text(min_size=1)) +def test_valid_key_label(s): + cleaned = u.key_label(s) + assert_valid_key_label(cleaned) + + +@given(st.text(min_size=1)) +def test_valid_value_label(s): + cleaned = u.value_label(s) + assert_valid_label(cleaned) + + +def assert_script_args_to_labels(s, m): + """Assertion that passes if the supplied string of arguments parses to a + dictionary that equals the supplied m, representing the expected kv pairs. + + """ + parsed_args = u.script_args_to_labels(s.split(" ")) + assert parsed_args == m + + +def test_script_args_to_labels(): + """unit tests of our script_args_to_labels function behavior.""" + + # Args like --!!! that parse keys to the empty string should not make it + # through. + assert_script_args_to_labels("--lr 1 --!!! 2 --face 3", { + "lr": "1", + "face": "3" + }) + + # Duplicates get overwritten. + assert_script_args_to_labels("--lr 1 --lr 2", {"lr": "2"}) + + assert_script_args_to_labels("--LR 1 --item-label --fa!!ce cake --a", { + "lr": "1", + "item-label": "", + "face": "cake", + "a": "", + }) + + # Multiple values are dropped for now, for the purpose of creating labels. + assert_script_args_to_labels("--lr 1 2 3 --item_underscoRE!! --face cake --a", + { + "lr": "1", + "item_underscore": "", + "face": "cake", + "a": "", + }) + + assert_script_args_to_labels("--face", {"face": ""}) + + # single arguments get ignore if they're not boolean flags. + assert_script_args_to_labels("face", {}) + + +def test_sanitize_labels_kill_empty(): + """Keys that are sanitized to the empty string should NOT make it through.""" + assert {} == u.sanitize_labels([["--!!", "face"]]) + + +@given( + st.one_of(non_empty_dict(st.text()), + st.lists(st.tuples(st.text(), st.text())))) +def test_sanitize_labels(pairs): + """Test that any input we could possibly be provided, as long as it parses into + kv pairs, will only make it into a dict of labels if it's properly + sanitized. + + Checks that the functions works for dicts OR for lists of pairs. + + """ + for k, v in u.sanitize_labels(pairs).items(): + assert_valid_key_label(k) + assert_valid_label(v) + + +@given(st.lists(st.tuples(st.text(), st.text()))) +def test_sanitize_labels_second_noop(pairs): + """Test that passing the output of sanitize_labels back into the function + returns its input. Sanitizing a set of sanitized kv pairs should have no + effect. + + """ + once = u.sanitize_labels(pairs) + twice = u.sanitize_labels(once) + assert once == twice diff --git a/tests/caliban/platform/gke/test_utils.py b/tests/caliban/platform/gke/test_util.py similarity index 84% rename from tests/caliban/platform/gke/test_utils.py rename to tests/caliban/platform/gke/test_util.py index 59d3eba..c27f38d 100644 --- a/tests/caliban/platform/gke/test_utils.py +++ b/tests/caliban/platform/gke/test_util.py @@ -24,9 +24,9 @@ import caliban.platform.cloud.types as ct import caliban.platform.gke.constants as k -import caliban.platform.gke.utils as utils +import caliban.platform.gke.util as util from caliban.platform.gke.types import NodeImage, OpStatus -from caliban.platform.gke.utils import trap +from caliban.platform.gke.util import trap # ---------------------------------------------------------------------------- @@ -42,8 +42,8 @@ def everything_except(excluded_types): # ---------------------------------------------------------------------------- -class UtilsTestSuite(unittest.TestCase): - """tests for caliban.platform.gke.utils""" +class UtilTestSuite(unittest.TestCase): + """tests for caliban.platform.gke.util""" # -------------------------------------------------------------------------- @given( @@ -64,7 +64,7 @@ def test_validate_gpu_spec_against_limits( (gpu_list[i], limits[i]) for i in range(len(limits)) if limits[i] ]) spec = ct.GPUSpec(gpu_type, count) - valid = utils.validate_gpu_spec_against_limits(spec, gpu_limits, 'test') + valid = util.validate_gpu_spec_against_limits(spec, gpu_limits, 'test') if spec.gpu not in gpu_limits: self.assertFalse(valid) @@ -79,7 +79,7 @@ def test_nvidia_daemonset_url(self): VALID_NODE_IMAGES = [NodeImage.COS, NodeImage.UBUNTU] for n in NodeImage: - url = utils.nvidia_daemonset_url(n) + url = util.nvidia_daemonset_url(n) if n in VALID_NODE_IMAGES: self.assertIsNotNone(url) @@ -89,7 +89,7 @@ def test_nvidia_daemonset_url(self): return # -------------------------------------------------------------------------- - @mock.patch('caliban.platform.gke.utils.input', create=True) + @mock.patch('caliban.platform.gke.util.input', create=True) @given(st.lists(st.from_regex('^[^yYnN]+$'), min_size=0, max_size=8)) def test_user_verify( self, @@ -103,18 +103,18 @@ def test_user_verify( # default input mocked_input.side_effect = [''] - self.assertEqual(utils.user_verify('test default', default=default), + self.assertEqual(util.user_verify('test default', default=default), default) # upper/lower true input for x in ['y', 'Y']: mocked_input.side_effect = invalid_strings + [x] - self.assertTrue(utils.user_verify('y input', default=default)) + self.assertTrue(util.user_verify('y input', default=default)) # upper/lower false input for x in ['n', 'N']: mocked_input.side_effect = invalid_strings + [x] - self.assertFalse(utils.user_verify('n input', default=default)) + self.assertFalse(util.user_verify('n input', default=default)) return @@ -204,11 +204,11 @@ def _return_results(): # to take about a factor of 100 longer # empty condition list - self.assertIsNone(utils.wait_for_operation(api, 'name', [], spinner=False)) + self.assertIsNone(util.wait_for_operation(api, 'name', [], spinner=False)) # exception self.assertIsNone( - utils.wait_for_operation(api, 'name', list(conds), 0, spinner=False)) + util.wait_for_operation(api, 'name', list(conds), 0, spinner=False)) # normal operation rsp_generator = _return_results() @@ -222,14 +222,14 @@ def _return_results(): if expected_response is not None: self.assertEqual({'status': expected_response}, - utils.wait_for_operation(api, - 'name', - list(conds), - 0, - spinner=False)) + util.wait_for_operation(api, + 'name', + list(conds), + 0, + spinner=False)) else: self.assertIsNone( - utils.wait_for_operation(api, 'name', list(conds), 0, spinner=False)) + util.wait_for_operation(api, 'name', list(conds), 0, spinner=False)) return @@ -277,11 +277,11 @@ def _invalid_response(): # exception handling api.execute = _raises - self.assertIsNone(utils.get_zone_tpu_types(api, 'p', 'z')) + self.assertIsNone(util.get_zone_tpu_types(api, 'p', 'z')) # invalid response api.execute = _invalid_response - self.assertIsNone(utils.get_zone_tpu_types(api, 'p', 'z')) + self.assertIsNone(util.get_zone_tpu_types(api, 'p', 'z')) # normal mode api.execute = _response @@ -289,7 +289,7 @@ def _invalid_response(): sorted(tpus), sorted([ '{}-{}'.format(x.name.lower(), x.count) - for x in utils.get_zone_tpu_types(api, 'p', 'z') + for x in util.get_zone_tpu_types(api, 'p', 'z') ])) return @@ -302,7 +302,7 @@ def test_sanitize_job_name(self, job_name): def valid(x): return k.DNS_1123_RE.match(x) is not None - sanitized = utils.sanitize_job_name(job_name) + sanitized = util.sanitize_job_name(job_name) if valid(job_name): self.assertEqual(job_name, sanitized) @@ -310,7 +310,7 @@ def valid(x): self.assertTrue(valid(sanitized)) # idempotency check - self.assertEqual(sanitized, utils.sanitize_job_name(sanitized)) + self.assertEqual(sanitized, util.sanitize_job_name(sanitized)) return @@ -359,11 +359,11 @@ def _invalid_response(): # exception handling api.execute = _raises - self.assertIsNone(utils.get_zone_gpu_types(api, 'p', 'z')) + self.assertIsNone(util.get_zone_gpu_types(api, 'p', 'z')) # invalid response api.execute = _invalid_response - self.assertIsNone(utils.get_zone_gpu_types(api, 'p', 'z')) + self.assertIsNone(util.get_zone_gpu_types(api, 'p', 'z')) # normal execution api.execute = _response @@ -374,7 +374,7 @@ def _invalid_response(): ]), sorted([ 'nvidia-tesla-{}-{}'.format(x.gpu.name.lower(), x.count) - for x in utils.get_zone_gpu_types(api, 'p', 'z') + for x in util.get_zone_gpu_types(api, 'p', 'z') ])) return @@ -414,16 +414,15 @@ def _invalid(): # exception handling api.execute = _raises - self.assertIsNone(utils.get_region_quotas(api, 'p', 'r')) + self.assertIsNone(util.get_region_quotas(api, 'p', 'r')) # invalid return api.execute = _invalid - self.assertEqual([], utils.get_region_quotas(api, 'p', 'r')) + self.assertEqual([], util.get_region_quotas(api, 'p', 'r')) # normal execution api.execute = _normal - self.assertEqual(_normal()['quotas'], - utils.get_region_quotas(api, 'p', 'r')) + self.assertEqual(_normal()['quotas'], util.get_region_quotas(api, 'p', 'r')) return @@ -462,11 +461,11 @@ def _invalid(): # exception handling api.execute = _raises - self.assertIsNone(utils.generate_resource_limits(api, 'p', 'r')) + self.assertIsNone(util.generate_resource_limits(api, 'p', 'r')) # invalid return api.execute = _invalid - self.assertEqual([], utils.generate_resource_limits(api, 'p', 'r')) + self.assertEqual([], util.generate_resource_limits(api, 'p', 'r')) # normal execution api.execute = _normal @@ -482,7 +481,7 @@ def _invalid(): 'maximum': str(quotas[1]['limit']) }]) - self.assertEqual(expected, utils.generate_resource_limits(api, 'p', 'r')) + self.assertEqual(expected, util.generate_resource_limits(api, 'p', 'r')) return @@ -516,16 +515,16 @@ def list_clusters(self, project_id, zone): api.throws = True # exception handling - self.assertIsNone(utils.get_gke_cluster(api, 'foo', 'p')) + self.assertIsNone(util.get_gke_cluster(api, 'foo', 'p')) api.throws = False # single cluster if len(names) > 0: cname = names[random.randint(0, len(names) - 1)] - self.assertEqual(cname, utils.get_gke_cluster(api, cname, 'p').name) + self.assertEqual(cname, util.get_gke_cluster(api, cname, 'p').name) # name not in name list - self.assertIsNone(utils.get_gke_cluster(api, invalid, 'p')) + self.assertIsNone(util.get_gke_cluster(api, invalid, 'p')) return @@ -568,13 +567,13 @@ def _validate_nonnull_list(self, lst: list, ref: list): values=everything(), )) def test_nonnull_dict(self, input_dict): - self._validate_nonnull_dict(utils.nonnull_dict(input_dict), input_dict) + self._validate_nonnull_dict(util.nonnull_dict(input_dict), input_dict) return # -------------------------------------------------------------------------- @given(st.lists(everything())) def test_nonnull_list(self, input_list): - self._validate_nonnull_list(utils.nonnull_list(input_list), input_list) + self._validate_nonnull_list(util.nonnull_list(input_list), input_list) return # -------------------------------------------------------------------------- @@ -609,12 +608,12 @@ def _invalid(): # exception handling api.execute = _raises - self.assertIsNone(utils.get_zones_in_region(api, 'p', region)) + self.assertIsNone(util.get_zones_in_region(api, 'p', region)) # invalid return api.execute = _invalid - self.assertIsNone(utils.get_zones_in_region(api, 'p', region)) + self.assertIsNone(util.get_zones_in_region(api, 'p', region)) # normal execution api.execute = _normal - self.assertEqual(zones, utils.get_zones_in_region(api, 'p', region)) + self.assertEqual(zones, util.get_zones_in_region(api, 'p', region)) diff --git a/tests/caliban/test_builder.py b/tests/caliban/test_builder.py deleted file mode 100644 index e69de29..0000000 diff --git a/tests/caliban/util/test_argparse.py b/tests/caliban/util/test_argparse.py new file mode 100644 index 0000000..977227d --- /dev/null +++ b/tests/caliban/util/test_argparse.py @@ -0,0 +1,42 @@ +#!/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. + +from collections import OrderedDict + +import caliban.util.argparse as ua + + +def test_expand_args(): + m = OrderedDict([("a", "item"), ("b", None), ("c", "d")]) + expanded = ua.expand_args(m) + + # None is excluded from the results. + assert expanded == ["a", "item", "b", "c", "d"] + + +def test_is_key(): + """A key is anything that starts with a dash; nothing else! + + """ + assert ua.is_key("--face") + assert ua.is_key("-f") + assert not ua.is_key("") + assert not ua.is_key("face") + assert not ua.is_key("f") + + # this should never happen, but what the heck, why not test that it's a + # fine thing, accepted yet strange. + assert ua.is_key("-----face") diff --git a/tests/caliban/util/test_fs.py b/tests/caliban/util/test_fs.py index b9ab88f..1eece50 100644 --- a/tests/caliban/util/test_fs.py +++ b/tests/caliban/util/test_fs.py @@ -19,6 +19,71 @@ import caliban.util.fs as ufs +def test_module_to_path(): + """verify that we can go the other way and turn modules back into expected + relative paths. + + """ + m = { + # normal modules get nesting. + "face.cake": "face/cake.py", + + # root-level modules just get a py extension. + "face": "face.py", + + # This will get treated as a module nested inside of a folder, which is + # clearly invalid; marking this behavior in the tests. + "face/cake.py": "face/cake/py.py" + } + for k in m: + assert ufs.module_to_path(k) == m[k] + + +def test_generate_package(): + """validate that the generate_package function can handle all sorts of inputs + and generate valid Package objects. + + """ + m = { + # normal module syntax should just work. + "caliban.cli": + ufs.module_package("caliban.cli"), + + # This one is controversial, maybe... if something exists as a module + # if you replace slashes with dots, THEN it will also parse as a + # module. If it exists as a file in its own right this won't happen. + # + # TODO get a test in for this final claim using temp directories. + "caliban/cli": + ufs.module_package("caliban.cli"), + + # root scripts or packages should require the entire local directory. + "setup": + ufs.module_package("setup"), + "cake.py": + ufs.script_package("cake.py", "python"), + + # This is busted but should still parse. + "face.cake.py": + ufs.script_package("face.cake.py", "python"), + + # Paths into directories should parse properly into modules and include + # the root as their required package to import. + "face/cake.py": + ufs.script_package("face/cake.py", "python"), + + # Deeper nesting works. + "face/cake/cheese.py": + ufs.script_package("face/cake/cheese.py", "python"), + + # Other executables work. + "face/cake/cheese.sh": + ufs.script_package("face/cake/cheese.sh"), + } + for k in m: + assert ufs.generate_package(k) == m[k] + + def test_capture_stdout(): buf = io.StringIO() ret_string, code = ufs.capture_stdout(["echo", "hello!"], file=buf) diff --git a/tests/caliban/util/test_tqdm.py b/tests/caliban/util/test_tqdm.py index ebd6e37..901d9ef 100644 --- a/tests/caliban/util/test_tqdm.py +++ b/tests/caliban/util/test_tqdm.py @@ -16,7 +16,7 @@ import io -from tqdm._utils import _term_move_up +from tqdm.utils import _term_move_up import caliban.util.tqdm as ut diff --git a/tests/caliban/util/test_util.py b/tests/caliban/util/test_util.py index 2c99318..e9777da 100644 --- a/tests/caliban/util/test_util.py +++ b/tests/caliban/util/test_util.py @@ -14,19 +14,16 @@ # See the License for the specific language governing permissions and # limitations under the License. -import io import itertools -import re -import unittest from collections import OrderedDict from enum import Enum from typing import Union import hypothesis.strategies as st from hypothesis import given -from tqdm._utils import _term_move_up import caliban.util as u +import pytest text_set = st.sets(st.text(), min_size=1) ne_text_set = st.sets(st.text(min_size=1), min_size=1) @@ -36,476 +33,254 @@ def non_empty_dict(vgen): return st.dictionaries(st.text(), vgen, min_size=1) -def test_carriage_return(): - - def through(xs): - buf = io.StringIO() - f = u.TqdmFile(file=buf) - - for x in xs: - f.write(x) - f.flush() - - return buf.getvalue() - - # Strings pass through tqdmfile with no newline attached. - assert through(["Yo!"]) == "Yo!" - - # Empty lines do nothing. - assert through(["", "", ""]) == "" - - # A carriage return is converted to a newline, but the next line, if it's - # written, will have the proper prefix to trigger a carriage return. - assert through(["Yo!\r"]) == "Yo!\n" - - # Boom, triggered. - assert through(["Yo!\r", "continue"]) == f"Yo!\n{_term_move_up()}\rcontinue" - - -class UtilTestSuite(unittest.TestCase): - """Tests for the util package.""" - - @given(ne_text_set, ne_text_set) - def test_enum_vals(self, ks, vs): - """Setup ensures that the values are unique.""" - m = dict(zip(ks, vs)) - enum = Enum('TestEnum', m) - - # enum_vals returns the values from the enum. - self.assertListEqual(list(m.values()), u.enum_vals(enum)) - - def test_any_of_unit(self): - MyEnum = Enum('MyEnum', {"a": "a_string", "b": "b_string"}) - SecondEnum = Enum('SecondEnum', {"c": "c_cake", "d": "d_face"}) - SomeEnum = Union[MyEnum, SecondEnum] - - # Asking for a value not in ANY enum raises a value error. - with self.assertRaises(ValueError): - u.any_of("face", SomeEnum) - - @given(ne_text_set, ne_text_set, ne_text_set, ne_text_set) - def test_any_of(self, k1, v1, k2, v2): - m1 = dict(zip(k1, v1)) - m2 = dict(zip(k2, v2)) - enum1 = Enum('enum1', m1) - enum2 = Enum('enum2', m2) - union = Union[enum1, enum2] - - # If the item appears in the first map any_of will return it. - for k, v in m1.items(): - self.assertEqual(u.any_of(v, union), enum1(v)) - - for k, v in m2.items(): - # If a value from the second enum appears in enum1 any_of will return it; - # else, it'll return the value from enum2. - try: - expected = enum1(v) - except ValueError: - expected = enum2(v) - - self.assertEqual(u.any_of(v, union), expected) - - def test_dict_product(self): - input_m = OrderedDict([("a", [1, 2, 3]), ("b", [4, 5]), ("c", "d")]) - result = list(u.dict_product(input_m)) - - expected = [{ - 'a': 1, - 'b': 4, - 'c': 'd' - }, { - 'a': 1, - 'b': 5, - 'c': 'd' - }, { - 'a': 2, - 'b': 4, - 'c': 'd' - }, { - 'a': 2, - 'b': 5, - 'c': 'd' - }, { - 'a': 3, - 'b': 4, - 'c': 'd' - }, { - 'a': 3, - 'b': 5, - 'c': 'd' - }] - - self.assertListEqual(result, expected) - - @given(st.dictionaries(st.text(), st.text()), - st.dictionaries(st.text(), st.text())) - def test_merge(self, m1, m2): - merged = u.merge(m1, m2) - - # Every item from the second map should be in the merged map. - for k, v in m2.items(): - self.assertEqual(merged[k], v) - - # Every item from the first map should be in the merged map, OR, if it - # shares a key with m2, m2's value will have bumped it. - for k, v in m1.items(): - self.assertEqual(merged[k], m2.get(k, v)) - - def test_flipm_unit(self): - m = {"a": {1: "a_one", 2: "a_two"}, "b": {1: "b_one", 3: "b_three"}} - expected = { - 1: { - "a": "a_one", - "b": "b_one" - }, - 2: { - "a": "a_two" - }, - 3: { - "b": "b_three" - } - } - - # Flipping does what we expect! - self.assertDictEqual(u.flipm(m), expected) - - @given( - st.dictionaries(st.text(), - st.dictionaries(st.text(), st.text(), min_size=1))) - def test_flipm(self, m): - # As long as an inner dictionary isn't empty, flipping is invertible. - self.assertDictEqual(m, u.flipm(u.flipm(m))) - - @given(st.sets(st.text())) - def test_flipm_empty_values(self, ks): - """Flipping a dictionary with empty values always equals the empty map.""" - m = u.dict_by(ks, lambda k: {}) - self.assertDictEqual({}, u.flipm(m)) - - def test_invertm_unit(self): - m = {"a": [1, 2, 3], "b": [2, 3, 4]} - expected = { - 1: {"a"}, - 2: {"a", "b"}, - 3: {"a", "b"}, - 4: {"b"}, - } - self.assertDictEqual(u.invertm(m), expected) - - @given(non_empty_dict(non_empty_dict(text_set))) - def test_reorderm(self, m): - - def invert_inner(d): - return {k: u.invertm(v) for k, v in d.items()} - - flipped = u.flipm(m) - - # flipping the inner map. - self.assertDictEqual(u.reorderm(m, (0, 2, 1)), invert_inner(m)) - - # Flipping the outer keys is equivalent to calling flipm once. - self.assertDictEqual(u.reorderm(m, (1, 0, 2)), flipped) - - # Reordering the inner keys is equiv to flipping the outer keys the - # flipping the new inner dictionary. - self.assertDictEqual(u.reorderm(m, (1, 2, 0)), invert_inner(flipped)) - - # Flipping again brings the original list entry out. - self.assertDictEqual(u.reorderm(m, (2, 1, 0)), - u.flipm(invert_inner(flipped))) - - @given(non_empty_dict(text_set)) - def test_invertm(self, m): - self.assertDictEqual(m, u.invertm(u.invertm(m))) - - @given(st.sets(st.text())) - def test_dict_by(self, xs): - """dict_by should apply a function to each item in a set to generate the values - of the returned dict. - - """ - m = u.dict_by(xs, len) - - # every value is properly constructed - for k, v in m.items(): - self.assertEqual(len(k), v) - - # the key set of the dict is equal to the incoming original set. - self.assertSetEqual(set(m.keys()), xs) - - def test_expand_args(self): - m = OrderedDict([("a", "item"), ("b", None), ("c", "d")]) - expanded = u.expand_args(m) - - # None is excluded from the results. - self.assertListEqual(expanded, ["a", "item", "b", "c", "d"]) - - def test_generate_package(self): - """validate that the generate_package function can handle all sorts of inputs - and generate valid Package objects. - - """ - m = { - # normal module syntax should just work. - "caliban.util": - u.module_package("caliban.util"), - - # This one is controversial, maybe... if something exists as a module - # if you replace slashes with dots, THEN it will also parse as a - # module. If it exists as a file in its own right this won't happen. - # - # TODO get a test in for this final claim using temp directories. - "caliban/util": - u.module_package("caliban.util"), - - # root scripts or packages should require the entire local directory. - "setup": - u.module_package("setup"), - "cake.py": - u.script_package("cake.py", "python"), - - # This is busted but should still parse. - "face.cake.py": - u.script_package("face.cake.py", "python"), - - # Paths into directories should parse properly into modules and include - # the root as their required package to import. - "face/cake.py": - u.script_package("face/cake.py", "python"), - - # Deeper nesting works. - "face/cake/cheese.py": - u.script_package("face/cake/cheese.py", "python"), - - # Other executables work. - "face/cake/cheese.sh": - u.script_package("face/cake/cheese.sh"), - } - for k in m: - self.assertEqual(u.generate_package(k), m[k]) - - def test_module_to_path(self): - """verify that we can go the other way and turn modules back into expected - relative paths. - - """ - m = { - # normal modules get nesting. - "face.cake": "face/cake.py", - - # root-level modules just get a py extension. - "face": "face.py", - - # This will get treated as a module nested inside of a folder, which is - # clearly invalid; marking this behavior in the tests. - "face/cake.py": "face/cake/py.py" - } - for k in m: - self.assertEqual(u.module_to_path(k), m[k]) - - def test_is_key(self): - """A key is anything that starts with a dash; nothing else! - - """ - self.assertTrue(u._is_key("--face")) - self.assertTrue(u._is_key("-f")) - self.assertFalse(u._is_key("")) - self.assertFalse(u._is_key("face")) - self.assertFalse(u._is_key("f")) - - # this should never happen, but what the heck, why not test that it's a - # fine thing, accepted yet strange. - self.assertTrue(u._is_key("-----face")) - - def test_key_value_label(self): - """unit tests for specific cases of key and value label conversion.""" - self.assertEqual("face", u.key_label("--face")) - self.assertEqual("fa_ce", u.key_label("------fA.!! ce")) - - # Empty string roundtrips. - self.assertEqual("", u.key_label("")) - self.assertEqual("", u.value_label("")) - - # keys can't have leading digits, just letters, so we append a k. - self.assertEqual("k0helper", u.key_label("--0helper")) - - # values CAN have leading digits and underscores. - self.assertEqual("0helper", u.value_label("--0helper")) - self.assertEqual("_helper", u.value_label("--_helper")) - - def assertValidLabel(self, label): - """Assertion that passes if the supplied string is a valid label by all the - rules of Cloud. - - """ - self.assertLessEqual(len(label), 63) - - if label != "": - # check that the output has only lowercase, letters, dashes or - # underscores. - self.assertTrue(re.match('^[a-z0-9_-]+$', label)) - - def assertValidKeyLabel(self, k): - """Assertion that passes if the input is a valid key by all the rules of - Cloud. - - """ - self.assertValidLabel(k) - if k != "": - self.assertTrue(k[0].isalpha()) - - @given(st.text(min_size=1)) - def test_valid_key_label(self, s): - cleaned = u.key_label(s) - self.assertValidKeyLabel(cleaned) - - @given(st.text(min_size=1)) - def test_valid_value_label(self, s): - cleaned = u.value_label(s) - self.assertValidLabel(cleaned) - - @given(st.lists(st.integers()), st.integers(min_value=1, max_value=500)) - def test_n_chunks(self, xs, n): - singletons = list(map(lambda x: [x], xs)) - - # If the chunks equal the length we get all singletons. - self.assertListEqual(u.n_chunks(xs, len(xs)), singletons) - - # one chunk returns a single singleton. - self.assertListEqual(u.n_chunks(xs, 1), [xs]) - - sharded = u.n_chunks(xs, n) - recombined = list(itertools.chain(*sharded)) - - # The ordering might not be the same, but the total number of items is the - # same if we break down and recombine. - self.assertEqual(len(recombined), len(xs)) - - # And the items are equal too. - self.assertSetEqual(set(xs), set(recombined)) - - def test_chunks_below_limit(self): - xs = [0, 1, 2, 3, 4, 5] - - # Below the limit, there's no breakdown. - self.assertListEqual([xs], u.chunks_below_limit(xs, 100)) - - # Below the limit, there's no breakdown. - shards = [[0, 2, 4], [1, 3, 5]] - self.assertListEqual(shards, u.chunks_below_limit(xs, 5)) - - # You can recover the original list by zipping together the shards (if they - # happen to be equal in length, as here.) - self.assertListEqual(xs, list(itertools.chain(*list(zip(*shards))))) - - @given(st.lists(st.integers(), min_size=1)) - def test_partition_first_items(self, xs): - """retrieving the first item of each grouping recovers the original list.""" - rt = list(map(lambda pair: pair[0], u.partition(xs, 1))) - self.assertListEqual(rt, xs) - - @given(st.lists(st.integers(), min_size=1)) - def test_partition_by_one_gives_singletons(self, xs): - """Partition by 1 gives a list of singletons.""" - singletons = list(u.partition(xs, 1)) - self.assertListEqual(singletons, [[x] for x in xs]) - - @given(st.lists(st.integers(), min_size=1), st.integers(min_value=0)) - def test_partition_by_big_gives_singleton(self, xs, n): - """partitioning by a number >= the list length returns a singleton containing - just the list. - - """ - one_entry = list(u.partition(xs, len(xs) + n)) - self.assertListEqual(one_entry, [xs]) - - def test_partition(self): - """Various unittests of partition.""" - - # partitioning by 1 generates singletons. - self.assertListEqual(list(u.partition([1, 2, 3, 4, 5], 1)), - [[1], [2], [3], [4], [5]]) - - # partition works into groups of 2 - self.assertListEqual(list(u.partition([1, 2, 3, 4], 2)), - [[1, 2], [2, 3], [3, 4]]) - - # >= case - self.assertListEqual(list(u.partition([1, 2, 3], 3)), [[1, 2, 3]]) - self.assertListEqual(list(u.partition([1, 2, 3], 10)), [[1, 2, 3]]) - - def assertScriptArgsToLabels(self, s, m): - """Assertion that passes if the supplied string of arguments parses to a - dictionary that equals the supplied m, representing the expected kv pairs. - - """ - parsed_args = u.script_args_to_labels(s.split(" ")) - self.assertDictEqual(parsed_args, m) - - def test_script_args_to_labels(self): - """unit tests of our script_args_to_labels function behavior.""" - - # Args like --!!! that parse keys to the empty string should not make it - # through. - self.assertScriptArgsToLabels("--lr 1 --!!! 2 --face 3", { - "lr": "1", - "face": "3" - }) - - # Duplicates get overwritten. - self.assertScriptArgsToLabels("--lr 1 --lr 2", {"lr": "2"}) - - self.assertScriptArgsToLabels("--LR 1 --item-label --fa!!ce cake --a", { - "lr": "1", - "item-label": "", - "face": "cake", - "a": "", - }) - - # Multiple values are dropped for now, for the purpose of creating labels. - self.assertScriptArgsToLabels( - "--lr 1 2 3 --item_underscoRE!! --face cake --a", { - "lr": "1", - "item_underscore": "", - "face": "cake", - "a": "", - }) - - self.assertScriptArgsToLabels("--face", {"face": ""}) - - # single arguments get ignore if they're not boolean flags. - self.assertScriptArgsToLabels("face", {}) - - def test_sanitize_labels_kill_empty(self): - """Keys that are sanitized to the empty string should NOT make it through.""" - self.assertDictEqual({}, u.sanitize_labels([["--!!", "face"]])) - - @given( - st.one_of(non_empty_dict(st.text()), - st.lists(st.tuples(st.text(), st.text())))) - def test_sanitize_labels(self, pairs): - """Test that any input we could possibly be provided, as long as it parses into - kv pairs, will only make it into a dict of labels if it's properly - sanitized. - - Checks that the functions works for dicts OR for lists of pairs. - - """ - for k, v in u.sanitize_labels(pairs).items(): - self.assertValidKeyLabel(k) - self.assertValidLabel(v) - - @given(st.lists(st.tuples(st.text(), st.text()))) - def test_sanitize_labels_second_noop(self, pairs): - """Test that passing the output of sanitize_labels back into the function - returns its input. Sanitizing a set of sanitized kv pairs should have no - effect. - - """ - once = u.sanitize_labels(pairs) - twice = u.sanitize_labels(once) - self.assertDictEqual(once, twice) - - -if __name__ == '__main__': - unittest.main() +@given(ne_text_set, ne_text_set) +def test_enum_vals(ks, vs): + """Setup ensures that the values are unique.""" + m = dict(zip(ks, vs)) + enum = Enum('TestEnum', m) + + # enum_vals returns the values from the enum. + assert list(m.values()) == u.enum_vals(enum) + + +def test_any_of_unit(): + MyEnum = Enum('MyEnum', {"a": "a_string", "b": "b_string"}) + SecondEnum = Enum('SecondEnum', {"c": "c_cake", "d": "d_face"}) + SomeEnum = Union[MyEnum, SecondEnum] + + # Asking for a value not in ANY enum raises a value error. + with pytest.raises(ValueError): + u.any_of("face", SomeEnum) + + +@given(ne_text_set, ne_text_set, ne_text_set, ne_text_set) +def test_any_of(k1, v1, k2, v2): + m1 = dict(zip(k1, v1)) + m2 = dict(zip(k2, v2)) + enum1 = Enum('enum1', m1) + enum2 = Enum('enum2', m2) + union = Union[enum1, enum2] + + # If the item appears in the first map any_of will return it. + for k, v in m1.items(): + assert u.any_of(v, union) == enum1(v) + + for k, v in m2.items(): + # If a value from the second enum appears in enum1 any_of will return it; + # else, it'll return the value from enum2. + try: + expected = enum1(v) + except ValueError: + expected = enum2(v) + + assert u.any_of(v, union) == expected + + +def test_dict_product(): + input_m = OrderedDict([("a", [1, 2, 3]), ("b", [4, 5]), ("c", "d")]) + result = list(u.dict_product(input_m)) + + expected = [{ + 'a': 1, + 'b': 4, + 'c': 'd' + }, { + 'a': 1, + 'b': 5, + 'c': 'd' + }, { + 'a': 2, + 'b': 4, + 'c': 'd' + }, { + 'a': 2, + 'b': 5, + 'c': 'd' + }, { + 'a': 3, + 'b': 4, + 'c': 'd' + }, { + 'a': 3, + 'b': 5, + 'c': 'd' + }] + + assert result == expected + + +@given(st.dictionaries(st.text(), st.text()), + st.dictionaries(st.text(), st.text())) +def test_merge(m1, m2): + merged = u.merge(m1, m2) + + # Every item from the second map should be in the merged map. + for k, v in m2.items(): + assert merged[k] == v + + # Every item from the first map should be in the merged map, OR, if it + # shares a key with m2, m2's value will have bumped it. + for k, v in m1.items(): + assert merged[k] == m2.get(k, v) + + +def test_flipm_unit(): + m = {"a": {1: "a_one", 2: "a_two"}, "b": {1: "b_one", 3: "b_three"}} + expected = { + 1: { + "a": "a_one", + "b": "b_one" + }, + 2: { + "a": "a_two" + }, + 3: { + "b": "b_three" + } + } + + # Flipping does what we expect! + assert u.flipm(m) == expected + + +@given( + st.dictionaries(st.text(), st.dictionaries(st.text(), st.text(), + min_size=1))) +def test_flipm(m): + # As long as an inner dictionary isn't empty, flipping is invertible. + assert m == u.flipm(u.flipm(m)) + + +@given(st.sets(st.text())) +def test_flipm_empty_values(ks): + """Flipping a dictionary with empty values always equals the empty map.""" + m = u.dict_by(ks, lambda k: {}) + assert {} == u.flipm(m) + + +def test_invertm_unit(): + m = {"a": [1, 2, 3], "b": [2, 3, 4]} + expected = { + 1: {"a"}, + 2: {"a", "b"}, + 3: {"a", "b"}, + 4: {"b"}, + } + assert u.invertm(m) == expected + + +@given(non_empty_dict(non_empty_dict(text_set))) +def test_reorderm(m): + + def invert_inner(d): + return {k: u.invertm(v) for k, v in d.items()} + + flipped = u.flipm(m) + + # flipping the inner map. + assert u.reorderm(m, (0, 2, 1)) == invert_inner(m) + + # Flipping the outer keys is equivalent to calling flipm once. + assert u.reorderm(m, (1, 0, 2)) == flipped + + # Reordering the inner keys is equiv to flipping the outer keys the + # flipping the new inner dictionary. + assert u.reorderm(m, (1, 2, 0)) == invert_inner(flipped) + + # Flipping again brings the original list entry out. + assert u.reorderm(m, (2, 1, 0)) == u.flipm(invert_inner(flipped)) + + +@given(non_empty_dict(text_set)) +def test_invertm(m): + assert m == u.invertm(u.invertm(m)) + + +@given(st.sets(st.text())) +def test_dict_by(xs): + """dict_by should apply a function to each item in a set to generate the values + of the returned dict. + + """ + m = u.dict_by(xs, len) + + # every value is properly constructed + for k, v in m.items(): + assert len(k) == v + + # the key set of the dict is equal to the incoming original set. + assert set(m.keys()) == xs + + +@given(st.lists(st.integers()), st.integers(min_value=1, max_value=500)) +def test_n_chunks(xs, n): + singletons = list(map(lambda x: [x], xs)) + + # If the chunks equal the length we get all singletons. + assert u.n_chunks(xs, len(xs)) == singletons + + # one chunk returns a single singleton. + assert u.n_chunks(xs, 1) == [xs] + + sharded = u.n_chunks(xs, n) + recombined = list(itertools.chain(*sharded)) + + # The ordering might not be the same, but the total number of items is the + # same if we break down and recombine. + assert len(recombined) == len(xs) + + # And the items are equal too. + assert set(xs) == set(recombined) + + +def test_chunks_below_limit(): + xs = [0, 1, 2, 3, 4, 5] + + # Below the limit, there's no breakdown. + assert [xs] == u.chunks_below_limit(xs, 100) + + # Below the limit, there's no breakdown. + shards = [[0, 2, 4], [1, 3, 5]] + assert shards == u.chunks_below_limit(xs, 5) + + # You can recover the original list by zipping together the shards (if they + # happen to be equal in length, as here.) + assert xs == list(itertools.chain(*list(zip(*shards)))) + + +@given(st.lists(st.integers(), min_size=1)) +def test_partition_first_items(xs): + """retrieving the first item of each grouping recovers the original list.""" + rt = list(map(lambda pair: pair[0], u.partition(xs, 1))) + assert rt == xs + + +@given(st.lists(st.integers(), min_size=1)) +def test_partition_by_one_gives_singletons(xs): + """Partition by 1 gives a list of singletons.""" + singletons = list(u.partition(xs, 1)) + assert singletons == [[x] for x in xs] + + +@given(st.lists(st.integers(), min_size=1), st.integers(min_value=0)) +def test_partition_by_big_gives_singleton(xs, n): + """partitioning by a number >= the list length returns a singleton containing + just the list. + + """ + one_entry = list(u.partition(xs, len(xs) + n)) + assert one_entry == [xs] + + +def test_partition(): + """Various unittests of partition.""" + + # partitioning by 1 generates singletons. + assert list(u.partition([1, 2, 3, 4, 5], 1)) == [[1], [2], [3], [4], [5]] + + # partition works into groups of 2 + assert list(u.partition([1, 2, 3, 4], 2)) == [[1, 2], [2, 3], [3, 4]] + + # >= case + assert list(u.partition([1, 2, 3], 3)), [[1, 2, 3]] + assert list(u.partition([1, 2, 3], 10)) == [[1, 2, 3]] From 88f675ee674a94f4f26614d3f2043cf550412baa Mon Sep 17 00:00:00 2001 From: Sam Ritchie Date: Sat, 11 Jul 2020 13:25:44 -0600 Subject: [PATCH 07/15] getting closer --- caliban/platform/notebook.py | 922 +--------------------------------- caliban/platform/run.py | 648 +----------------------- caliban/platform/shell.py | 924 +---------------------------------- 3 files changed, 24 insertions(+), 2470 deletions(-) diff --git a/caliban/platform/notebook.py b/caliban/platform/notebook.py index 19041f2..a7111ba 100644 --- a/caliban/platform/notebook.py +++ b/caliban/platform/notebook.py @@ -18,931 +18,17 @@ """ -from __future__ import absolute_import, division, print_function +from typing import List, Optional -import json -import os -import subprocess -import sys -from enum import Enum -from pathlib import Path -from typing import (Any, Callable, Dict, Iterable, List, NamedTuple, NewType, - Optional, Union) - -import tqdm -from absl import logging from blessings import Terminal -from tqdm.utils import _screen_shape_wrapper import caliban.config as c -import caliban.config.experiment as ce +import caliban.docker.build as b import caliban.platform.shell as ps -import caliban.util as u import caliban.util.fs as ufs -from caliban.history.types import Experiment, Job, JobSpec, JobStatus, Platform -from caliban.history.util import (create_experiments, generate_container_spec, - get_mem_engine, get_sql_engine, session_scope) t = Terminal() -DEV_CONTAINER_ROOT = "gcr.io/blueshift-playground/blueshift" -TF_VERSIONS = {"2.2.0", "1.12.3", "1.14.0", "1.15.0"} -DEFAULT_WORKDIR = "/usr/app" -CREDS_DIR = "/.creds" -CONDA_BIN = "/opt/conda/bin/conda" - -ImageId = NewType('ImageId', str) -ArgSeq = NewType('ArgSeq', List[str]) - - -class DockerError(Exception): - """Exception that passes info on a failed Docker command.""" - - def __init__(self, message, cmd, ret_code): - super().__init__(message) - self.message = message - self.cmd = cmd - self.ret_code = ret_code - - @property - def command(self): - return " ".join(self.cmd) - - -class NotebookInstall(Enum): - """Flag to decide what to do .""" - none = 'none' - lab = 'lab' - jupyter = 'jupyter' - - def __str__(self) -> str: - return self.value - - -class Shell(Enum): - """Add new shells here and below, in SHELL_DICT.""" - bash = 'bash' - zsh = 'zsh' - - def __str__(self) -> str: - return self.value - - -# Tuple to track the information required to install and execute some custom -# shell into a container. -ShellData = NamedTuple("ShellData", [("executable", str), - ("packages", List[str])]) - - -def apt_install(*packages: str) -> str: - """Returns a command that will install the supplied list of packages without - requiring confirmation or any user interaction. - """ - package_str = ' '.join(packages) - no_prompt = "DEBIAN_FRONTEND=noninteractive" - return f"{no_prompt} apt-get install --yes --no-install-recommends {package_str}" - - -def apt_command(commands: List[str]) -> List[str]: - """Pre-and-ap-pends the supplied commands with the appropriate in-container and - cleanup command for aptitude. - - """ - update = ["apt-get update"] - cleanup = ["apt-get clean", "rm -rf /var/lib/apt/lists/*"] - return update + commands + cleanup - - -# Dict linking a particular supported shell to the data required to run and -# install the shell inside a container. -# -# : Dict[Shell, ShellData] -SHELL_DICT = { - Shell.bash: ShellData("/bin/bash", []), - Shell.zsh: ShellData("/bin/zsh", ["zsh"]) -} - - -def default_shell() -> Shell: - """Returns the shell to load into the container. Defaults to Shell.bash, but if - the user's SHELL variable refers to a supported sub-shell, returns that - instead. - - """ - ret = Shell.bash - - if "zsh" in os.environ.get("SHELL"): - ret = Shell.zsh - - return ret - - -def adc_location(home_dir: Optional[str] = None) -> str: - """Returns the location for application default credentials, INSIDE the - container (so, hardcoded unix separators), given the supplied home directory. - - """ - if home_dir is None: - home_dir = Path.home() - - return "{}/.config/gcloud/application_default_credentials.json".format( - home_dir) - - -def container_home(): - """Returns the location of the home directory inside the generated - container. - - """ - return "/home/{}".format(u.current_user()) - - -def tf_base_image(job_mode: c.JobMode, tensorflow_version: str) -> str: - """Returns the base image to use, depending on whether or not we're using a - GPU. This is JUST for building our base images for Blueshift; not for - actually using in a job. - - List of available tags: https://hub.docker.com/r/tensorflow/tensorflow/tags - - """ - if tensorflow_version not in TF_VERSIONS: - raise Exception("""{} is not a valid tensorflow version. - Try one of: {}""".format(tensorflow_version, TF_VERSIONS)) - - gpu = "-gpu" if c.gpu(job_mode) else "" - return "tensorflow/tensorflow:{}{}-py3".format(tensorflow_version, gpu) - - -def base_image_suffix(job_mode: c.JobMode) -> str: - return "gpu" if c.gpu(job_mode) else "cpu" - - -def base_image_id(job_mode: c.JobMode) -> str: - """Returns the default base image for all caliban Dockerfiles.""" - base_suffix = base_image_suffix(job_mode) - return "{}:{}".format(DEV_CONTAINER_ROOT, base_suffix) - - -def extras_string(extras: List[str]) -> str: - """Returns the argument passed to `pip install` to install a project from its - setup.py and target a specific set of extras_require dependencies. - - Args: - extras: (potentially empty) list of extra_requires deps. - """ - ret = "." - if len(extras) > 0: - ret += "[{}]".format(','.join(extras)) - return ret - - -def base_extras(job_mode: c.JobMode, path: str, - extras: Optional[List[str]]) -> Optional[List[str]]: - """Returns None if the supplied path doesn't exist (it's assumed it points to a - setup.py file). - - If the path DOES exist, generates a list of extras to install. gpu or cpu are - always added to the beginning of the list, depending on the mode. - - """ - ret = None - - if os.path.exists(path): - base = extras or [] - extra = 'gpu' if c.gpu(job_mode) else 'cpu' - ret = base if extra in base else [extra] + base - - return ret - - -def _dependency_entries(workdir: str, - user_id: int, - user_group: int, - requirements_path: Optional[str] = None, - conda_env_path: Optional[str] = None, - setup_extras: Optional[List[str]] = None) -> str: - """Returns the Dockerfile entries required to install dependencies from either: - - - a requirements.txt file, path supplied by requirements_path - - a conda environment.yml file, path supplied by conda_env_path. - - a setup.py file, if some sequence of dependencies is supplied. - - An empty list for setup_extras means, run `pip install -c .` with no extras. - None for this argument means do nothing. If a list of strings is supplied, - they'll be treated as extras dependency sets. - """ - ret = "" - - if setup_extras is not None: - ret += f""" -COPY --chown={user_id}:{user_group} setup.py {workdir} -RUN /bin/bash -c "pip install --no-cache-dir {extras_string(setup_extras)}" -""" - - if conda_env_path is not None: - ret += f""" -COPY --chown={user_id}:{user_group} {conda_env_path} {workdir} -RUN /bin/bash -c "{CONDA_BIN} env update \ - --quiet --name caliban \ - --file {conda_env_path} && \ - {CONDA_BIN} clean -y -q --all" -""" - - if requirements_path is not None: - ret += f""" -COPY --chown={user_id}:{user_group} {requirements_path} {workdir} -RUN /bin/bash -c "pip install --no-cache-dir -r {requirements_path}" -""" - - return ret - - -def _package_entries(workdir: str, user_id: int, user_group: int, - package: u.Package) -> str: - """Returns the Dockerfile entries required to: - - - copy a directory of code into a docker container - - inject an entrypoint that executes a python module inside that directory. - - Python code runs as modules vs scripts so that we can enforce import hygiene - between files inside a project. - - """ - owner = "{}:{}".format(user_id, user_group) - - arg = package.main_module or package.script_path - - # This needs to use json so that quotes print as double quotes, not single - # quotes. - entrypoint_s = json.dumps(package.executable + [arg]) - - return """ -# Copy project code into the docker container. -COPY --chown={owner} {package_path} {workdir}/{package_path} - -# Declare an entrypoint that actually runs the container. -ENTRYPOINT {entrypoint_s} - """.format_map({ - "owner": owner, - "package_path": package.package_path, - "workdir": workdir, - "entrypoint_s": entrypoint_s - }) - - -def _service_account_entry(user_id: int, user_group: int, credentials_path: str, - docker_credentials_dir: str, - write_adc_placeholder: bool): - """Generates the Dockerfile entries required to transfer a set of Cloud service - account credentials into the Docker container. - - NOTE the write_adc_placeholder variable is here because the "ctpu" script - that we use to interact with TPUs has a bug in it, as of 1/21/2020, where the - script will fail if the application_default_credentials.json file isn't - present, EVEN THOUGH it properly uses the service account credentials - registered with gcloud instead of ADC creds. - - If a service account is present, we write a placeholder string to get past - this problem. This shouldn't matter for anyone else since adc isn't used if a - service account is present. - - """ - container_creds = "{}/credentials.json".format(docker_credentials_dir) - ret = """ -COPY --chown={user_id}:{user_group} {credentials_path} {container_creds} - -# Use the credentials file to activate gcloud, gsutil inside the container. -RUN gcloud auth activate-service-account --key-file={container_creds} && \ - git config --global credential.'https://source.developers.google.com'.helper gcloud.sh - -ENV GOOGLE_APPLICATION_CREDENTIALS={container_creds} -""".format_map({ - "user_id": user_id, - "user_group": user_group, - "credentials_path": credentials_path, - "container_creds": container_creds - }) - - if write_adc_placeholder: - ret += """ -RUN echo "placeholder" >> {} -""".format(adc_location(container_home())) - - return ret - - -def _adc_entry(user_id: int, user_group: int, adc_path: str): - """Returns the Dockerfile line required to transfer the - application_default_credentials.json file into the container's home - directory. - - """ - return """ -COPY --chown={user_id}:{user_group} {adc_path} {adc_loc} - """.format_map({ - "user_id": user_id, - "user_group": user_group, - "adc_path": adc_path, - "adc_loc": adc_location(container_home()) - }) - - -def _credentials_entries(user_id: int, - user_group: int, - adc_path: Optional[str], - credentials_path: Optional[str], - docker_credentials_dir: Optional[str] = None) -> str: - """Returns the Dockerfile entries necessary to copy a user's Cloud credentials - into the Docker container. - - - adc_path is the relative path inside the current directory to an - application_default_credentials.json file containing... well, you get it. - - credentials_path is the relative path inside the current directory to a - JSON credentials file. - - docker_credentials_dir is the relative path inside the docker container - where the JSON file will be copied on build. - - """ - if docker_credentials_dir is None: - docker_credentials_dir = CREDS_DIR - - ret = "" - if credentials_path is not None: - ret += _service_account_entry(user_id, - user_group, - credentials_path, - docker_credentials_dir, - write_adc_placeholder=adc_path is None) - - if adc_path is not None: - ret += _adc_entry(user_id, user_group, adc_path) - - return ret - - -def _notebook_entries(lab: bool = False, version: Optional[str] = None) -> str: - """Returns the Dockerfile entries necessary to install Jupyter{lab}. - - Optionally takes a version string. - - """ - version_suffix = "" - - if version is not None: - version_suffix = "=={}".format(version) - - library = "jupyterlab" if lab else "jupyter" - - return """ -RUN pip install {}{} -""".format(library, version_suffix) - - -def _custom_packages( - user_id: int, - user_group: int, - packages: Optional[List[str]] = None, - shell: Optional[Shell] = None, -) -> str: - """Returns the Dockerfile entries necessary to install custom dependencies for - the supplied shell and sequence of aptitude packages. - - """ - if packages is None: - packages = [] - - if shell is None: - shell = Shell.bash - - ret = "" - - to_install = sorted(packages + SHELL_DICT[shell].packages) - - if len(to_install) != 0: - commands = apt_command([apt_install(*to_install)]) - ret = """ -USER root - -RUN {commands} - -USER {user_id}:{user_group} -""".format_map({ - "commands": " && ".join(commands), - "user_id": user_id, - "user_group": user_group - }) - - return ret - - -def _copy_dir_entry(workdir: str, user_id: int, user_group: int, - dirname: str) -> str: - """Returns the Dockerfile entry necessary to copy a single extra subdirectory - from the current directory into a docker container during build. - - """ - owner = "{}:{}".format(user_id, user_group) - return """# Copy {dirname} into the Docker container. -COPY --chown={owner} {dirname} {workdir}/{dirname} -""".format_map({ - "owner": owner, - "workdir": workdir, - "dirname": dirname - }) - - -def _extra_dir_entries(workdir: str, user_id: int, user_group: int, - extra_dirs: List[str]) -> str: - """Returns the Dockerfile entries necessary to copy all directories in the - extra_dirs list into a docker container during build. - - """ - ret = "" - for d in extra_dirs: - ret += "\n{}".format(_copy_dir_entry(workdir, user_id, user_group, d)) - return ret - - -def _dockerfile_template( - job_mode: c.JobMode, - workdir: Optional[str] = None, - base_image_fn: Optional[Callable[[c.JobMode], str]] = None, - package: Optional[Union[List, u.Package]] = None, - requirements_path: Optional[str] = None, - conda_env_path: Optional[str] = None, - setup_extras: Optional[List[str]] = None, - adc_path: Optional[str] = None, - credentials_path: Optional[str] = None, - jupyter_version: Optional[str] = None, - inject_notebook: NotebookInstall = NotebookInstall.none, - shell: Optional[Shell] = None, - extra_dirs: Optional[List[str]] = None, - caliban_config: Optional[Dict[str, Any]] = None) -> str: - """Returns a Dockerfile that builds on a local CPU or GPU base image (depending - on the value of job_mode) to create a container that: - - - installs any dependency specified in a requirements.txt file living at - requirements_path, a conda environment at conda_env_path, or any - dependencies in a setup.py file, including extra dependencies, if - setup_extras isn't None - - injects gcloud credentials into the container, so Cloud interaction works - just like it does locally - - potentially installs a custom shell, or jupyterlab for notebook support - - copies all source needed by the main module specified by package, and - potentially injects an entrypoint that, on run, will run that main module - - Most functions that call _dockerfile_template pass along any kwargs that they - receive. It should be enough to add kwargs here, then rely on that mechanism - to pass them along, vs adding kwargs all the way down the call chain. - - Supply a custom base_image_fn (function from job_mode -> image ID) to inject - more complex Docker commands into the Caliban environments by, for example, - building your own image on top of the TF base images, then using that. - - """ - uid = os.getuid() - gid = os.getgid() - username = u.current_user() - - if isinstance(package, list): - package = u.Package(*package) - - if workdir is None: - workdir = DEFAULT_WORKDIR - - if base_image_fn is None: - base_image_fn = base_image_id - - base_image = base_image_fn(job_mode) - - dockerfile = """ -FROM {base_image} - -# Create the same group we're using on the host machine. -RUN [ $(getent group {gid}) ] || groupadd --gid {gid} {gid} - -# Create the user by name. --no-log-init guards against a crash with large user -# IDs. -RUN useradd --no-log-init --no-create-home -u {uid} -g {gid} --shell /bin/bash {username} - -# The directory is created by root. This sets permissions so that any user can -# access the folder. -RUN mkdir -m 777 {workdir} {creds_dir} {c_home} - -ENV HOME={c_home} - -WORKDIR {workdir} - -USER {uid}:{gid} -""".format_map({ - "base_image": base_image, - "username": username, - "uid": uid, - "gid": gid, - "workdir": workdir, - "c_home": container_home(), - "creds_dir": CREDS_DIR - }) - dockerfile += _credentials_entries(uid, - gid, - adc_path=adc_path, - credentials_path=credentials_path) - - dockerfile += _dependency_entries(workdir, - uid, - gid, - requirements_path=requirements_path, - conda_env_path=conda_env_path, - setup_extras=setup_extras) - - if inject_notebook.value != 'none': - install_lab = inject_notebook == NotebookInstall.lab - dockerfile += _notebook_entries(lab=install_lab, version=jupyter_version) - - if extra_dirs is not None: - dockerfile += _extra_dir_entries(workdir, uid, gid, extra_dirs) - - dockerfile += _custom_packages(uid, - gid, - packages=c.apt_packages( - caliban_config, job_mode), - shell=shell) - - if package is not None: - # The actual entrypoint and final copied code. - dockerfile += _package_entries(workdir, uid, gid, package) - - return dockerfile - - -def docker_image_id(output: str) -> ImageId: - """Accepts a string containing the output of a successful `docker build` - command and parses the Docker image ID from the stream. - - NOTE this is probably quite brittle! I can imagine this breaking quite easily - on a Docker upgrade. - - """ - return ImageId(output.splitlines()[-1].split()[-1]) - - -def build_image(job_mode: c.JobMode, - build_path: str, - credentials_path: Optional[str] = None, - adc_path: Optional[str] = None, - no_cache: bool = False, - **kwargs) -> str: - """Builds a Docker image by generating a Dockerfile and passing it to `docker - build` via stdin. All output from the `docker build` process prints to - stdout. - - Returns the image ID of the new docker container; if the command fails, - throws on error with information about the command and any issues that caused - the problem. - - """ - with u.TempCopy(credentials_path, - tmp_name=".caliban_default_creds.json") as creds: - with u.TempCopy(adc_path, tmp_name=".caliban_adc_creds.json") as adc: - cache_args = ["--no-cache"] if no_cache else [] - cmd = ["docker", "build"] + cache_args + ["--rm", "-f-", build_path] - - dockerfile = _dockerfile_template(job_mode, - credentials_path=creds, - adc_path=adc, - **kwargs) - - joined_cmd = " ".join(cmd) - logging.info("Running command: {}".format(joined_cmd)) - - try: - output, ret_code = ufs.capture_stdout(cmd, input_str=dockerfile) - if ret_code == 0: - return docker_image_id(output) - else: - error_msg = "Docker failed with error code {}.".format(ret_code) - raise DockerError(error_msg, cmd, ret_code) - - except subprocess.CalledProcessError as e: - logging.error(e.output) - logging.error(e.stderr) - - -def _image_tag_for_project(project_id: str, image_id: str) -> str: - """Generate the GCR Docker image tag for the supplied pair of project_id and - image_id. - - This function properly handles "domain scoped projects", where the project ID - contains a domain name and project ID separated by : - https://cloud.google.com/container-registry/docs/overview#domain-scoped_projects. - - """ - project_s = project_id.replace(":", "/") - return "gcr.io/{}/{}:latest".format(project_s, image_id) - - -def push_uuid_tag(project_id: str, image_id: str) -> str: - """Takes a base image and tags it for upload, then pushes it to a remote Google - Container Registry. - - Returns the tag on a successful push. - - TODO should this just check first before attempting to push if the image - exists? Immutable names means that if the tag is up there, we're done. - Potentially use docker-py for this. - - """ - image_tag = _image_tag_for_project(project_id, image_id) - subprocess.run(["docker", "tag", image_id, image_tag], check=True) - subprocess.run(["docker", "push", image_tag], check=True) - return image_tag - - -def _run_cmd(job_mode: c.JobMode, - run_args: Optional[List[str]] = None) -> List[str]: - """Returns the sequence of commands for the subprocess run functions required - to execute `docker run`. in CPU or GPU mode, depending on the value of - job_mode. - - Keyword args: - - run_args: list of args to pass to docker run. - - """ - if run_args is None: - run_args = [] - - runtime = ["--runtime", "nvidia"] if c.gpu(job_mode) else [] - return ["docker", "run"] + runtime + ["--ipc", "host"] + run_args - - -def _home_mount_cmds(enable_home_mount: bool) -> List[str]: - """Returns the argument needed by Docker to mount a user's local home directory - into the home directory location inside their container. - - If enable_home_mount is false returns an empty list. - - """ - ret = [] - if enable_home_mount: - ret = ["-v", "{}:{}".format(Path.home(), container_home())] - return ret - - -def _interactive_opts(workdir: str) -> List[str]: - """Returns the basic arguments we want to run a docker process locally. - - """ - return [ - "-w", workdir, \ - "-u", "{}:{}".format(os.getuid(), os.getgid()), \ - "-v", "{}:{}".format(os.getcwd(), workdir) \ - ] - - -def log_job_spec_instance(job_spec: JobSpec, i: int) -> JobSpec: - """Prints logging as a side effect for the supplied sequence of job specs - generated from an experiment definition; returns the input job spec. - - """ - args = ce.experiment_to_args(job_spec.experiment.kwargs, - job_spec.experiment.args) - logging.info("") - logging.info("Job {} - Experiment args: {}".format(i, t.yellow(str(args)))) - return job_spec - - -def logged_job_specs(job_specs: Iterable[JobSpec]) -> Iterable[JobSpec]: - """Generates an iterable of job specs that should be passed to `docker run` to - execute the experiments defined by the supplied iterable. - - """ - for i, s in enumerate(job_specs, 1): - yield log_job_spec_instance(s, i) - - -def execute_dry_run(job_specs: Iterable[JobSpec]) -> None: - """Expands the supplied sequence of experiments into sequences of args and logs - the jobs that WOULD have been executed, had the dry run flag not been - applied. - - """ - list(logged_job_specs(job_specs)) - - logging.info('') - logging.info( - t.yellow("To build your image and execute these jobs, \ -run your command again without {}.".format(c.DRY_RUN_FLAG))) - logging.info('') - return None - - -def local_callback(idx: int, job: Job) -> None: - """Provides logging feedback for jobs run locally. If the return code is 0, - logs success; else, logs the failure as an error and logs the script args - that provided the failure. - - """ - if job.status == JobStatus.SUCCEEDED: - logging.info(t.green(f'Job {idx} succeeded!')) - else: - logging.error( - t.red(f'Job {idx} failed with return code {job.details["ret_code"]}.')) - args = ce.experiment_to_args(job.spec.experiment.kwargs, - job.spec.experiment.args) - logging.error(t.red(f'Failing args for job {idx}: {args}')) - - -def window_size_env_cmds(): - """Returns a sequence of `docker run` arguments that will internally configure - the terminal columns and lines, so that progress bars and other terminal - interactions will work properly. - - These aren't required for interactive Docker commands like those triggered by - `caliban shell`. - - """ - ret = [] - cols, lines = _screen_shape_wrapper()(0) - if cols: - ret += ["-e", f"COLUMNS={cols}"] - if lines: - ret += ["-e", f"LINES={lines}"] - return ret - - -# ---------------------------------------------------------------------------- -def _create_job_spec_dict( - experiment: Experiment, - job_mode: c.JobMode, - image_id: str, - run_args: Optional[List[str]] = None, -) -> Dict[str, Any]: - '''creates a job spec dictionary for a local job''' - - # Without the unbuffered environment variable, stderr and stdout won't be - # emitted in the proper order from inside the container. - terminal_cmds = ["-e" "PYTHONUNBUFFERED=1"] + window_size_env_cmds() - - base_cmd = _run_cmd(job_mode, run_args) + terminal_cmds + [image_id] - command = base_cmd + ce.experiment_to_args(experiment.kwargs, experiment.args) - return {'command': command, 'container': image_id} - - -# ---------------------------------------------------------------------------- -def execute_jobs( - job_specs: Iterable[JobSpec], - dry_run: bool = False, -): - '''executes a sequence of jobs based on job specs - - Arg: - job_specs: specifications for jobs to be executed - dry_run: if True, only print what would be done - ''' - - with u.tqdm_logging() as orig_stream: - pbar = tqdm.tqdm(logged_job_specs(job_specs), - file=orig_stream, - total=len(job_specs), - ascii=True, - unit="experiment", - desc="Executing") - for idx, job_spec in enumerate(pbar, 1): - command = job_spec.spec['command'] - logging.info(f'Running command: {" ".join(command)}') - if not dry_run: - _, ret_code = ufs.capture_stdout(command, "", u.TqdmFile(sys.stderr)) - else: - ret_code = 0 - j = Job(spec=job_spec, - container=job_spec.spec['container'], - details={'ret_code': ret_code}, - status=JobStatus.SUCCEEDED if ret_code == 0 else JobStatus.FAILED) - local_callback(idx=idx, job=j) - - if dry_run: - logging.info( - t.yellow(f'\nTo build your image and execute these jobs, ' - f'run your command again without {c.DRY_RUN_FLAG}\n')) - - return None - - -def run_experiments(job_mode: c.JobMode, - run_args: Optional[List[str]] = None, - script_args: Optional[List[str]] = None, - image_id: Optional[str] = None, - dry_run: bool = False, - experiment_config: Optional[ce.ExpConf] = None, - xgroup: Optional[str] = None, - **build_image_kwargs) -> None: - """Builds an image using the supplied **build_image_kwargs and calls `docker - run` on the resulting image using sensible defaults. - - Keyword args: - - - job_mode: c.JobMode. - - - run_args: extra arguments to supply to `docker run` after our defaults. - - script_args: extra arguments to supply to the entrypoint. (You can - - override the default container entrypoint by supplying a new one inside - run_args.) - - image_id: ID of the image to run. Supplying this will skip an image build. - - experiment_config: dict of string to list, boolean, string or int. Any - lists will trigger a cartesian product out with the rest of the config. A - job will be executed for every combination of parameters in the experiment - config. - - dry_run: if True, no actual jobs will be executed and docker won't - actually build; logging side effects will show the user what will happen - without dry_run=True. - - any extra kwargs supplied are passed through to build_image. - """ - if run_args is None: - run_args = [] - - if script_args is None: - script_args = [] - - if experiment_config is None: - experiment_config = {} - - docker_args = {k: v for k, v in build_image_kwargs.items()} - docker_args['job_mode'] = job_mode - - engine = get_mem_engine() if dry_run else get_sql_engine() - - with session_scope(engine) as session: - container_spec = generate_container_spec(session, docker_args, image_id) - - if image_id is None: - if dry_run: - logging.info("Dry run - skipping actual 'docker build'.") - image_id = 'dry_run_tag' - else: - image_id = build_image(**docker_args) - - experiments = create_experiments( - session=session, - container_spec=container_spec, - script_args=script_args, - experiment_config=experiment_config, - xgroup=xgroup, - ) - - job_specs = [ - JobSpec.get_or_create( - experiment=x, - spec=_create_job_spec_dict( - experiment=x, - job_mode=job_mode, - run_args=run_args, - image_id=image_id, - ), - platform=Platform.LOCAL, - ) for x in experiments - ] - - try: - execute_jobs(job_specs=job_specs, dry_run=dry_run) - except Exception as e: - logging.error(f'exception: {e}') - session.commit() # commit here, otherwise will be rolled back - - -def run(job_mode: c.JobMode, - run_args: Optional[List[str]] = None, - script_args: Optional[List[str]] = None, - image_id: Optional[str] = None, - **build_image_kwargs) -> None: - """Builds an image using the supplied **build_image_kwargs and calls `docker - run` on the resulting image using sensible defaults. - Keyword args: - - job_mode: c.JobMode. - - run_args: extra arguments to supply to `docker run` after our defaults. - - script_args: extra arguments to supply to the entrypoint. (You can - - override the default container entrypoint by supplying a new one inside - run_args.) - - image_id: ID of the image to run. Supplying this will skip an image build. - any extra kwargs supplied are passed through to build_image. - """ - if run_args is None: - run_args = [] - - if script_args is None: - script_args = [] - - if image_id is None: - image_id = build_image(job_mode, **build_image_kwargs) - - base_cmd = _run_cmd(job_mode, run_args) - - command = base_cmd + [image_id] + script_args - - logging.info("Running command: {}".format(' '.join(command))) - subprocess.call(command) - return None - def run_notebook(job_mode: c.JobMode, port: Optional[int] = None, @@ -969,7 +55,7 @@ def run_notebook(job_mode: c.JobMode, """ if port is None: - port = u.next_free_port(8888) + port = ufs.next_free_port(8888) if lab is None: lab = False @@ -977,7 +63,7 @@ def run_notebook(job_mode: c.JobMode, if run_args is None: run_args = [] - inject_arg = NotebookInstall.lab if lab else NotebookInstall.jupyter + inject_arg = b.NotebookInstall.lab if lab else b.NotebookInstall.jupyter jupyter_cmd = "lab" if lab else "notebook" jupyter_args = [ "-m", "jupyter", jupyter_cmd, \ diff --git a/caliban/platform/run.py b/caliban/platform/run.py index 9304ee4..c55beb9 100644 --- a/caliban/platform/run.py +++ b/caliban/platform/run.py @@ -20,14 +20,9 @@ from __future__ import absolute_import, division, print_function -import json -import os import subprocess import sys -from enum import Enum -from pathlib import Path -from typing import (Any, Callable, Dict, Iterable, List, NamedTuple, NewType, - Optional, Union) +from typing import Any, Dict, Iterable, List, Optional import tqdm from absl import logging @@ -36,6 +31,7 @@ import caliban.config as c import caliban.config.experiment as ce +import caliban.docker.build as b import caliban.util as u import caliban.util.fs as ufs from caliban.history.types import Experiment, Job, JobSpec, JobStatus, Platform @@ -44,618 +40,6 @@ t = Terminal() -DEV_CONTAINER_ROOT = "gcr.io/blueshift-playground/blueshift" -TF_VERSIONS = {"2.2.0", "1.12.3", "1.14.0", "1.15.0"} -DEFAULT_WORKDIR = "/usr/app" -CREDS_DIR = "/.creds" -CONDA_BIN = "/opt/conda/bin/conda" - -ImageId = NewType('ImageId', str) -ArgSeq = NewType('ArgSeq', List[str]) - - -class DockerError(Exception): - """Exception that passes info on a failed Docker command.""" - - def __init__(self, message, cmd, ret_code): - super().__init__(message) - self.message = message - self.cmd = cmd - self.ret_code = ret_code - - @property - def command(self): - return " ".join(self.cmd) - - -class NotebookInstall(Enum): - """Flag to decide what to do .""" - none = 'none' - lab = 'lab' - jupyter = 'jupyter' - - def __str__(self) -> str: - return self.value - - -class Shell(Enum): - """Add new shells here and below, in SHELL_DICT.""" - bash = 'bash' - zsh = 'zsh' - - def __str__(self) -> str: - return self.value - - -# Tuple to track the information required to install and execute some custom -# shell into a container. -ShellData = NamedTuple("ShellData", [("executable", str), - ("packages", List[str])]) - - -def apt_install(*packages: str) -> str: - """Returns a command that will install the supplied list of packages without - requiring confirmation or any user interaction. - """ - package_str = ' '.join(packages) - no_prompt = "DEBIAN_FRONTEND=noninteractive" - return f"{no_prompt} apt-get install --yes --no-install-recommends {package_str}" - - -def apt_command(commands: List[str]) -> List[str]: - """Pre-and-ap-pends the supplied commands with the appropriate in-container and - cleanup command for aptitude. - - """ - update = ["apt-get update"] - cleanup = ["apt-get clean", "rm -rf /var/lib/apt/lists/*"] - return update + commands + cleanup - - -# Dict linking a particular supported shell to the data required to run and -# install the shell inside a container. -# -# : Dict[Shell, ShellData] -SHELL_DICT = { - Shell.bash: ShellData("/bin/bash", []), - Shell.zsh: ShellData("/bin/zsh", ["zsh"]) -} - - -def default_shell() -> Shell: - """Returns the shell to load into the container. Defaults to Shell.bash, but if - the user's SHELL variable refers to a supported sub-shell, returns that - instead. - - """ - ret = Shell.bash - - if "zsh" in os.environ.get("SHELL"): - ret = Shell.zsh - - return ret - - -def adc_location(home_dir: Optional[str] = None) -> str: - """Returns the location for application default credentials, INSIDE the - container (so, hardcoded unix separators), given the supplied home directory. - - """ - if home_dir is None: - home_dir = Path.home() - - return "{}/.config/gcloud/application_default_credentials.json".format( - home_dir) - - -def container_home(): - """Returns the location of the home directory inside the generated - container. - - """ - return "/home/{}".format(u.current_user()) - - -def tf_base_image(job_mode: c.JobMode, tensorflow_version: str) -> str: - """Returns the base image to use, depending on whether or not we're using a - GPU. This is JUST for building our base images for Blueshift; not for - actually using in a job. - - List of available tags: https://hub.docker.com/r/tensorflow/tensorflow/tags - - """ - if tensorflow_version not in TF_VERSIONS: - raise Exception("""{} is not a valid tensorflow version. - Try one of: {}""".format(tensorflow_version, TF_VERSIONS)) - - gpu = "-gpu" if c.gpu(job_mode) else "" - return "tensorflow/tensorflow:{}{}-py3".format(tensorflow_version, gpu) - - -def base_image_suffix(job_mode: c.JobMode) -> str: - return "gpu" if c.gpu(job_mode) else "cpu" - - -def base_image_id(job_mode: c.JobMode) -> str: - """Returns the default base image for all caliban Dockerfiles.""" - base_suffix = base_image_suffix(job_mode) - return "{}:{}".format(DEV_CONTAINER_ROOT, base_suffix) - - -def extras_string(extras: List[str]) -> str: - """Returns the argument passed to `pip install` to install a project from its - setup.py and target a specific set of extras_require dependencies. - - Args: - extras: (potentially empty) list of extra_requires deps. - """ - ret = "." - if len(extras) > 0: - ret += "[{}]".format(','.join(extras)) - return ret - - -def base_extras(job_mode: c.JobMode, path: str, - extras: Optional[List[str]]) -> Optional[List[str]]: - """Returns None if the supplied path doesn't exist (it's assumed it points to a - setup.py file). - - If the path DOES exist, generates a list of extras to install. gpu or cpu are - always added to the beginning of the list, depending on the mode. - - """ - ret = None - - if os.path.exists(path): - base = extras or [] - extra = 'gpu' if c.gpu(job_mode) else 'cpu' - ret = base if extra in base else [extra] + base - - return ret - - -def _dependency_entries(workdir: str, - user_id: int, - user_group: int, - requirements_path: Optional[str] = None, - conda_env_path: Optional[str] = None, - setup_extras: Optional[List[str]] = None) -> str: - """Returns the Dockerfile entries required to install dependencies from either: - - - a requirements.txt file, path supplied by requirements_path - - a conda environment.yml file, path supplied by conda_env_path. - - a setup.py file, if some sequence of dependencies is supplied. - - An empty list for setup_extras means, run `pip install -c .` with no extras. - None for this argument means do nothing. If a list of strings is supplied, - they'll be treated as extras dependency sets. - """ - ret = "" - - if setup_extras is not None: - ret += f""" -COPY --chown={user_id}:{user_group} setup.py {workdir} -RUN /bin/bash -c "pip install --no-cache-dir {extras_string(setup_extras)}" -""" - - if conda_env_path is not None: - ret += f""" -COPY --chown={user_id}:{user_group} {conda_env_path} {workdir} -RUN /bin/bash -c "{CONDA_BIN} env update \ - --quiet --name caliban \ - --file {conda_env_path} && \ - {CONDA_BIN} clean -y -q --all" -""" - - if requirements_path is not None: - ret += f""" -COPY --chown={user_id}:{user_group} {requirements_path} {workdir} -RUN /bin/bash -c "pip install --no-cache-dir -r {requirements_path}" -""" - - return ret - - -def _package_entries(workdir: str, user_id: int, user_group: int, - package: u.Package) -> str: - """Returns the Dockerfile entries required to: - - - copy a directory of code into a docker container - - inject an entrypoint that executes a python module inside that directory. - - Python code runs as modules vs scripts so that we can enforce import hygiene - between files inside a project. - - """ - owner = "{}:{}".format(user_id, user_group) - - arg = package.main_module or package.script_path - - # This needs to use json so that quotes print as double quotes, not single - # quotes. - entrypoint_s = json.dumps(package.executable + [arg]) - - return """ -# Copy project code into the docker container. -COPY --chown={owner} {package_path} {workdir}/{package_path} - -# Declare an entrypoint that actually runs the container. -ENTRYPOINT {entrypoint_s} - """.format_map({ - "owner": owner, - "package_path": package.package_path, - "workdir": workdir, - "entrypoint_s": entrypoint_s - }) - - -def _service_account_entry(user_id: int, user_group: int, credentials_path: str, - docker_credentials_dir: str, - write_adc_placeholder: bool): - """Generates the Dockerfile entries required to transfer a set of Cloud service - account credentials into the Docker container. - - NOTE the write_adc_placeholder variable is here because the "ctpu" script - that we use to interact with TPUs has a bug in it, as of 1/21/2020, where the - script will fail if the application_default_credentials.json file isn't - present, EVEN THOUGH it properly uses the service account credentials - registered with gcloud instead of ADC creds. - - If a service account is present, we write a placeholder string to get past - this problem. This shouldn't matter for anyone else since adc isn't used if a - service account is present. - - """ - container_creds = "{}/credentials.json".format(docker_credentials_dir) - ret = """ -COPY --chown={user_id}:{user_group} {credentials_path} {container_creds} - -# Use the credentials file to activate gcloud, gsutil inside the container. -RUN gcloud auth activate-service-account --key-file={container_creds} && \ - git config --global credential.'https://source.developers.google.com'.helper gcloud.sh - -ENV GOOGLE_APPLICATION_CREDENTIALS={container_creds} -""".format_map({ - "user_id": user_id, - "user_group": user_group, - "credentials_path": credentials_path, - "container_creds": container_creds - }) - - if write_adc_placeholder: - ret += """ -RUN echo "placeholder" >> {} -""".format(adc_location(container_home())) - - return ret - - -def _adc_entry(user_id: int, user_group: int, adc_path: str): - """Returns the Dockerfile line required to transfer the - application_default_credentials.json file into the container's home - directory. - - """ - return """ -COPY --chown={user_id}:{user_group} {adc_path} {adc_loc} - """.format_map({ - "user_id": user_id, - "user_group": user_group, - "adc_path": adc_path, - "adc_loc": adc_location(container_home()) - }) - - -def _credentials_entries(user_id: int, - user_group: int, - adc_path: Optional[str], - credentials_path: Optional[str], - docker_credentials_dir: Optional[str] = None) -> str: - """Returns the Dockerfile entries necessary to copy a user's Cloud credentials - into the Docker container. - - - adc_path is the relative path inside the current directory to an - application_default_credentials.json file containing... well, you get it. - - credentials_path is the relative path inside the current directory to a - JSON credentials file. - - docker_credentials_dir is the relative path inside the docker container - where the JSON file will be copied on build. - - """ - if docker_credentials_dir is None: - docker_credentials_dir = CREDS_DIR - - ret = "" - if credentials_path is not None: - ret += _service_account_entry(user_id, - user_group, - credentials_path, - docker_credentials_dir, - write_adc_placeholder=adc_path is None) - - if adc_path is not None: - ret += _adc_entry(user_id, user_group, adc_path) - - return ret - - -def _notebook_entries(lab: bool = False, version: Optional[str] = None) -> str: - """Returns the Dockerfile entries necessary to install Jupyter{lab}. - - Optionally takes a version string. - - """ - version_suffix = "" - - if version is not None: - version_suffix = "=={}".format(version) - - library = "jupyterlab" if lab else "jupyter" - - return """ -RUN pip install {}{} -""".format(library, version_suffix) - - -def _custom_packages( - user_id: int, - user_group: int, - packages: Optional[List[str]] = None, - shell: Optional[Shell] = None, -) -> str: - """Returns the Dockerfile entries necessary to install custom dependencies for - the supplied shell and sequence of aptitude packages. - - """ - if packages is None: - packages = [] - - if shell is None: - shell = Shell.bash - - ret = "" - - to_install = sorted(packages + SHELL_DICT[shell].packages) - - if len(to_install) != 0: - commands = apt_command([apt_install(*to_install)]) - ret = """ -USER root - -RUN {commands} - -USER {user_id}:{user_group} -""".format_map({ - "commands": " && ".join(commands), - "user_id": user_id, - "user_group": user_group - }) - - return ret - - -def _copy_dir_entry(workdir: str, user_id: int, user_group: int, - dirname: str) -> str: - """Returns the Dockerfile entry necessary to copy a single extra subdirectory - from the current directory into a docker container during build. - - """ - owner = "{}:{}".format(user_id, user_group) - return """# Copy {dirname} into the Docker container. -COPY --chown={owner} {dirname} {workdir}/{dirname} -""".format_map({ - "owner": owner, - "workdir": workdir, - "dirname": dirname - }) - - -def _extra_dir_entries(workdir: str, user_id: int, user_group: int, - extra_dirs: List[str]) -> str: - """Returns the Dockerfile entries necessary to copy all directories in the - extra_dirs list into a docker container during build. - - """ - ret = "" - for d in extra_dirs: - ret += "\n{}".format(_copy_dir_entry(workdir, user_id, user_group, d)) - return ret - - -def _dockerfile_template( - job_mode: c.JobMode, - workdir: Optional[str] = None, - base_image_fn: Optional[Callable[[c.JobMode], str]] = None, - package: Optional[Union[List, u.Package]] = None, - requirements_path: Optional[str] = None, - conda_env_path: Optional[str] = None, - setup_extras: Optional[List[str]] = None, - adc_path: Optional[str] = None, - credentials_path: Optional[str] = None, - jupyter_version: Optional[str] = None, - inject_notebook: NotebookInstall = NotebookInstall.none, - shell: Optional[Shell] = None, - extra_dirs: Optional[List[str]] = None, - caliban_config: Optional[Dict[str, Any]] = None) -> str: - """Returns a Dockerfile that builds on a local CPU or GPU base image (depending - on the value of job_mode) to create a container that: - - - installs any dependency specified in a requirements.txt file living at - requirements_path, a conda environment at conda_env_path, or any - dependencies in a setup.py file, including extra dependencies, if - setup_extras isn't None - - injects gcloud credentials into the container, so Cloud interaction works - just like it does locally - - potentially installs a custom shell, or jupyterlab for notebook support - - copies all source needed by the main module specified by package, and - potentially injects an entrypoint that, on run, will run that main module - - Most functions that call _dockerfile_template pass along any kwargs that they - receive. It should be enough to add kwargs here, then rely on that mechanism - to pass them along, vs adding kwargs all the way down the call chain. - - Supply a custom base_image_fn (function from job_mode -> image ID) to inject - more complex Docker commands into the Caliban environments by, for example, - building your own image on top of the TF base images, then using that. - - """ - uid = os.getuid() - gid = os.getgid() - username = u.current_user() - - if isinstance(package, list): - package = u.Package(*package) - - if workdir is None: - workdir = DEFAULT_WORKDIR - - if base_image_fn is None: - base_image_fn = base_image_id - - base_image = base_image_fn(job_mode) - - dockerfile = """ -FROM {base_image} - -# Create the same group we're using on the host machine. -RUN [ $(getent group {gid}) ] || groupadd --gid {gid} {gid} - -# Create the user by name. --no-log-init guards against a crash with large user -# IDs. -RUN useradd --no-log-init --no-create-home -u {uid} -g {gid} --shell /bin/bash {username} - -# The directory is created by root. This sets permissions so that any user can -# access the folder. -RUN mkdir -m 777 {workdir} {creds_dir} {c_home} - -ENV HOME={c_home} - -WORKDIR {workdir} - -USER {uid}:{gid} -""".format_map({ - "base_image": base_image, - "username": username, - "uid": uid, - "gid": gid, - "workdir": workdir, - "c_home": container_home(), - "creds_dir": CREDS_DIR - }) - dockerfile += _credentials_entries(uid, - gid, - adc_path=adc_path, - credentials_path=credentials_path) - - dockerfile += _dependency_entries(workdir, - uid, - gid, - requirements_path=requirements_path, - conda_env_path=conda_env_path, - setup_extras=setup_extras) - - if inject_notebook.value != 'none': - install_lab = inject_notebook == NotebookInstall.lab - dockerfile += _notebook_entries(lab=install_lab, version=jupyter_version) - - if extra_dirs is not None: - dockerfile += _extra_dir_entries(workdir, uid, gid, extra_dirs) - - dockerfile += _custom_packages(uid, - gid, - packages=c.apt_packages( - caliban_config, job_mode), - shell=shell) - - if package is not None: - # The actual entrypoint and final copied code. - dockerfile += _package_entries(workdir, uid, gid, package) - - return dockerfile - - -def docker_image_id(output: str) -> ImageId: - """Accepts a string containing the output of a successful `docker build` - command and parses the Docker image ID from the stream. - - NOTE this is probably quite brittle! I can imagine this breaking quite easily - on a Docker upgrade. - - """ - return ImageId(output.splitlines()[-1].split()[-1]) - - -def build_image(job_mode: c.JobMode, - build_path: str, - credentials_path: Optional[str] = None, - adc_path: Optional[str] = None, - no_cache: bool = False, - **kwargs) -> str: - """Builds a Docker image by generating a Dockerfile and passing it to `docker - build` via stdin. All output from the `docker build` process prints to - stdout. - - Returns the image ID of the new docker container; if the command fails, - throws on error with information about the command and any issues that caused - the problem. - - """ - with u.TempCopy(credentials_path, - tmp_name=".caliban_default_creds.json") as creds: - with u.TempCopy(adc_path, tmp_name=".caliban_adc_creds.json") as adc: - cache_args = ["--no-cache"] if no_cache else [] - cmd = ["docker", "build"] + cache_args + ["--rm", "-f-", build_path] - - dockerfile = _dockerfile_template(job_mode, - credentials_path=creds, - adc_path=adc, - **kwargs) - - joined_cmd = " ".join(cmd) - logging.info("Running command: {}".format(joined_cmd)) - - try: - output, ret_code = ufs.capture_stdout(cmd, input_str=dockerfile) - if ret_code == 0: - return docker_image_id(output) - else: - error_msg = "Docker failed with error code {}.".format(ret_code) - raise DockerError(error_msg, cmd, ret_code) - - except subprocess.CalledProcessError as e: - logging.error(e.output) - logging.error(e.stderr) - - -def _image_tag_for_project(project_id: str, image_id: str) -> str: - """Generate the GCR Docker image tag for the supplied pair of project_id and - image_id. - - This function properly handles "domain scoped projects", where the project ID - contains a domain name and project ID separated by : - https://cloud.google.com/container-registry/docs/overview#domain-scoped_projects. - - """ - project_s = project_id.replace(":", "/") - return "gcr.io/{}/{}:latest".format(project_s, image_id) - - -def push_uuid_tag(project_id: str, image_id: str) -> str: - """Takes a base image and tags it for upload, then pushes it to a remote Google - Container Registry. - - Returns the tag on a successful push. - - TODO should this just check first before attempting to push if the image - exists? Immutable names means that if the tag is up there, we're done. - Potentially use docker-py for this. - - """ - image_tag = _image_tag_for_project(project_id, image_id) - subprocess.run(["docker", "tag", image_id, image_tag], check=True) - subprocess.run(["docker", "push", image_tag], check=True) - return image_tag - def _run_cmd(job_mode: c.JobMode, run_args: Optional[List[str]] = None) -> List[str]: @@ -674,30 +58,6 @@ def _run_cmd(job_mode: c.JobMode, return ["docker", "run"] + runtime + ["--ipc", "host"] + run_args -def _home_mount_cmds(enable_home_mount: bool) -> List[str]: - """Returns the argument needed by Docker to mount a user's local home directory - into the home directory location inside their container. - - If enable_home_mount is false returns an empty list. - - """ - ret = [] - if enable_home_mount: - ret = ["-v", "{}:{}".format(Path.home(), container_home())] - return ret - - -def _interactive_opts(workdir: str) -> List[str]: - """Returns the basic arguments we want to run a docker process locally. - - """ - return [ - "-w", workdir, \ - "-u", "{}:{}".format(os.getuid(), os.getgid()), \ - "-v", "{}:{}".format(os.getcwd(), workdir) \ - ] - - def log_job_spec_instance(job_spec: JobSpec, i: int) -> JobSpec: """Prints logging as a side effect for the supplied sequence of job specs generated from an experiment definition; returns the input job spec. @@ -879,7 +239,7 @@ def run_experiments(job_mode: c.JobMode, logging.info("Dry run - skipping actual 'docker build'.") image_id = 'dry_run_tag' else: - image_id = build_image(**docker_args) + image_id = b.build_image(**docker_args) experiments = create_experiments( session=session, @@ -932,7 +292,7 @@ def run(job_mode: c.JobMode, script_args = [] if image_id is None: - image_id = build_image(job_mode, **build_image_kwargs) + image_id = b.build_image(job_mode, **build_image_kwargs) base_cmd = _run_cmd(job_mode, run_args) diff --git a/caliban/platform/shell.py b/caliban/platform/shell.py index e15ff83..26e8856 100644 --- a/caliban/platform/shell.py +++ b/caliban/platform/shell.py @@ -18,660 +18,13 @@ """ -from __future__ import absolute_import, division, print_function - -import json import os -import subprocess -import sys -from enum import Enum from pathlib import Path -from typing import (Any, Callable, Dict, Iterable, List, NamedTuple, NewType, - Optional, Union) - -import tqdm -from absl import logging -from blessings import Terminal -from tqdm.utils import _screen_shape_wrapper +from typing import List, Optional import caliban.config as c -import caliban.config.experiment as ce -import caliban.util as u -import caliban.util.fs as ufs -from caliban.history.types import Experiment, Job, JobSpec, JobStatus, Platform -from caliban.history.util import (create_experiments, generate_container_spec, - get_mem_engine, get_sql_engine, session_scope) - -t = Terminal() - -DEV_CONTAINER_ROOT = "gcr.io/blueshift-playground/blueshift" -TF_VERSIONS = {"2.2.0", "1.12.3", "1.14.0", "1.15.0"} -DEFAULT_WORKDIR = "/usr/app" -CREDS_DIR = "/.creds" -CONDA_BIN = "/opt/conda/bin/conda" - -ImageId = NewType('ImageId', str) -ArgSeq = NewType('ArgSeq', List[str]) - - -class DockerError(Exception): - """Exception that passes info on a failed Docker command.""" - - def __init__(self, message, cmd, ret_code): - super().__init__(message) - self.message = message - self.cmd = cmd - self.ret_code = ret_code - - @property - def command(self): - return " ".join(self.cmd) - - -class NotebookInstall(Enum): - """Flag to decide what to do .""" - none = 'none' - lab = 'lab' - jupyter = 'jupyter' - - def __str__(self) -> str: - return self.value - - -class Shell(Enum): - """Add new shells here and below, in SHELL_DICT.""" - bash = 'bash' - zsh = 'zsh' - - def __str__(self) -> str: - return self.value - - -# Tuple to track the information required to install and execute some custom -# shell into a container. -ShellData = NamedTuple("ShellData", [("executable", str), - ("packages", List[str])]) - - -def apt_install(*packages: str) -> str: - """Returns a command that will install the supplied list of packages without - requiring confirmation or any user interaction. - """ - package_str = ' '.join(packages) - no_prompt = "DEBIAN_FRONTEND=noninteractive" - return f"{no_prompt} apt-get install --yes --no-install-recommends {package_str}" - - -def apt_command(commands: List[str]) -> List[str]: - """Pre-and-ap-pends the supplied commands with the appropriate in-container and - cleanup command for aptitude. - - """ - update = ["apt-get update"] - cleanup = ["apt-get clean", "rm -rf /var/lib/apt/lists/*"] - return update + commands + cleanup - - -# Dict linking a particular supported shell to the data required to run and -# install the shell inside a container. -# -# : Dict[Shell, ShellData] -SHELL_DICT = { - Shell.bash: ShellData("/bin/bash", []), - Shell.zsh: ShellData("/bin/zsh", ["zsh"]) -} - - -def default_shell() -> Shell: - """Returns the shell to load into the container. Defaults to Shell.bash, but if - the user's SHELL variable refers to a supported sub-shell, returns that - instead. - - """ - ret = Shell.bash - - if "zsh" in os.environ.get("SHELL"): - ret = Shell.zsh - - return ret - - -def adc_location(home_dir: Optional[str] = None) -> str: - """Returns the location for application default credentials, INSIDE the - container (so, hardcoded unix separators), given the supplied home directory. - - """ - if home_dir is None: - home_dir = Path.home() - - return "{}/.config/gcloud/application_default_credentials.json".format( - home_dir) - - -def container_home(): - """Returns the location of the home directory inside the generated - container. - - """ - return "/home/{}".format(u.current_user()) - - -def tf_base_image(job_mode: c.JobMode, tensorflow_version: str) -> str: - """Returns the base image to use, depending on whether or not we're using a - GPU. This is JUST for building our base images for Blueshift; not for - actually using in a job. - - List of available tags: https://hub.docker.com/r/tensorflow/tensorflow/tags - - """ - if tensorflow_version not in TF_VERSIONS: - raise Exception("""{} is not a valid tensorflow version. - Try one of: {}""".format(tensorflow_version, TF_VERSIONS)) - - gpu = "-gpu" if c.gpu(job_mode) else "" - return "tensorflow/tensorflow:{}{}-py3".format(tensorflow_version, gpu) - - -def base_image_suffix(job_mode: c.JobMode) -> str: - return "gpu" if c.gpu(job_mode) else "cpu" - - -def base_image_id(job_mode: c.JobMode) -> str: - """Returns the default base image for all caliban Dockerfiles.""" - base_suffix = base_image_suffix(job_mode) - return "{}:{}".format(DEV_CONTAINER_ROOT, base_suffix) - - -def extras_string(extras: List[str]) -> str: - """Returns the argument passed to `pip install` to install a project from its - setup.py and target a specific set of extras_require dependencies. - - Args: - extras: (potentially empty) list of extra_requires deps. - """ - ret = "." - if len(extras) > 0: - ret += "[{}]".format(','.join(extras)) - return ret - - -def base_extras(job_mode: c.JobMode, path: str, - extras: Optional[List[str]]) -> Optional[List[str]]: - """Returns None if the supplied path doesn't exist (it's assumed it points to a - setup.py file). - - If the path DOES exist, generates a list of extras to install. gpu or cpu are - always added to the beginning of the list, depending on the mode. - - """ - ret = None - - if os.path.exists(path): - base = extras or [] - extra = 'gpu' if c.gpu(job_mode) else 'cpu' - ret = base if extra in base else [extra] + base - - return ret - - -def _dependency_entries(workdir: str, - user_id: int, - user_group: int, - requirements_path: Optional[str] = None, - conda_env_path: Optional[str] = None, - setup_extras: Optional[List[str]] = None) -> str: - """Returns the Dockerfile entries required to install dependencies from either: - - - a requirements.txt file, path supplied by requirements_path - - a conda environment.yml file, path supplied by conda_env_path. - - a setup.py file, if some sequence of dependencies is supplied. - - An empty list for setup_extras means, run `pip install -c .` with no extras. - None for this argument means do nothing. If a list of strings is supplied, - they'll be treated as extras dependency sets. - """ - ret = "" - - if setup_extras is not None: - ret += f""" -COPY --chown={user_id}:{user_group} setup.py {workdir} -RUN /bin/bash -c "pip install --no-cache-dir {extras_string(setup_extras)}" -""" - - if conda_env_path is not None: - ret += f""" -COPY --chown={user_id}:{user_group} {conda_env_path} {workdir} -RUN /bin/bash -c "{CONDA_BIN} env update \ - --quiet --name caliban \ - --file {conda_env_path} && \ - {CONDA_BIN} clean -y -q --all" -""" - - if requirements_path is not None: - ret += f""" -COPY --chown={user_id}:{user_group} {requirements_path} {workdir} -RUN /bin/bash -c "pip install --no-cache-dir -r {requirements_path}" -""" - - return ret - - -def _package_entries(workdir: str, user_id: int, user_group: int, - package: u.Package) -> str: - """Returns the Dockerfile entries required to: - - - copy a directory of code into a docker container - - inject an entrypoint that executes a python module inside that directory. - - Python code runs as modules vs scripts so that we can enforce import hygiene - between files inside a project. - - """ - owner = "{}:{}".format(user_id, user_group) - - arg = package.main_module or package.script_path - - # This needs to use json so that quotes print as double quotes, not single - # quotes. - entrypoint_s = json.dumps(package.executable + [arg]) - - return """ -# Copy project code into the docker container. -COPY --chown={owner} {package_path} {workdir}/{package_path} - -# Declare an entrypoint that actually runs the container. -ENTRYPOINT {entrypoint_s} - """.format_map({ - "owner": owner, - "package_path": package.package_path, - "workdir": workdir, - "entrypoint_s": entrypoint_s - }) - - -def _service_account_entry(user_id: int, user_group: int, credentials_path: str, - docker_credentials_dir: str, - write_adc_placeholder: bool): - """Generates the Dockerfile entries required to transfer a set of Cloud service - account credentials into the Docker container. - - NOTE the write_adc_placeholder variable is here because the "ctpu" script - that we use to interact with TPUs has a bug in it, as of 1/21/2020, where the - script will fail if the application_default_credentials.json file isn't - present, EVEN THOUGH it properly uses the service account credentials - registered with gcloud instead of ADC creds. - - If a service account is present, we write a placeholder string to get past - this problem. This shouldn't matter for anyone else since adc isn't used if a - service account is present. - - """ - container_creds = "{}/credentials.json".format(docker_credentials_dir) - ret = """ -COPY --chown={user_id}:{user_group} {credentials_path} {container_creds} - -# Use the credentials file to activate gcloud, gsutil inside the container. -RUN gcloud auth activate-service-account --key-file={container_creds} && \ - git config --global credential.'https://source.developers.google.com'.helper gcloud.sh - -ENV GOOGLE_APPLICATION_CREDENTIALS={container_creds} -""".format_map({ - "user_id": user_id, - "user_group": user_group, - "credentials_path": credentials_path, - "container_creds": container_creds - }) - - if write_adc_placeholder: - ret += """ -RUN echo "placeholder" >> {} -""".format(adc_location(container_home())) - - return ret - - -def _adc_entry(user_id: int, user_group: int, adc_path: str): - """Returns the Dockerfile line required to transfer the - application_default_credentials.json file into the container's home - directory. - - """ - return """ -COPY --chown={user_id}:{user_group} {adc_path} {adc_loc} - """.format_map({ - "user_id": user_id, - "user_group": user_group, - "adc_path": adc_path, - "adc_loc": adc_location(container_home()) - }) - - -def _credentials_entries(user_id: int, - user_group: int, - adc_path: Optional[str], - credentials_path: Optional[str], - docker_credentials_dir: Optional[str] = None) -> str: - """Returns the Dockerfile entries necessary to copy a user's Cloud credentials - into the Docker container. - - - adc_path is the relative path inside the current directory to an - application_default_credentials.json file containing... well, you get it. - - credentials_path is the relative path inside the current directory to a - JSON credentials file. - - docker_credentials_dir is the relative path inside the docker container - where the JSON file will be copied on build. - - """ - if docker_credentials_dir is None: - docker_credentials_dir = CREDS_DIR - - ret = "" - if credentials_path is not None: - ret += _service_account_entry(user_id, - user_group, - credentials_path, - docker_credentials_dir, - write_adc_placeholder=adc_path is None) - - if adc_path is not None: - ret += _adc_entry(user_id, user_group, adc_path) - - return ret - - -def _notebook_entries(lab: bool = False, version: Optional[str] = None) -> str: - """Returns the Dockerfile entries necessary to install Jupyter{lab}. - - Optionally takes a version string. - - """ - version_suffix = "" - - if version is not None: - version_suffix = "=={}".format(version) - - library = "jupyterlab" if lab else "jupyter" - - return """ -RUN pip install {}{} -""".format(library, version_suffix) - - -def _custom_packages( - user_id: int, - user_group: int, - packages: Optional[List[str]] = None, - shell: Optional[Shell] = None, -) -> str: - """Returns the Dockerfile entries necessary to install custom dependencies for - the supplied shell and sequence of aptitude packages. - - """ - if packages is None: - packages = [] - - if shell is None: - shell = Shell.bash - - ret = "" - - to_install = sorted(packages + SHELL_DICT[shell].packages) - - if len(to_install) != 0: - commands = apt_command([apt_install(*to_install)]) - ret = """ -USER root - -RUN {commands} - -USER {user_id}:{user_group} -""".format_map({ - "commands": " && ".join(commands), - "user_id": user_id, - "user_group": user_group - }) - - return ret - - -def _copy_dir_entry(workdir: str, user_id: int, user_group: int, - dirname: str) -> str: - """Returns the Dockerfile entry necessary to copy a single extra subdirectory - from the current directory into a docker container during build. - - """ - owner = "{}:{}".format(user_id, user_group) - return """# Copy {dirname} into the Docker container. -COPY --chown={owner} {dirname} {workdir}/{dirname} -""".format_map({ - "owner": owner, - "workdir": workdir, - "dirname": dirname - }) - - -def _extra_dir_entries(workdir: str, user_id: int, user_group: int, - extra_dirs: List[str]) -> str: - """Returns the Dockerfile entries necessary to copy all directories in the - extra_dirs list into a docker container during build. - - """ - ret = "" - for d in extra_dirs: - ret += "\n{}".format(_copy_dir_entry(workdir, user_id, user_group, d)) - return ret - - -def _dockerfile_template( - job_mode: c.JobMode, - workdir: Optional[str] = None, - base_image_fn: Optional[Callable[[c.JobMode], str]] = None, - package: Optional[Union[List, u.Package]] = None, - requirements_path: Optional[str] = None, - conda_env_path: Optional[str] = None, - setup_extras: Optional[List[str]] = None, - adc_path: Optional[str] = None, - credentials_path: Optional[str] = None, - jupyter_version: Optional[str] = None, - inject_notebook: NotebookInstall = NotebookInstall.none, - shell: Optional[Shell] = None, - extra_dirs: Optional[List[str]] = None, - caliban_config: Optional[Dict[str, Any]] = None) -> str: - """Returns a Dockerfile that builds on a local CPU or GPU base image (depending - on the value of job_mode) to create a container that: - - - installs any dependency specified in a requirements.txt file living at - requirements_path, a conda environment at conda_env_path, or any - dependencies in a setup.py file, including extra dependencies, if - setup_extras isn't None - - injects gcloud credentials into the container, so Cloud interaction works - just like it does locally - - potentially installs a custom shell, or jupyterlab for notebook support - - copies all source needed by the main module specified by package, and - potentially injects an entrypoint that, on run, will run that main module - - Most functions that call _dockerfile_template pass along any kwargs that they - receive. It should be enough to add kwargs here, then rely on that mechanism - to pass them along, vs adding kwargs all the way down the call chain. - - Supply a custom base_image_fn (function from job_mode -> image ID) to inject - more complex Docker commands into the Caliban environments by, for example, - building your own image on top of the TF base images, then using that. - - """ - uid = os.getuid() - gid = os.getgid() - username = u.current_user() - - if isinstance(package, list): - package = u.Package(*package) - - if workdir is None: - workdir = DEFAULT_WORKDIR - - if base_image_fn is None: - base_image_fn = base_image_id - - base_image = base_image_fn(job_mode) - - dockerfile = """ -FROM {base_image} - -# Create the same group we're using on the host machine. -RUN [ $(getent group {gid}) ] || groupadd --gid {gid} {gid} - -# Create the user by name. --no-log-init guards against a crash with large user -# IDs. -RUN useradd --no-log-init --no-create-home -u {uid} -g {gid} --shell /bin/bash {username} - -# The directory is created by root. This sets permissions so that any user can -# access the folder. -RUN mkdir -m 777 {workdir} {creds_dir} {c_home} - -ENV HOME={c_home} - -WORKDIR {workdir} - -USER {uid}:{gid} -""".format_map({ - "base_image": base_image, - "username": username, - "uid": uid, - "gid": gid, - "workdir": workdir, - "c_home": container_home(), - "creds_dir": CREDS_DIR - }) - dockerfile += _credentials_entries(uid, - gid, - adc_path=adc_path, - credentials_path=credentials_path) - - dockerfile += _dependency_entries(workdir, - uid, - gid, - requirements_path=requirements_path, - conda_env_path=conda_env_path, - setup_extras=setup_extras) - - if inject_notebook.value != 'none': - install_lab = inject_notebook == NotebookInstall.lab - dockerfile += _notebook_entries(lab=install_lab, version=jupyter_version) - - if extra_dirs is not None: - dockerfile += _extra_dir_entries(workdir, uid, gid, extra_dirs) - - dockerfile += _custom_packages(uid, - gid, - packages=c.apt_packages( - caliban_config, job_mode), - shell=shell) - - if package is not None: - # The actual entrypoint and final copied code. - dockerfile += _package_entries(workdir, uid, gid, package) - - return dockerfile - - -def docker_image_id(output: str) -> ImageId: - """Accepts a string containing the output of a successful `docker build` - command and parses the Docker image ID from the stream. - - NOTE this is probably quite brittle! I can imagine this breaking quite easily - on a Docker upgrade. - - """ - return ImageId(output.splitlines()[-1].split()[-1]) - - -def build_image(job_mode: c.JobMode, - build_path: str, - credentials_path: Optional[str] = None, - adc_path: Optional[str] = None, - no_cache: bool = False, - **kwargs) -> str: - """Builds a Docker image by generating a Dockerfile and passing it to `docker - build` via stdin. All output from the `docker build` process prints to - stdout. - - Returns the image ID of the new docker container; if the command fails, - throws on error with information about the command and any issues that caused - the problem. - - """ - with u.TempCopy(credentials_path, - tmp_name=".caliban_default_creds.json") as creds: - with u.TempCopy(adc_path, tmp_name=".caliban_adc_creds.json") as adc: - cache_args = ["--no-cache"] if no_cache else [] - cmd = ["docker", "build"] + cache_args + ["--rm", "-f-", build_path] - - dockerfile = _dockerfile_template(job_mode, - credentials_path=creds, - adc_path=adc, - **kwargs) - - joined_cmd = " ".join(cmd) - logging.info("Running command: {}".format(joined_cmd)) - - try: - output, ret_code = ufs.capture_stdout(cmd, input_str=dockerfile) - if ret_code == 0: - return docker_image_id(output) - else: - error_msg = "Docker failed with error code {}.".format(ret_code) - raise DockerError(error_msg, cmd, ret_code) - - except subprocess.CalledProcessError as e: - logging.error(e.output) - logging.error(e.stderr) - - -def _image_tag_for_project(project_id: str, image_id: str) -> str: - """Generate the GCR Docker image tag for the supplied pair of project_id and - image_id. - - This function properly handles "domain scoped projects", where the project ID - contains a domain name and project ID separated by : - https://cloud.google.com/container-registry/docs/overview#domain-scoped_projects. - - """ - project_s = project_id.replace(":", "/") - return "gcr.io/{}/{}:latest".format(project_s, image_id) - - -def push_uuid_tag(project_id: str, image_id: str) -> str: - """Takes a base image and tags it for upload, then pushes it to a remote Google - Container Registry. - - Returns the tag on a successful push. - - TODO should this just check first before attempting to push if the image - exists? Immutable names means that if the tag is up there, we're done. - Potentially use docker-py for this. - - """ - image_tag = _image_tag_for_project(project_id, image_id) - subprocess.run(["docker", "tag", image_id, image_tag], check=True) - subprocess.run(["docker", "push", image_tag], check=True) - return image_tag - - -def _run_cmd(job_mode: c.JobMode, - run_args: Optional[List[str]] = None) -> List[str]: - """Returns the sequence of commands for the subprocess run functions required - to execute `docker run`. in CPU or GPU mode, depending on the value of - job_mode. - - Keyword args: - - run_args: list of args to pass to docker run. - - """ - if run_args is None: - run_args = [] - - runtime = ["--runtime", "nvidia"] if c.gpu(job_mode) else [] - return ["docker", "run"] + runtime + ["--ipc", "host"] + run_args +import caliban.docker.build as b +import caliban.platform.run as r def _home_mount_cmds(enable_home_mount: bool) -> List[str]: @@ -683,7 +36,7 @@ def _home_mount_cmds(enable_home_mount: bool) -> List[str]: """ ret = [] if enable_home_mount: - ret = ["-v", "{}:{}".format(Path.home(), container_home())] + ret = ["-v", "{}:{}".format(Path.home(), b.container_home())] return ret @@ -692,263 +45,18 @@ def _interactive_opts(workdir: str) -> List[str]: """ return [ - "-w", workdir, \ + "-w", workdir, "-u", "{}:{}".format(os.getuid(), os.getgid()), \ "-v", "{}:{}".format(os.getcwd(), workdir) \ ] -def log_job_spec_instance(job_spec: JobSpec, i: int) -> JobSpec: - """Prints logging as a side effect for the supplied sequence of job specs - generated from an experiment definition; returns the input job spec. - - """ - args = ce.experiment_to_args(job_spec.experiment.kwargs, - job_spec.experiment.args) - logging.info("") - logging.info("Job {} - Experiment args: {}".format(i, t.yellow(str(args)))) - return job_spec - - -def logged_job_specs(job_specs: Iterable[JobSpec]) -> Iterable[JobSpec]: - """Generates an iterable of job specs that should be passed to `docker run` to - execute the experiments defined by the supplied iterable. - - """ - for i, s in enumerate(job_specs, 1): - yield log_job_spec_instance(s, i) - - -def execute_dry_run(job_specs: Iterable[JobSpec]) -> None: - """Expands the supplied sequence of experiments into sequences of args and logs - the jobs that WOULD have been executed, had the dry run flag not been - applied. - - """ - list(logged_job_specs(job_specs)) - - logging.info('') - logging.info( - t.yellow("To build your image and execute these jobs, \ -run your command again without {}.".format(c.DRY_RUN_FLAG))) - logging.info('') - return None - - -def local_callback(idx: int, job: Job) -> None: - """Provides logging feedback for jobs run locally. If the return code is 0, - logs success; else, logs the failure as an error and logs the script args - that provided the failure. - - """ - if job.status == JobStatus.SUCCEEDED: - logging.info(t.green(f'Job {idx} succeeded!')) - else: - logging.error( - t.red(f'Job {idx} failed with return code {job.details["ret_code"]}.')) - args = ce.experiment_to_args(job.spec.experiment.kwargs, - job.spec.experiment.args) - logging.error(t.red(f'Failing args for job {idx}: {args}')) - - -def window_size_env_cmds(): - """Returns a sequence of `docker run` arguments that will internally configure - the terminal columns and lines, so that progress bars and other terminal - interactions will work properly. - - These aren't required for interactive Docker commands like those triggered by - `caliban shell`. - - """ - ret = [] - cols, lines = _screen_shape_wrapper()(0) - if cols: - ret += ["-e", f"COLUMNS={cols}"] - if lines: - ret += ["-e", f"LINES={lines}"] - return ret - - -# ---------------------------------------------------------------------------- -def _create_job_spec_dict( - experiment: Experiment, - job_mode: c.JobMode, - image_id: str, - run_args: Optional[List[str]] = None, -) -> Dict[str, Any]: - '''creates a job spec dictionary for a local job''' - - # Without the unbuffered environment variable, stderr and stdout won't be - # emitted in the proper order from inside the container. - terminal_cmds = ["-e" "PYTHONUNBUFFERED=1"] + window_size_env_cmds() - - base_cmd = _run_cmd(job_mode, run_args) + terminal_cmds + [image_id] - command = base_cmd + ce.experiment_to_args(experiment.kwargs, experiment.args) - return {'command': command, 'container': image_id} - - -# ---------------------------------------------------------------------------- -def execute_jobs( - job_specs: Iterable[JobSpec], - dry_run: bool = False, -): - '''executes a sequence of jobs based on job specs - - Arg: - job_specs: specifications for jobs to be executed - dry_run: if True, only print what would be done - ''' - - with u.tqdm_logging() as orig_stream: - pbar = tqdm.tqdm(logged_job_specs(job_specs), - file=orig_stream, - total=len(job_specs), - ascii=True, - unit="experiment", - desc="Executing") - for idx, job_spec in enumerate(pbar, 1): - command = job_spec.spec['command'] - logging.info(f'Running command: {" ".join(command)}') - if not dry_run: - _, ret_code = ufs.capture_stdout(command, "", u.TqdmFile(sys.stderr)) - else: - ret_code = 0 - j = Job(spec=job_spec, - container=job_spec.spec['container'], - details={'ret_code': ret_code}, - status=JobStatus.SUCCEEDED if ret_code == 0 else JobStatus.FAILED) - local_callback(idx=idx, job=j) - - if dry_run: - logging.info( - t.yellow(f'\nTo build your image and execute these jobs, ' - f'run your command again without {c.DRY_RUN_FLAG}\n')) - - return None - - -def run_experiments(job_mode: c.JobMode, - run_args: Optional[List[str]] = None, - script_args: Optional[List[str]] = None, - image_id: Optional[str] = None, - dry_run: bool = False, - experiment_config: Optional[ce.ExpConf] = None, - xgroup: Optional[str] = None, - **build_image_kwargs) -> None: - """Builds an image using the supplied **build_image_kwargs and calls `docker - run` on the resulting image using sensible defaults. - - Keyword args: - - - job_mode: c.JobMode. - - - run_args: extra arguments to supply to `docker run` after our defaults. - - script_args: extra arguments to supply to the entrypoint. (You can - - override the default container entrypoint by supplying a new one inside - run_args.) - - image_id: ID of the image to run. Supplying this will skip an image build. - - experiment_config: dict of string to list, boolean, string or int. Any - lists will trigger a cartesian product out with the rest of the config. A - job will be executed for every combination of parameters in the experiment - config. - - dry_run: if True, no actual jobs will be executed and docker won't - actually build; logging side effects will show the user what will happen - without dry_run=True. - - any extra kwargs supplied are passed through to build_image. - """ - if run_args is None: - run_args = [] - - if script_args is None: - script_args = [] - - if experiment_config is None: - experiment_config = {} - - docker_args = {k: v for k, v in build_image_kwargs.items()} - docker_args['job_mode'] = job_mode - - engine = get_mem_engine() if dry_run else get_sql_engine() - - with session_scope(engine) as session: - container_spec = generate_container_spec(session, docker_args, image_id) - - if image_id is None: - if dry_run: - logging.info("Dry run - skipping actual 'docker build'.") - image_id = 'dry_run_tag' - else: - image_id = build_image(**docker_args) - - experiments = create_experiments( - session=session, - container_spec=container_spec, - script_args=script_args, - experiment_config=experiment_config, - xgroup=xgroup, - ) - - job_specs = [ - JobSpec.get_or_create( - experiment=x, - spec=_create_job_spec_dict( - experiment=x, - job_mode=job_mode, - run_args=run_args, - image_id=image_id, - ), - platform=Platform.LOCAL, - ) for x in experiments - ] - - try: - execute_jobs(job_specs=job_specs, dry_run=dry_run) - except Exception as e: - logging.error(f'exception: {e}') - session.commit() # commit here, otherwise will be rolled back - - -def run(job_mode: c.JobMode, - run_args: Optional[List[str]] = None, - script_args: Optional[List[str]] = None, - image_id: Optional[str] = None, - **build_image_kwargs) -> None: - """Builds an image using the supplied **build_image_kwargs and calls `docker - run` on the resulting image using sensible defaults. - Keyword args: - - job_mode: c.JobMode. - - run_args: extra arguments to supply to `docker run` after our defaults. - - script_args: extra arguments to supply to the entrypoint. (You can - - override the default container entrypoint by supplying a new one inside - run_args.) - - image_id: ID of the image to run. Supplying this will skip an image build. - any extra kwargs supplied are passed through to build_image. - """ - if run_args is None: - run_args = [] - - if script_args is None: - script_args = [] - - if image_id is None: - image_id = build_image(job_mode, **build_image_kwargs) - - base_cmd = _run_cmd(job_mode, run_args) - - command = base_cmd + [image_id] + script_args - - logging.info("Running command: {}".format(' '.join(command))) - subprocess.call(command) - return None - - def run_interactive(job_mode: c.JobMode, workdir: Optional[str] = None, image_id: Optional[str] = None, run_args: Optional[List[str]] = None, mount_home: Optional[bool] = None, - shell: Optional[Shell] = None, + shell: Optional[b.Shell] = None, entrypoint: Optional[str] = None, entrypoint_args: Optional[List[str]] = None, **build_image_kwargs) -> None: @@ -972,7 +80,7 @@ def run_interactive(job_mode: c.JobMode, """ if workdir is None: - workdir = DEFAULT_WORKDIR + workdir = b.DEFAULT_WORKDIR if run_args is None: run_args = [] @@ -986,20 +94,20 @@ def run_interactive(job_mode: c.JobMode, if shell is None: # Only set a default shell if we're also mounting the home volume. # Otherwise a custom shell won't have access to the user's profile. - shell = default_shell() if mount_home else Shell.bash + shell = b.default_shell() if mount_home else b.Shell.bash if entrypoint is None: - entrypoint = SHELL_DICT[shell].executable + entrypoint = b.SHELL_DICT[shell].executable interactive_run_args = _interactive_opts(workdir) + [ "-it", \ "--entrypoint", entrypoint ] + _home_mount_cmds(mount_home) + run_args - run(job_mode=job_mode, - run_args=interactive_run_args, - script_args=entrypoint_args, - image_id=image_id, - shell=shell, - workdir=workdir, - **build_image_kwargs) + r.run(job_mode=job_mode, + run_args=interactive_run_args, + script_args=entrypoint_args, + image_id=image_id, + shell=shell, + workdir=workdir, + **build_image_kwargs) From f613c1132a49aeeafa1453c5b02b52de1e6ae84e Mon Sep 17 00:00:00 2001 From: Sam Ritchie Date: Sat, 11 Jul 2020 13:37:19 -0600 Subject: [PATCH 08/15] all works --- caliban/cli.py | 18 +++++++++--------- caliban/main.py | 4 ++-- caliban/platform/cloud/core.py | 8 +++++--- caliban/platform/gke/cli.py | 3 ++- caliban/platform/run.py | 16 ++++++++-------- caliban/util/argparse.py | 5 +++-- 6 files changed, 29 insertions(+), 25 deletions(-) diff --git a/caliban/cli.py b/caliban/cli.py index 7d18ad6..f90894b 100644 --- a/caliban/cli.py +++ b/caliban/cli.py @@ -26,13 +26,14 @@ import caliban.config as conf import caliban.config.experiment as ce -import caliban.docker.build as docker +import caliban.docker.build as b import caliban.platform.cloud.types as ct import caliban.platform.gke as gke 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.util as u +import caliban.util.argparse as ua from caliban import __version__ t = Terminal() @@ -127,7 +128,7 @@ def add_script_args(parser): def require_module(parser): parser.add_argument( "module", - type=u.validated_package, + type=ua.validated_package, help= "Code to execute, in either trainer.train' or 'trainer/train.py' format. " "Accepts python scripts, modules or a path to an arbitrary script.") @@ -158,7 +159,7 @@ def extra_dirs(parser): "-d", "--dir", action="append", - type=u.validated_directory, + type=ua.validated_directory, help="Extra directories to include. List these from large to small " "to take full advantage of Docker's build cache.") @@ -190,7 +191,7 @@ def region_arg(parser): def cloud_key_arg(parser): parser.add_argument("--cloud_key", - type=u.validated_file, + type=ua.validated_file, help="Path to GCloud service account key. " "(Defaults to $GOOGLE_APPLICATION_CREDENTIALS.)") @@ -264,8 +265,8 @@ def shell_parser(base): docker_run_arg(parser) parser.add_argument( "--shell", - choices=docker.Shell, - type=docker.Shell, + choices=b.Shell, + type=b.Shell, help= """This argument sets the shell used inside the container to one of Caliban's supported shells. Defaults to the shell specified by the $SHELL environment @@ -362,7 +363,7 @@ def label_arg(parser): "--label", metavar="KEY=VALUE", action="append", - type=u.parse_kv_pair, + type=ua.parse_kv_pair, help="Extra label k=v pair to submit to Cloud.") @@ -540,8 +541,7 @@ def generate_docker_args(job_mode: conf.JobMode, # Get extra dependencies in case you want to install your requirements via a # setup.py file. - setup_extras = docker.build.base_extras(job_mode, "setup.py", - args.get("extras")) + setup_extras = b.base_extras(job_mode, "setup.py", args.get("extras")) # Google application credentials, from the CLI or from an env variable. creds_path = conf.extract_cloud_key(args) diff --git a/caliban/main.py b/caliban/main.py index 1b34e0d..d4a52f5 100644 --- a/caliban/main.py +++ b/caliban/main.py @@ -28,12 +28,12 @@ import caliban.docker.build as b import caliban.history.cli import caliban.platform.cloud.core as cloud +import caliban.platform.cloud.util as cu import caliban.platform.gke as gke import caliban.platform.gke.cli import caliban.platform.notebook as pn import caliban.platform.run as pr import caliban.platform.shell as ps -import caliban.util as u ll.getLogger('caliban.main').setLevel(logging.ERROR) t = Terminal() @@ -123,7 +123,7 @@ def run_app(arg_input): image_tag = args.get("image_tag") machine_type = args.get("machine_type") exp_config = args.get("experiment_config") - labels = u.sanitize_labels(args.get("label") or []) + labels = cu.sanitize_labels(args.get("label") or []) xgroup = args.get('xgroup') # Arguments to internally build the image required to submit to Cloud. diff --git a/caliban/platform/cloud/core.py b/caliban/platform/cloud/core.py index 056c2be..6af5af7 100644 --- a/caliban/platform/cloud/core.py +++ b/caliban/platform/cloud/core.py @@ -37,7 +37,9 @@ import caliban.docker.push as dp import caliban.history.types as ht import caliban.platform.cloud.types as ct +import caliban.platform.cloud.util as cu import caliban.util as u +import caliban.util.tqdm as ut from caliban.history.util import (create_experiments, generate_container_spec, get_mem_engine, get_sql_engine, session_scope) @@ -368,7 +370,7 @@ def execute_requests( generated handle any response or exception. """ - with u.tqdm_logging() as orig_stream: + with ut.tqdm_logging() as orig_stream: pbar = tqdm.tqdm( requests, file=orig_stream, @@ -437,8 +439,8 @@ def _job_spec( "jobId": job_id, "trainingInput": training_input, "labels": { - **u.sanitize_labels(labels), - **u.script_args_to_labels(job_args) + **cu.sanitize_labels(labels), + **cu.script_args_to_labels(job_args) } }, platform=ht.Platform.CAIP, diff --git a/caliban/platform/gke/cli.py b/caliban/platform/gke/cli.py index 625c4b2..85b2ad3 100644 --- a/caliban/platform/gke/cli.py +++ b/caliban/platform/gke/cli.py @@ -28,6 +28,7 @@ import caliban.cli as cli import caliban.config as conf +import caliban.platform.cloud.util as cu import caliban.platform.gke.constants as k import caliban.platform.gke.util as util import caliban.util as u @@ -361,7 +362,7 @@ def _job_submit(args: dict, cluster: Cluster) -> None: labels = args.get('label') if labels is not None: - labels = dict(u.sanitize_labels(args.get('label'))) + labels = dict(cu.sanitize_labels(args.get('label'))) # Arguments to internally build the image required to submit to Cloud. docker_m = {'job_mode': job_mode, 'package': package, **docker_args} diff --git a/caliban/platform/run.py b/caliban/platform/run.py index c55beb9..df20d37 100644 --- a/caliban/platform/run.py +++ b/caliban/platform/run.py @@ -32,8 +32,8 @@ import caliban.config as c import caliban.config.experiment as ce import caliban.docker.build as b -import caliban.util as u import caliban.util.fs as ufs +import caliban.util.tqdm as ut from caliban.history.types import Experiment, Job, JobSpec, JobStatus, Platform from caliban.history.util import (create_experiments, generate_container_spec, get_mem_engine, get_sql_engine, session_scope) @@ -63,8 +63,8 @@ def log_job_spec_instance(job_spec: JobSpec, i: int) -> JobSpec: generated from an experiment definition; returns the input job spec. """ - args = c.experiment_to_args(job_spec.experiment.kwargs, - job_spec.experiment.args) + args = ce.experiment_to_args(job_spec.experiment.kwargs, + job_spec.experiment.args) logging.info("") logging.info("Job {} - Experiment args: {}".format(i, t.yellow(str(args)))) return job_spec @@ -106,8 +106,8 @@ def local_callback(idx: int, job: Job) -> None: else: logging.error( t.red(f'Job {idx} failed with return code {job.details["ret_code"]}.')) - args = c.experiment_to_args(job.spec.experiment.kwargs, - job.spec.experiment.args) + args = ce.experiment_to_args(job.spec.experiment.kwargs, + job.spec.experiment.args) logging.error(t.red(f'Failing args for job {idx}: {args}')) @@ -143,7 +143,7 @@ def _create_job_spec_dict( terminal_cmds = ["-e" "PYTHONUNBUFFERED=1"] + window_size_env_cmds() base_cmd = _run_cmd(job_mode, run_args) + terminal_cmds + [image_id] - command = base_cmd + c.experiment_to_args(experiment.kwargs, experiment.args) + command = base_cmd + ce.experiment_to_args(experiment.kwargs, experiment.args) return {'command': command, 'container': image_id} @@ -159,7 +159,7 @@ def execute_jobs( dry_run: if True, only print what would be done ''' - with u.tqdm_logging() as orig_stream: + with ut.tqdm_logging() as orig_stream: pbar = tqdm.tqdm(logged_job_specs(job_specs), file=orig_stream, total=len(job_specs), @@ -170,7 +170,7 @@ def execute_jobs( command = job_spec.spec['command'] logging.info(f'Running command: {" ".join(command)}') if not dry_run: - _, ret_code = ufs.capture_stdout(command, "", u.TqdmFile(sys.stderr)) + _, ret_code = ufs.capture_stdout(command, "", ut.TqdmFile(sys.stderr)) else: ret_code = 0 j = Job(spec=job_spec, diff --git a/caliban/util/argparse.py b/caliban/util/argparse.py index 53f0b2f..efd6e34 100644 --- a/caliban/util/argparse.py +++ b/caliban/util/argparse.py @@ -24,6 +24,7 @@ from blessings import Terminal import caliban.util as u +import caliban.util.fs as ufs t = Terminal() @@ -43,7 +44,7 @@ def validated_package(path: str) -> u.Package: don't actually exist in the filesystem. """ - p = u.generate_package(path) + p = ufs.generate_package(path) if not os.path.isdir(p.package_path): raise argparse.ArgumentTypeError( @@ -52,7 +53,7 @@ def validated_package(path: str) -> u.Package: p.package_path)) filename = p.script_path - if not file_exists_in_cwd(filename): + if not ufs.file_exists_in_cwd(filename): raise argparse.ArgumentTypeError( """File '{}' doesn't exist locally as a script or python module; code must live inside the current directory.""".format(filename)) From 3329f462579159b4d0f9ac0e98604e9c0bfb154b Mon Sep 17 00:00:00 2001 From: Sam Ritchie Date: Sat, 11 Jul 2020 13:45:31 -0600 Subject: [PATCH 09/15] refactor auth --- caliban/history/util.py | 9 +++-- caliban/platform/cloud/core.py | 35 ++----------------- caliban/util/auth.py | 62 ++++++++++++++++++++++++++++++++++ 3 files changed, 71 insertions(+), 35 deletions(-) create mode 100644 caliban/util/auth.py diff --git a/caliban/history/util.py b/caliban/history/util.py index c701390..f414d19 100644 --- a/caliban/history/util.py +++ b/caliban/history/util.py @@ -29,6 +29,7 @@ from sqlalchemy.orm import Session, sessionmaker import caliban.config.experiment as ce +import caliban.util.auth as ua from caliban.history.types import (ContainerSpec, Experiment, ExperimentGroup, Job, JobSpec, JobStatus, Platform, init_db) from caliban.platform.cloud.types import JobStatus as CloudStatus @@ -243,8 +244,12 @@ def _get_caip_job_name(j: Job) -> str: # ---------------------------------------------------------------------------- -def _get_caip_job_api() -> Any: - return discovery.build('ml', 'v1', cache_discovery=False).projects().jobs() +def _get_caip_job_api(credentials_path: Optional[str] = None) -> Any: + credentials = ua.gcloud_credentials(credentials_path) + return discovery.build('ml', + 'v1', + cache_discovery=False, + credentials=credentials).projects().jobs() # ---------------------------------------------------------------------------- diff --git a/caliban/platform/cloud/core.py b/caliban/platform/cloud/core.py index 6af5af7..1e75cc7 100644 --- a/caliban/platform/cloud/core.py +++ b/caliban/platform/cloud/core.py @@ -39,6 +39,7 @@ import caliban.platform.cloud.types as ct import caliban.platform.cloud.util as cu import caliban.util as u +import caliban.util.auth as ua import caliban.util.tqdm as ut from caliban.history.util import (create_experiments, generate_container_spec, get_mem_engine, get_sql_engine, session_scope) @@ -46,28 +47,6 @@ t = Terminal() -def auth_access_token() -> Optional[str]: - """Attempts to fetch the local Oauth2 access token from the user's environment. - Returns the token if it exists, or None if not - - """ - try: - return check_output(['gcloud', 'auth', 'print-access-token'], - encoding='utf8').rstrip() - except CalledProcessError: - return None - - -def gcloud_auth_credentials() -> Optional[Credentials]: - """Attempt to generate credentials from the oauth2 workflow triggered by - `gcloud auth login`. Returns - - """ - token = auth_access_token() - if token: - return Credentials(token) - - def get_accelerator_config(gpu_spec: Optional[ct.GPUSpec]) -> Dict[str, Any]: """Returns the accelerator config for the supplied GPUSpec if present; else, returns the default accelerator config. @@ -282,17 +261,7 @@ def ml_api(credentials_path: Optional[str] = None): The actual details are a little byzantine. """ - credentials = None - - if credentials_path is not None: - credentials = service_account.Credentials.from_service_account_file( - credentials_path) - else: - # attempt to fetch credentials acquired via `gcloud auth login`. If this - # fails, the following API object will attempt to use application default - # credentials. - credentials = gcloud_auth_credentials() - + credentials = ua.gcloud_credentials(credentials_path) return discovery.build('ml', 'v1', cache_discovery=False, diff --git a/caliban/util/auth.py b/caliban/util/auth.py new file mode 100644 index 0000000..77a4940 --- /dev/null +++ b/caliban/util/auth.py @@ -0,0 +1,62 @@ +#!/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. +"""Utilities for interacting with the filesystem and packages. + +""" + +from subprocess import CalledProcessError, check_output +from typing import Optional + +from google.oauth2 import service_account +from google.oauth2.credentials import Credentials + + +def auth_access_token() -> Optional[str]: + """Attempts to fetch the local Oauth2 access token from the user's environment. + Returns the token if it exists, or None if not + + """ + try: + return check_output(['gcloud', 'auth', 'print-access-token'], + encoding='utf8').rstrip() + except CalledProcessError: + return None + + +def gcloud_auth_credentials() -> Optional[Credentials]: + """Attempt to generate credentials from the oauth2 workflow triggered by + `gcloud auth login`. Returns + + """ + token = auth_access_token() + if token: + return Credentials(token) + + +def gcloud_credentials( + credentials_path: Optional[str] = None) -> Optional[Credentials]: + credentials = None + + if credentials_path is not None: + credentials = service_account.Credentials.from_service_account_file( + credentials_path) + else: + # attempt to fetch credentials acquired via `gcloud auth login`. If this + # fails, the following API object will attempt to use application default + # credentials. + credentials = gcloud_auth_credentials() + + return credentials From 36ff6b1233c57de4993b8caa29f10eb017dc9f95 Mon Sep 17 00:00:00 2001 From: Sam Ritchie Date: Mon, 13 Jul 2020 09:47:01 -0600 Subject: [PATCH 10/15] add basic calibanconfig schema --- caliban/config/__init__.py | 19 +++++++++++++++++++ setup.py | 1 + 2 files changed, 20 insertions(+) diff --git a/caliban/config/__init__.py b/caliban/config/__init__.py index ed3a05e..8b579e1 100644 --- a/caliban/config/__init__.py +++ b/caliban/config/__init__.py @@ -27,6 +27,7 @@ import commentjson import yaml +from schema import And, Optional, Or, Schema, Use import caliban.platform.cloud.types as ct @@ -59,6 +60,24 @@ class JobMode(str, Enum): "type": "ACCELERATOR_TYPE_UNSPECIFIED" } +# Schema for Caliban Config + +AptPackages = Schema( + Or([str], { + Optional("gpu", default=list): [str], + Optional("cpu", default=list): [str] + })) + +CalibanConfig = Schema({ + Optional("project_id"): And(str, len), + Optional("cloud_key"): And(str, len), + Optional("base_image"): str, + Optional("apt_packages", default=dict): AptPackages, + Optional(str): str, +}) + +# Accessors + def gpu(job_mode: JobMode) -> bool: """Returns True if the supplied JobMode is JobMode.GPU, False otherwise. diff --git a/setup.py b/setup.py index a72392d..a66b300 100644 --- a/setup.py +++ b/setup.py @@ -49,6 +49,7 @@ def readme(): 'google-cloud-core>=1.0.3', 'google-cloud-container>=0.3.0', 'psycopg2-binary==2.8.5', + 'schema==0.7.2', 'urllib3>=1.25.7', 'yaspin>=0.16.0', # This is not a real dependency of ours, but we need it to override the From f317b23007d1becf2ccbb41ff5d6e57580dc8d2d Mon Sep 17 00:00:00 2001 From: Sam Ritchie Date: Tue, 14 Jul 2020 13:00:19 -0600 Subject: [PATCH 11/15] making it happen --- caliban/cli.py | 3 -- caliban/config/__init__.py | 81 ++++++++++++++++++++++------- caliban/main.py | 3 +- tutorials/basic/.calibanconfig.json | 1 + 4 files changed, 64 insertions(+), 24 deletions(-) create mode 100644 tutorials/basic/.calibanconfig.json diff --git a/caliban/cli.py b/caliban/cli.py index f90894b..8a180cc 100644 --- a/caliban/cli.py +++ b/caliban/cli.py @@ -22,7 +22,6 @@ import google.auth._cloud_sdk as csdk from absl.flags import argparse_flags -from blessings import Terminal import caliban.config as conf import caliban.config.experiment as ce @@ -36,8 +35,6 @@ import caliban.util.argparse as ua from caliban import __version__ -t = Terminal() - def _job_mode(use_gpu: bool, gpu_spec: Optional[ct.GPUSpec], tpu_spec: Optional[ct.TPUSpec]) -> conf.JobMode: diff --git a/caliban/config/__init__.py b/caliban/config/__init__.py index 8b579e1..4ce9160 100644 --- a/caliban/config/__init__.py +++ b/caliban/config/__init__.py @@ -17,25 +17,31 @@ Utilities for our job runner, for working with configs. """ -from __future__ import absolute_import, division, print_function - import argparse import os import sys +from contextlib import contextmanager from enum import Enum from typing import Any, Dict, List, Optional import commentjson +import schema as s import yaml -from schema import And, Optional, Or, Schema, Use import caliban.platform.cloud.types as ct +import caliban.util as u class JobMode(str, Enum): + """Represents the two modes that you can use to execute a Caliban job.""" CPU = 'CPU' GPU = 'GPU' + @staticmethod + def parse(label): + + return JobMode(label.upper()) + # Special config for Caliban. CalibanConfig = Dict[str, Any] @@ -62,18 +68,31 @@ class JobMode(str, Enum): # Schema for Caliban Config -AptPackages = Schema( - Or([str], { - Optional("gpu", default=list): [str], - Optional("cpu", default=list): [str] - })) - -CalibanConfig = Schema({ - Optional("project_id"): And(str, len), - Optional("cloud_key"): And(str, len), - Optional("base_image"): str, - Optional("apt_packages", default=dict): AptPackages, - Optional(str): str, +AptPackages = s.Schema( + s.Or( + [str], { + s.Optional("gpu", default=list): [str], + s.Optional("cpu", default=list): [str] + })) + +CCSchema = s.Schema({ + s.Optional("build_time_credentials", default=False): + bool, + s.Optional("default_mode", default=JobMode.CPU): + s.Use(JobMode.parse), + s.Optional("project_id"): + s.And(str, len), + s.Optional("cloud_key"): + s.And(str, len), + s.Optional("base_image"): + str, + s.Optional("apt_packages", default=dict): + AptPackages, + + # Allow extra entries without killing the schema to allow for backwards + # compatibility. + s.Optional(str): + str, }) # Accessors @@ -200,13 +219,35 @@ def apt_packages(conf: CalibanConfig, mode: JobMode) -> List[str]: format(CALIBAN_CONFIG, packages)) -def caliban_config() -> CalibanConfig: +def caliban_config(conf_path: str = CALIBAN_CONFIG) -> CalibanConfig: """Returns a dict that represents a `.calibanconfig.json` file if present, empty dictionary otherwise. + """ - if not os.path.isfile(CALIBAN_CONFIG): + if not os.path.isfile(conf_path): return {} - with open(CALIBAN_CONFIG) as f: - conf = commentjson.load(f) - return conf + conf = load_config(conf_path, mode='json') + return CCSchema.validate(conf) + + +@contextmanager +def argparse_schema(*args, **kwds): + """This function should work as a context manager that will trap a SchemaError + and return an argparse.ArgumentTypeError instead. + + TODO move to util.argparse + + """ + try: + yield + except s.SchemaError as e: + raise argparse.ArgumentTypeError(e.code) from None + + +@contextmanager +def error_schema(*args, **kwds): + try: + yield + except s.SchemaError as e: + u.err(e.code) diff --git a/caliban/main.py b/caliban/main.py index d4a52f5..07df504 100644 --- a/caliban/main.py +++ b/caliban/main.py @@ -154,7 +154,8 @@ def run_app(arg_input): def main(): logging.use_python_logging() try: - app.run(run_app, flags_parser=cli.parse_flags) + with c.argparse_schema(): + app.run(run_app, flags_parser=cli.parse_flags) except KeyboardInterrupt: logging.info('Shutting down.') sys.exit(0) diff --git a/tutorials/basic/.calibanconfig.json b/tutorials/basic/.calibanconfig.json new file mode 100644 index 0000000..7a62321 --- /dev/null +++ b/tutorials/basic/.calibanconfig.json @@ -0,0 +1 @@ +{"apt_packages": "cake"} From c1f9aa8b8c1205c976ddfacbd7c046fc06d356fa Mon Sep 17 00:00:00 2001 From: Sam Ritchie Date: Tue, 14 Jul 2020 21:18:36 -0600 Subject: [PATCH 12/15] convert argparse to schema --- caliban/cli.py | 5 +- caliban/config/__init__.py | 57 +++++--------------- caliban/config/experiment.py | 7 ++- caliban/main.py | 3 +- caliban/util/argparse.py | 42 +++++++-------- caliban/util/schema.py | 81 +++++++++++++++++++++++++++++ tutorials/basic/.calibanconfig.json | 2 +- 7 files changed, 125 insertions(+), 72 deletions(-) create mode 100644 caliban/util/schema.py diff --git a/caliban/cli.py b/caliban/cli.py index 8a180cc..5ea0dac 100644 --- a/caliban/cli.py +++ b/caliban/cli.py @@ -33,6 +33,7 @@ import caliban.platform.gke.util as gke_u import caliban.util as u import caliban.util.argparse as ua +import caliban.util.schema as us from caliban import __version__ @@ -156,7 +157,7 @@ def extra_dirs(parser): "-d", "--dir", action="append", - type=ua.validated_directory, + type=ua.argparse_schema(us.Directory), help="Extra directories to include. List these from large to small " "to take full advantage of Docker's build cache.") @@ -188,7 +189,7 @@ def region_arg(parser): def cloud_key_arg(parser): parser.add_argument("--cloud_key", - type=ua.validated_file, + type=ua.argparse_schema(us.File), help="Path to GCloud service account key. " "(Defaults to $GOOGLE_APPLICATION_CREDENTIALS.)") diff --git a/caliban/config/__init__.py b/caliban/config/__init__.py index 4ce9160..baa5755 100644 --- a/caliban/config/__init__.py +++ b/caliban/config/__init__.py @@ -20,7 +20,6 @@ import argparse import os import sys -from contextlib import contextmanager from enum import Enum from typing import Any, Dict, List, Optional @@ -29,7 +28,7 @@ import yaml import caliban.platform.cloud.types as ct -import caliban.util as u +import caliban.util.schema as us class JobMode(str, Enum): @@ -39,13 +38,9 @@ class JobMode(str, Enum): @staticmethod def parse(label): - return JobMode(label.upper()) -# Special config for Caliban. -CalibanConfig = Dict[str, Any] - DRY_RUN_FLAG = "--dry_run" CALIBAN_CONFIG = ".calibanconfig.json" @@ -68,14 +63,14 @@ def parse(label): # Schema for Caliban Config -AptPackages = s.Schema( - s.Or( - [str], { - s.Optional("gpu", default=list): [str], - s.Optional("cpu", default=list): [str] - })) +AptPackages = s.Or( + [str], { + s.Optional("gpu", default=list): [str], + s.Optional("cpu", default=list): [str] + }, + error=""""apt_packages" entry must be a dictionary or list, not '{}'""") -CCSchema = s.Schema({ +CalibanConfig = s.Schema({ s.Optional("build_time_credentials", default=False): bool, s.Optional("default_mode", default=JobMode.CPU): @@ -210,44 +205,20 @@ def apt_packages(conf: CalibanConfig, mode: JobMode) -> List[str]: k = "gpu" if gpu(mode) else "cpu" return packages.get(k, []) - elif isinstance(packages, list): - return packages - - else: - raise argparse.ArgumentTypeError( - """{}'s "apt_packages" entry must be a dictionary or list, not '{}'""". - format(CALIBAN_CONFIG, packages)) + return packages def caliban_config(conf_path: str = CALIBAN_CONFIG) -> CalibanConfig: """Returns a dict that represents a `.calibanconfig.json` file if present, empty dictionary otherwise. + If the supplied conf_path is present, but doesn't pass the supplied schema, + errors and kills the program. + """ if not os.path.isfile(conf_path): return {} conf = load_config(conf_path, mode='json') - return CCSchema.validate(conf) - - -@contextmanager -def argparse_schema(*args, **kwds): - """This function should work as a context manager that will trap a SchemaError - and return an argparse.ArgumentTypeError instead. - - TODO move to util.argparse - - """ - try: - yield - except s.SchemaError as e: - raise argparse.ArgumentTypeError(e.code) from None - - -@contextmanager -def error_schema(*args, **kwds): - try: - yield - except s.SchemaError as e: - u.err(e.code) + with us.error_schema(f"{conf_path}"): + return CalibanConfig.validate(conf) diff --git a/caliban/config/experiment.py b/caliban/config/experiment.py index c0fecd5..b4352bc 100644 --- a/caliban/config/experiment.py +++ b/caliban/config/experiment.py @@ -29,6 +29,8 @@ import commentjson import caliban.util as u +import caliban.util.argparse as ua +import caliban.util.schema as us # int, str and bool are allowed in a final experiment; lists are markers for # expansion. @@ -266,8 +268,9 @@ def load_experiment_config(s): if s.lower() == 'stdin': json = commentjson.load(sys.stdin) else: - with open(u.validated_file(s)) as f: - json = commentjson.load(f) + with ua.argparse_schema(): + with open(us.File.validate(s)) as f: + json = commentjson.load(f) return validate_experiment_config(json) diff --git a/caliban/main.py b/caliban/main.py index 07df504..3dada18 100644 --- a/caliban/main.py +++ b/caliban/main.py @@ -34,6 +34,7 @@ import caliban.platform.notebook as pn import caliban.platform.run as pr import caliban.platform.shell as ps +import caliban.util.schema as cs ll.getLogger('caliban.main').setLevel(logging.ERROR) t = Terminal() @@ -154,7 +155,7 @@ def run_app(arg_input): def main(): logging.use_python_logging() try: - with c.argparse_schema(): + with cs.fatal_errors(): app.run(run_app, flags_parser=cli.parse_flags) except KeyboardInterrupt: logging.info('Shutting down.') diff --git a/caliban/util/argparse.py b/caliban/util/argparse.py index efd6e34..bda9f28 100644 --- a/caliban/util/argparse.py +++ b/caliban/util/argparse.py @@ -25,6 +25,7 @@ import caliban.util as u import caliban.util.fs as ufs +import schema as s t = Terminal() @@ -39,6 +40,24 @@ def expand_args(items: Dict[str, str]) -> List[str]: return list(it.chain.from_iterable(pairs)) +def argparse_schema(schema): + """Wrapper that performs validation and converts SchemaErrors into + ArgumentTypeErrors for better argument error reporting. + + """ + + def check(x): + try: + return schema.validate(x) + except s.SchemaError as e: + raise argparse.ArgumentTypeError(e.code) from None + + return check + + +# TODO: Now that we use schema, these should all be converted to schema instances. + + def validated_package(path: str) -> u.Package: """similar to generate_package but runs argparse validation on packages that don't actually exist in the filesystem. @@ -89,26 +108,3 @@ def is_key(k: Optional[str]) -> bool: """ return k is not None and len(k) > 0 and k[0] == "-" - - -def validated_directory(path: str) -> str: - """This validates that the supplied directory exists locally. - - """ - if not os.path.isdir(path): - raise argparse.ArgumentTypeError( - """Directory '{}' doesn't exist in this directory. Check yourself!""". - format(path)) - return path - - -def validated_file(path: str) -> str: - """This validates that the supplied file exists. Tilde expansion is supported. - - """ - expanded = os.path.expanduser(path) - if not os.path.isfile(expanded): - raise argparse.ArgumentTypeError( - """File '{}' isn't a valid file on your system. Try again!""".format( - path)) - return path diff --git a/caliban/util/schema.py b/caliban/util/schema.py new file mode 100644 index 0000000..00d434a --- /dev/null +++ b/caliban/util/schema.py @@ -0,0 +1,81 @@ +#!/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. +""" +Useful shared schemas. +""" +import os +import sys +from contextlib import contextmanager +from typing import Any, Dict, Optional + +import caliban.util as u +import schema as s + + +class FatalSchemaError(Exception): + """Wrapper for an exception that can bubble itself up to the top level of the +program.""" + + def __init__(self, message, context): + self.message = message + self.context = context + super().__init__(self.message) + + +@contextmanager +def error_schema(context: Optional[str] = None): + """Wrap functions that check schemas in this context manager to throw an + appropriate error with a nice message. + + """ + prefix = "" + if context is not None: + prefix = f"\nValidation error while parsing {context}:\n" + + try: + yield + except s.SchemaError as e: + raise FatalSchemaError(e.code, prefix) + + +@contextmanager +def fatal_errors(): + """Context manager meant to wrap an entire program and present schema errors in + an easy-to-read way. + + """ + try: + yield + except FatalSchemaError as e: + u.err(f"{e.context}\n{e.message}\n\n") + sys.exit(1) + except s.SchemaError as e: + u.err(f"\n{e.code}\n\n") + sys.exit(1) + + +# TODO Once a release with this patch happens: +# https://github.com/keleshev/schema/pull/238,, Change `Or` to `Schema`. This +# problem only occurs for callable validators. + +Directory = s.Or( + os.path.isdir, + False, + error="""Directory '{}' doesn't exist in this directory. Check yourself!""") + +File = s.Or(lambda path: os.path.isfile(os.path.expanduser(path)), + False, + error="""File '{}' isn't a valid file on your system. Try again!""") diff --git a/tutorials/basic/.calibanconfig.json b/tutorials/basic/.calibanconfig.json index 7a62321..e324ca9 100644 --- a/tutorials/basic/.calibanconfig.json +++ b/tutorials/basic/.calibanconfig.json @@ -1 +1 @@ -{"apt_packages": "cake"} +{"apt_packages": ["cake"]} From baa3eeb4f5d6f2cd1389c871c4c50e03aaec8ab3 Mon Sep 17 00:00:00 2001 From: Sam Ritchie Date: Tue, 14 Jul 2020 21:42:57 -0600 Subject: [PATCH 13/15] remove unused functions --- caliban/config/__init__.py | 44 ++++---------------------------------- caliban/util/argparse.py | 3 ++- caliban/util/schema.py | 14 +++++++++++- 3 files changed, 19 insertions(+), 42 deletions(-) diff --git a/caliban/config/__init__.py b/caliban/config/__init__.py index baa5755..a010e84 100644 --- a/caliban/config/__init__.py +++ b/caliban/config/__init__.py @@ -100,41 +100,6 @@ def gpu(job_mode: JobMode) -> bool: return job_mode == JobMode.GPU -def load_yaml_config(path): - """returns the config parsed based on the info in the flags. - - Grabs the config file, written in yaml, slurps it in. - """ - with open(path) as f: - config = yaml.load(f, Loader=yaml.FullLoader) - - return config - - -def load_config(path, mode='yaml'): - """Load a JSON or YAML config. - - """ - if mode == 'json': - with open(path) as f: - return commentjson.load(f) - - return load_yaml_config(path) - - -def valid_json(path: str) -> Dict[str, Any]: - """Loads JSON if the path points to a valid JSON file; otherwise, throws an - exception that's picked up by argparse. - - """ - try: - return load_config(path, mode='json') - except commentjson.JSONLibraryException: - raise argparse.ArgumentTypeError( - """File '{}' doesn't seem to contain valid JSON. Try again!""".format( - path)) - - def extract_script_args(m: Dict[str, Any]) -> List[str]: """Strip off the "--" argument if it was passed in as a separator.""" script_args = m.get("script_args") @@ -199,11 +164,11 @@ def apt_packages(conf: CalibanConfig, mode: JobMode) -> List[str]: the requests in the config. """ - packages = conf.get("apt_packages") or {} + packages = conf["apt_packages"] if isinstance(packages, dict): k = "gpu" if gpu(mode) else "cpu" - return packages.get(k, []) + return packages.get[k] return packages @@ -219,6 +184,5 @@ def caliban_config(conf_path: str = CALIBAN_CONFIG) -> CalibanConfig: if not os.path.isfile(conf_path): return {} - conf = load_config(conf_path, mode='json') - with us.error_schema(f"{conf_path}"): - return CalibanConfig.validate(conf) + with us.error_schema(conf_path): + return s.And(us.Json, CalibanConfig).validate(conf_path) diff --git a/caliban/util/argparse.py b/caliban/util/argparse.py index bda9f28..79f24bc 100644 --- a/caliban/util/argparse.py +++ b/caliban/util/argparse.py @@ -55,7 +55,8 @@ def check(x): return check -# TODO: Now that we use schema, these should all be converted to schema instances. +# TODO: Now that we use schema, validated_package and parse_kv_pair should be +# converted to schema instances. def validated_package(path: str) -> u.Package: diff --git a/caliban/util/schema.py b/caliban/util/schema.py index 00d434a..c15b839 100644 --- a/caliban/util/schema.py +++ b/caliban/util/schema.py @@ -19,7 +19,9 @@ import os import sys from contextlib import contextmanager -from typing import Any, Dict, Optional +from typing import Optional + +import commentjson import caliban.util as u import schema as s @@ -67,6 +69,11 @@ def fatal_errors(): sys.exit(1) +def load_json(path): + with open(path) as f: + return commentjson.load(f) + + # TODO Once a release with this patch happens: # https://github.com/keleshev/schema/pull/238,, Change `Or` to `Schema`. This # problem only occurs for callable validators. @@ -79,3 +86,8 @@ def fatal_errors(): File = s.Or(lambda path: os.path.isfile(os.path.expanduser(path)), False, error="""File '{}' isn't a valid file on your system. Try again!""") + +Json = s.And( + File, + s.Use(load_json, + error="""File '{}' doesn't seem to contain valid JSON. Try again!""")) From bf992326c2c80615384ce12e5972c0d264f0a277 Mon Sep 17 00:00:00 2001 From: Sam Ritchie Date: Tue, 14 Jul 2020 22:25:12 -0600 Subject: [PATCH 14/15] calibanconfig --- caliban/config/__init__.py | 2 +- caliban/config/experiment.py | 6 ++-- tests/caliban/config/test_config.py | 40 +++++++++++++++++++++++- tests/caliban/config/test_experiment.py | 26 ++++++++++++++++ tests/caliban/util/test_schema.py | 41 +++++++++++++++++++++++++ 5 files changed, 109 insertions(+), 6 deletions(-) create mode 100644 tests/caliban/util/test_schema.py diff --git a/caliban/config/__init__.py b/caliban/config/__init__.py index a010e84..0e0c6b5 100644 --- a/caliban/config/__init__.py +++ b/caliban/config/__init__.py @@ -168,7 +168,7 @@ def apt_packages(conf: CalibanConfig, mode: JobMode) -> List[str]: if isinstance(packages, dict): k = "gpu" if gpu(mode) else "cpu" - return packages.get[k] + return packages[k] return packages diff --git a/caliban/config/experiment.py b/caliban/config/experiment.py index b4352bc..4305d2d 100644 --- a/caliban/config/experiment.py +++ b/caliban/config/experiment.py @@ -265,12 +265,10 @@ def validate_experiment_config(items: ExpConf) -> ExpConf: def load_experiment_config(s): - if s.lower() == 'stdin': + if isinstance(s, str) and s.lower() == 'stdin': json = commentjson.load(sys.stdin) else: - with ua.argparse_schema(): - with open(us.File.validate(s)) as f: - json = commentjson.load(f) + json = ua.argparse_schema(us.Json)(s) return validate_experiment_config(json) diff --git a/tests/caliban/config/test_config.py b/tests/caliban/config/test_config.py index 126c21d..3f94229 100644 --- a/tests/caliban/config/test_config.py +++ b/tests/caliban/config/test_config.py @@ -14,14 +14,22 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json import os from argparse import ArgumentTypeError -import caliban.platform.cloud.types as ct import caliban.config as c +import caliban.platform.cloud.types as ct +import caliban.util.schema as us import pytest +def test_gpu(): + assert c.gpu(c.JobMode.GPU) + assert not c.gpu(c.JobMode.CPU) + assert not c.gpu("face") + + def test_extract_region(monkeypatch): if os.environ.get('REGION'): monkeypatch.delenv('REGION') @@ -42,3 +50,33 @@ def test_extract_region(monkeypatch): assert c.extract_region({}) == c.DEFAULT_REGION assert c.extract_region({"region": "us-west1"}) == ct.US.west1 + + +def test_caliban_config(tmpdir): + """Tests validation of the CalibanConfig schema and the method that returns the + parsed config. + + """ + valid = {"apt_packages": {"cpu": ["face"]}, "random": "entry"} + valid_path = tmpdir.join('valid.json') + + with open(valid_path, 'w') as f: + json.dump(valid, f) + + invalid = {"apt_packages": "face"} + invalid_path = tmpdir.join('invalid.json') + + with open(invalid_path, 'w') as f: + json.dump(invalid, f) + + # Failing the schema raises an error. + with pytest.raises(us.FatalSchemaError): + c.caliban_config(invalid_path) + + # paths that don't exist return an empty map: + assert c.caliban_config('random_path') == {} + + # If the config is valid, c.apt_packages can fetch the packages we specified. + config = c.caliban_config(valid_path) + assert c.apt_packages(config, c.JobMode.GPU) == [] + assert c.apt_packages(config, c.JobMode.CPU) == ["face"] diff --git a/tests/caliban/config/test_experiment.py b/tests/caliban/config/test_experiment.py index 5ca1e3d..ad4860c 100644 --- a/tests/caliban/config/test_experiment.py +++ b/tests/caliban/config/test_experiment.py @@ -14,10 +14,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json from argparse import ArgumentTypeError import caliban.config.experiment as c import caliban.util as u +import caliban.util.schema as us import pytest @@ -238,3 +240,27 @@ def test_compound_key_handling(): assert test['after_expansion'] == c.expand_experiment_config(test['input']) assert test['after_dictproduct'] == list( u.dict_product(test['after_tupleization'])) + + +def test_load_experiment_config(tmpdir): + valid = {"key": ['a', 'b'], "random": [True, False]} + valid_path = tmpdir.join('valid.json') + + with open(valid_path, 'w') as f: + json.dump(valid, f) + + invalid_path = tmpdir.join('invalid.json') + + with open(invalid_path, 'w') as f: + f.write("{{{I am not JSON!\n") + + # Failing the schema with invalid json raises an ARGPARSE error, not a schema + # error. We haven't converted experimentconfig to schema yet. + # + # We use schema to validate, but the ua.argparse_schema wrapper converts the + # error internally. + with pytest.raises(ArgumentTypeError): + c.load_experiment_config(invalid_path) + + # A valid config should round trip. + assert c.load_experiment_config(valid_path) == valid diff --git a/tests/caliban/util/test_schema.py b/tests/caliban/util/test_schema.py new file mode 100644 index 0000000..af9bb56 --- /dev/null +++ b/tests/caliban/util/test_schema.py @@ -0,0 +1,41 @@ +#!/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. + +import tempfile + +import schema as s + +import caliban.util.schema as us +import pytest + + +def test_directory(tmpdir): + # Proper directories pass validation. + assert us.Directory.validate(tmpdir) == tmpdir + + # random dirs that I made up dont! + with pytest.raises(s.SchemaError): + assert us.Directory.validate('random') + + +def test_file(): + with tempfile.NamedTemporaryFile() as tmp: + # Existing files pass validation. + assert us.File.validate(tmp.name) == tmp.name + + # random paths that I made up dont! + with pytest.raises(s.SchemaError): + assert us.File.validate('random') From d67ae412b062d53187a4257cdf4c69624f0adc09 Mon Sep 17 00:00:00 2001 From: Sam Ritchie Date: Tue, 14 Jul 2020 22:52:40 -0600 Subject: [PATCH 15/15] remove unused --- caliban/config/__init__.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/caliban/config/__init__.py b/caliban/config/__init__.py index 0e0c6b5..f85340c 100644 --- a/caliban/config/__init__.py +++ b/caliban/config/__init__.py @@ -17,15 +17,12 @@ Utilities for our job runner, for working with configs. """ -import argparse import os import sys from enum import Enum from typing import Any, Dict, List, Optional -import commentjson import schema as s -import yaml import caliban.platform.cloud.types as ct import caliban.util.schema as us