| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you 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 re |
| import subprocess |
| from io import StringIO |
| |
| from dotenv import dotenv_values |
| from ruamel.yaml import YAML |
| |
| from .utils.command import Command, default_bin |
| from .compat import _ensure_path |
| |
| |
| def flatten(node, parents=None): |
| parents = list(parents or []) |
| if isinstance(node, str): |
| yield (node, parents) |
| elif isinstance(node, list): |
| for value in node: |
| yield from flatten(value, parents=parents) |
| elif isinstance(node, dict): |
| for key, value in node.items(): |
| yield (key, parents) |
| yield from flatten(value, parents=parents + [key]) |
| else: |
| raise TypeError(node) |
| |
| |
| def _sanitize_command(cmd): |
| if isinstance(cmd, list): |
| cmd = " ".join(cmd) |
| return re.sub(r"\s+", " ", cmd) |
| |
| |
| class UndefinedImage(Exception): |
| pass |
| |
| |
| class ComposeConfig: |
| |
| def __init__(self, config_path, dotenv_path, compose_bin, params=None): |
| config_path = _ensure_path(config_path) |
| if dotenv_path: |
| dotenv_path = _ensure_path(dotenv_path) |
| else: |
| dotenv_path = config_path.parent / '.env' |
| self._read_env(dotenv_path, params) |
| self._read_config(config_path, compose_bin) |
| |
| def _read_env(self, dotenv_path, params): |
| """ |
| Read .env and merge it with explicitly passed parameters. |
| """ |
| self.dotenv = dotenv_values(str(dotenv_path)) |
| if params is None: |
| self.params = {} |
| else: |
| self.params = {k: v for k, v in params.items() if k in self.dotenv} |
| |
| # forward the process' environment variables |
| self.env = os.environ.copy() |
| # set the defaults from the dotenv files |
| self.env.update(self.dotenv) |
| # override the defaults passed as parameters |
| self.env.update(self.params) |
| |
| # translate docker's architecture notation to a more widely used one |
| arch = self.env.get('ARCH', 'amd64') |
| arch_aliases = { |
| 'amd64': 'x86_64', |
| 'arm64v8': 'aarch64', |
| 's390x': 's390x' |
| } |
| arch_short_aliases = { |
| 'amd64': 'x64', |
| 'arm64v8': 'arm64', |
| 's390x': 's390x' |
| } |
| self.env['ARCH_ALIAS'] = arch_aliases.get(arch, arch) |
| self.env['ARCH_SHORT_ALIAS'] = arch_short_aliases.get(arch, arch) |
| |
| def _read_config(self, config_path, compose_bin): |
| """ |
| Validate and read the docker-compose.yml |
| """ |
| yaml = YAML() |
| with config_path.open() as fp: |
| config = yaml.load(fp) |
| |
| services = config['services'].keys() |
| self.hierarchy = dict(flatten(config.get('x-hierarchy', {}))) |
| self.with_gpus = config.get('x-with-gpus', []) |
| nodes = self.hierarchy.keys() |
| errors = [] |
| |
| for name in self.with_gpus: |
| if name not in services: |
| errors.append( |
| 'Service `{}` defined in `x-with-gpus` bot not in ' |
| '`services`'.format(name) |
| ) |
| for name in nodes - services: |
| errors.append( |
| 'Service `{}` is defined in `x-hierarchy` bot not in ' |
| '`services`'.format(name) |
| ) |
| for name in services - nodes: |
| errors.append( |
| 'Service `{}` is defined in `services` but not in ' |
| '`x-hierarchy`'.format(name) |
| ) |
| |
| # trigger docker-compose's own validation |
| compose = Command('docker-compose') |
| args = ['--file', str(config_path), 'config'] |
| result = compose.run(*args, env=self.env, check=False, |
| stderr=subprocess.PIPE, stdout=subprocess.PIPE) |
| |
| if result.returncode != 0: |
| # strip the intro line of docker-compose errors |
| errors += result.stderr.decode().splitlines() |
| |
| if errors: |
| msg = '\n'.join([' - {}'.format(msg) for msg in errors]) |
| raise ValueError( |
| 'Found errors with docker-compose:\n{}'.format(msg) |
| ) |
| |
| rendered_config = StringIO(result.stdout.decode()) |
| self.path = config_path |
| self.config = yaml.load(rendered_config) |
| |
| def get(self, service_name): |
| try: |
| service = self.config['services'][service_name] |
| except KeyError: |
| raise UndefinedImage(service_name) |
| service['name'] = service_name |
| service['need_gpu'] = service_name in self.with_gpus |
| service['ancestors'] = self.hierarchy[service_name] |
| return service |
| |
| def __getitem__(self, service_name): |
| return self.get(service_name) |
| |
| |
| class Docker(Command): |
| |
| def __init__(self, docker_bin=None): |
| self.bin = default_bin(docker_bin, "docker") |
| |
| |
| class DockerCompose(Command): |
| |
| def __init__(self, config_path, dotenv_path=None, compose_bin=None, |
| params=None): |
| compose_bin = default_bin(compose_bin, 'docker-compose') |
| self.config = ComposeConfig(config_path, dotenv_path, compose_bin, |
| params) |
| self.bin = compose_bin |
| self.pull_memory = set() |
| |
| def clear_pull_memory(self): |
| self.pull_memory = set() |
| |
| def _execute_compose(self, *args, **kwargs): |
| # execute as a docker compose command |
| try: |
| result = super().run('--file', str(self.config.path), *args, |
| env=self.config.env, **kwargs) |
| result.check_returncode() |
| except subprocess.CalledProcessError as e: |
| def formatdict(d, template): |
| return '\n'.join( |
| template.format(k, v) for k, v in sorted(d.items()) |
| ) |
| msg = ( |
| "`{cmd}` exited with a non-zero exit code {code}, see the " |
| "process log above.\n\nThe docker-compose command was " |
| "invoked with the following parameters:\n\nDefaults defined " |
| "in .env:\n{dotenv}\n\nArchery was called with:\n{params}" |
| ) |
| raise RuntimeError( |
| msg.format( |
| cmd=' '.join(e.cmd), |
| code=e.returncode, |
| dotenv=formatdict(self.config.dotenv, template=' {}: {}'), |
| params=formatdict( |
| self.config.params, template=' export {}={}' |
| ) |
| ) |
| ) |
| |
| def _execute_docker(self, *args, **kwargs): |
| # execute as a plain docker cli command |
| try: |
| result = Docker().run(*args, **kwargs) |
| result.check_returncode() |
| except subprocess.CalledProcessError as e: |
| raise RuntimeError( |
| "{} exited with non-zero exit code {}".format( |
| ' '.join(e.cmd), e.returncode |
| ) |
| ) |
| |
| def pull(self, service_name, pull_leaf=True, using_docker=False): |
| def _pull(service): |
| args = ['pull'] |
| if service['image'] in self.pull_memory: |
| return |
| |
| if using_docker: |
| try: |
| self._execute_docker(*args, service['image']) |
| except Exception as e: |
| # better --ignore-pull-failures handling |
| print(e) |
| else: |
| args.append('--ignore-pull-failures') |
| self._execute_compose(*args, service['name']) |
| |
| self.pull_memory.add(service['image']) |
| |
| service = self.config.get(service_name) |
| for ancestor in service['ancestors']: |
| _pull(self.config.get(ancestor)) |
| if pull_leaf: |
| _pull(service) |
| |
| def build(self, service_name, use_cache=True, use_leaf_cache=True, |
| using_docker=False, using_buildx=False): |
| def _build(service, use_cache): |
| if 'build' not in service: |
| # nothing to do |
| return |
| |
| args = [] |
| cache_from = list(service.get('build', {}).get('cache_from', [])) |
| if use_cache: |
| for image in cache_from: |
| if image not in self.pull_memory: |
| try: |
| self._execute_docker('pull', image) |
| except Exception as e: |
| print(e) |
| finally: |
| self.pull_memory.add(image) |
| else: |
| args.append('--no-cache') |
| |
| # turn on inline build cache, this is a docker buildx feature |
| # used to bundle the image build cache to the pushed image manifest |
| # so the build cache can be reused across hosts, documented at |
| # https://github.com/docker/buildx#--cache-tonametypetypekeyvalue |
| if self.config.env.get('BUILDKIT_INLINE_CACHE') == '1': |
| args.extend(['--build-arg', 'BUILDKIT_INLINE_CACHE=1']) |
| |
| if using_buildx: |
| for k, v in service['build'].get('args', {}).items(): |
| args.extend(['--build-arg', '{}={}'.format(k, v)]) |
| |
| if use_cache: |
| cache_ref = '{}-cache'.format(service['image']) |
| cache_from = 'type=registry,ref={}'.format(cache_ref) |
| cache_to = ( |
| 'type=registry,ref={},mode=max'.format(cache_ref) |
| ) |
| args.extend([ |
| '--cache-from', cache_from, |
| '--cache-to', cache_to, |
| ]) |
| |
| args.extend([ |
| '--output', 'type=docker', |
| '-f', service['build']['dockerfile'], |
| '-t', service['image'], |
| service['build'].get('context', '.') |
| ]) |
| self._execute_docker("buildx", "build", *args) |
| elif using_docker: |
| # better for caching |
| for k, v in service['build'].get('args', {}).items(): |
| args.extend(['--build-arg', '{}={}'.format(k, v)]) |
| for img in cache_from: |
| args.append('--cache-from="{}"'.format(img)) |
| args.extend([ |
| '-f', service['build']['dockerfile'], |
| '-t', service['image'], |
| service['build'].get('context', '.') |
| ]) |
| self._execute_docker("build", *args) |
| else: |
| self._execute_compose("build", *args, service['name']) |
| |
| service = self.config.get(service_name) |
| # build ancestor services |
| for ancestor in service['ancestors']: |
| _build(self.config.get(ancestor), use_cache=use_cache) |
| # build the leaf/target service |
| _build(service, use_cache=use_cache and use_leaf_cache) |
| |
| def run(self, service_name, command=None, *, env=None, volumes=None, |
| user=None, using_docker=False): |
| service = self.config.get(service_name) |
| |
| args = [] |
| if user is not None: |
| args.extend(['-u', user]) |
| |
| if env is not None: |
| for k, v in env.items(): |
| args.extend(['-e', '{}={}'.format(k, v)]) |
| |
| if volumes is not None: |
| for volume in volumes: |
| args.extend(['--volume', volume]) |
| |
| if using_docker or service['need_gpu']: |
| # use gpus, requires docker>=19.03 |
| if service['need_gpu']: |
| args.extend(['--gpus', 'all']) |
| |
| if service.get('shm_size'): |
| args.extend(['--shm-size', service['shm_size']]) |
| |
| # append env variables from the compose conf |
| for k, v in service.get('environment', {}).items(): |
| args.extend(['-e', '{}={}'.format(k, v)]) |
| |
| # append volumes from the compose conf |
| for v in service.get('volumes', []): |
| if not isinstance(v, str): |
| # if not the compact string volume definition |
| v = "{}:{}".format(v['source'], v['target']) |
| args.extend(['-v', v]) |
| |
| # infer whether an interactive shell is desired or not |
| if command in ['cmd.exe', 'bash', 'sh', 'powershell']: |
| args.append('-it') |
| |
| # get the actual docker image name instead of the compose service |
| # name which we refer as image in general |
| args.append(service['image']) |
| |
| # add command from compose if it wasn't overridden |
| if command is not None: |
| args.append(command) |
| else: |
| # replace whitespaces from the preformatted compose command |
| cmd = _sanitize_command(service.get('command', '')) |
| if cmd: |
| args.append(cmd) |
| |
| # execute as a plain docker cli command |
| self._execute_docker('run', '--rm', *args) |
| else: |
| # execute as a docker-compose command |
| args.append(service_name) |
| if command is not None: |
| args.append(command) |
| self._execute_compose('run', '--rm', *args) |
| |
| def push(self, service_name, user=None, password=None, using_docker=False): |
| def _push(service): |
| if using_docker: |
| return self._execute_docker('push', service['image']) |
| else: |
| return self._execute_compose('push', service['name']) |
| |
| if user is not None: |
| try: |
| # TODO(kszucs): have an option for a prompt |
| self._execute_docker('login', '-u', user, '-p', password) |
| except subprocess.CalledProcessError: |
| # hide credentials |
| msg = ('Failed to push `{}`, check the passed credentials' |
| .format(service_name)) |
| raise RuntimeError(msg) from None |
| |
| service = self.config.get(service_name) |
| for ancestor in service['ancestors']: |
| _push(self.config.get(ancestor)) |
| _push(service) |
| |
| def images(self): |
| return sorted(self.config.hierarchy.keys()) |
