def start_containers(client: docker.APIClient): configs = tables('docker').data images = ['ubuntu', 'alpine', 'nginx'] ports_delta = 1 for image in images: base_config = { "image": image, "command": "sleep 1d", "detach": True} for conf in configs: if conf.startswith('vol'): if conf == 'vol1' and image != 'alpine': container = client.create_container( host_config=client.create_host_config(binds=configs[conf]), image=image, command=COMMAND, detach=True) else: container = client.create_container( host_config=client.create_host_config(binds=configs[conf]), **base_config) elif conf.startswith('ports'): ports = {} for p in range(configs[conf]): ports.update({9980 + ports_delta: 9980 + ports_delta}) ports.update({str(9981 + ports_delta) + '/udp': 9985 + ports_delta}) ports_delta += 1 container = client.create_container( host_config=client.create_host_config(port_bindings=ports), ports=[*ports], **base_config) elif conf.startswith('labels'): container = client.create_container( labels=configs[conf], **base_config) elif conf == 'privileged': container = client.create_container( host_config=client.create_host_config(privileged=configs[conf]), **base_config) else: entry_config = copy.copy(base_config) entry_config.pop('command') container = client.create_container( entrypoint=configs[conf], **entry_config) client.start(container)
class DockerOperator(BaseOperator): """ Execute a command inside a docker container. A temporary directory is created on the host and mounted into a container to allow storing files that together exceed the default disk size of 10GB in a container. The path to the mounted directory can be accessed via the environment variable ``AIRFLOW_TMP_DIR``. If a login to a private registry is required prior to pulling the image, a Docker connection needs to be configured in Airflow and the connection ID be provided with the parameter ``docker_conn_id``. :param image: Docker image from which to create the container. If image tag is omitted, "latest" will be used. :type image: str :param api_version: Remote API version. Set to ``auto`` to automatically detect the server's version. :type api_version: str :param auto_remove: Auto-removal of the container on daemon side when the container's process exits. The default is False. :type auto_remove: bool :param command: Command to be run in the container. (templated) :type command: str or list :param cpus: Number of CPUs to assign to the container. This value gets multiplied with 1024. See https://docs.docker.com/engine/reference/run/#cpu-share-constraint :type cpus: float :param dns: Docker custom DNS servers :type dns: list[str] :param dns_search: Docker custom DNS search domain :type dns_search: list[str] :param docker_url: URL of the host running the docker daemon. Default is unix://var/run/docker.sock :type docker_url: str :param environment: Environment variables to set in the container. (templated) :type environment: dict :param force_pull: Pull the docker image on every run. Default is False. :type force_pull: bool :param mem_limit: Maximum amount of memory the container can use. Either a float value, which represents the limit in bytes, or a string like ``128m`` or ``1g``. :type mem_limit: float or str :param network_mode: Network mode for the container. :type network_mode: str :param tls_ca_cert: Path to a PEM-encoded certificate authority to secure the docker connection. :type tls_ca_cert: str :param tls_client_cert: Path to the PEM-encoded certificate used to authenticate docker client. :type tls_client_cert: str :param tls_client_key: Path to the PEM-encoded key used to authenticate docker client. :type tls_client_key: str :param tls_hostname: Hostname to match against the docker server certificate or False to disable the check. :type tls_hostname: str or bool :param tls_ssl_version: Version of SSL to use when communicating with docker daemon. :type tls_ssl_version: str :param tmp_dir: Mount point inside the container to a temporary directory created on the host by the operator. The path is also made available via the environment variable ``AIRFLOW_TMP_DIR`` inside the container. :type tmp_dir: str :param user: Default user inside the docker container. :type user: int or str :param volumes: List of volumes to mount into the container, e.g. ``['/host/path:/container/path', '/host/path2:/container/path2:ro']``. :type volumes: list :param working_dir: Working directory to set on the container (equivalent to the -w switch the docker client) :type working_dir: str :param xcom_all: Push all the stdout or just the last line. The default is False (last line). :type xcom_all: bool :param docker_conn_id: ID of the Airflow connection to use :type docker_conn_id: str :param shm_size: Size of ``/dev/shm`` in bytes. The size must be greater than 0. If omitted uses system default. :type shm_size: int """ template_fields = ('command', 'environment',) template_ext = ('.sh', '.bash',) @apply_defaults def __init__( self, image, api_version=None, command=None, cpus=1.0, docker_url='unix://var/run/docker.sock', environment=None, force_pull=False, mem_limit=None, network_mode=None, tls_ca_cert=None, tls_client_cert=None, tls_client_key=None, tls_hostname=None, tls_ssl_version=None, tmp_dir='/tmp/airflow', user=None, volumes=None, working_dir=None, xcom_all=False, docker_conn_id=None, dns=None, dns_search=None, auto_remove=False, shm_size=None, *args, **kwargs): super().__init__(*args, **kwargs) self.api_version = api_version self.auto_remove = auto_remove self.command = command self.cpus = cpus self.dns = dns self.dns_search = dns_search self.docker_url = docker_url self.environment = environment or {} self.force_pull = force_pull self.image = image self.mem_limit = mem_limit self.network_mode = network_mode self.tls_ca_cert = tls_ca_cert self.tls_client_cert = tls_client_cert self.tls_client_key = tls_client_key self.tls_hostname = tls_hostname self.tls_ssl_version = tls_ssl_version self.tmp_dir = tmp_dir self.user = user self.volumes = volumes or [] self.working_dir = working_dir self.xcom_all = xcom_all self.docker_conn_id = docker_conn_id self.shm_size = shm_size if kwargs.get('xcom_push') is not None: raise AirflowException("'xcom_push' was deprecated, use 'BaseOperator.do_xcom_push' instead") self.cli = None self.container = None def get_hook(self): return DockerHook( docker_conn_id=self.docker_conn_id, base_url=self.docker_url, version=self.api_version, tls=self.__get_tls_config() ) def execute(self, context): self.log.info('Starting docker container from image %s', self.image) tls_config = self.__get_tls_config() if self.docker_conn_id: self.cli = self.get_hook().get_conn() else: self.cli = APIClient( base_url=self.docker_url, version=self.api_version, tls=tls_config ) if self.force_pull or len(self.cli.images(name=self.image)) == 0: self.log.info('Pulling docker image %s', self.image) for l in self.cli.pull(self.image, stream=True): output = json.loads(l.decode('utf-8').strip()) if 'status' in output: self.log.info("%s", output['status']) with TemporaryDirectory(prefix='airflowtmp') as host_tmp_dir: self.environment['AIRFLOW_TMP_DIR'] = self.tmp_dir self.volumes.append('{0}:{1}'.format(host_tmp_dir, self.tmp_dir)) self.container = self.cli.create_container( command=self.get_command(), environment=self.environment, host_config=self.cli.create_host_config( auto_remove=self.auto_remove, binds=self.volumes, network_mode=self.network_mode, shm_size=self.shm_size, dns=self.dns, dns_search=self.dns_search, cpu_shares=int(round(self.cpus * 1024)), mem_limit=self.mem_limit), image=self.image, user=self.user, working_dir=self.working_dir ) self.cli.start(self.container['Id']) line = '' for line in self.cli.logs(container=self.container['Id'], stream=True): line = line.strip() if hasattr(line, 'decode'): line = line.decode('utf-8') self.log.info(line) result = self.cli.wait(self.container['Id']) if result['StatusCode'] != 0: raise AirflowException('docker container failed: ' + repr(result)) # duplicated conditional logic because of expensive operation if self.do_xcom_push: return self.cli.logs(container=self.container['Id']) \ if self.xcom_all else line.encode('utf-8') def get_command(self): if self.command is not None and self.command.strip().find('[') == 0: commands = ast.literal_eval(self.command) else: commands = self.command return commands def on_kill(self): if self.cli is not None: self.log.info('Stopping docker container') self.cli.stop(self.container['Id']) def __get_tls_config(self): tls_config = None if self.tls_ca_cert and self.tls_client_cert and self.tls_client_key: tls_config = tls.TLSConfig( ca_cert=self.tls_ca_cert, client_cert=(self.tls_client_cert, self.tls_client_key), verify=True, ssl_version=self.tls_ssl_version, assert_hostname=self.tls_hostname ) self.docker_url = self.docker_url.replace('tcp://', 'https://') return tls_config
class DockerOperator(BaseOperator): """ Execute a command inside a docker container. A temporary directory is created on the host and mounted into a container to allow storing files that together exceed the default disk size of 10GB in a container. The path to the mounted directory can be accessed via the environment variable ``AIRFLOW_TMP_DIR``. :param image: Docker image from which to create the container. :type image: str :param api_version: Remote API version. :type api_version: str :param command: Command to be run in the container. :type command: str or list :param cpus: Number of CPUs to assign to the container. This value gets multiplied with 1024. See https://docs.docker.com/engine/reference/run/#cpu-share-constraint :type cpus: float :param docker_url: URL of the host running the docker daemon. :type docker_url: str :param environment: Environment variables to set in the container. :type environment: dict :param force_pull: Pull the docker image on every run. :type force_pull: bool :param mem_limit: Maximum amount of memory the container can use. Either a float value, which represents the limit in bytes, or a string like ``128m`` or ``1g``. :type mem_limit: float or str :param network_mode: Network mode for the container. :type network_mode: str :param tls_ca_cert: Path to a PEM-encoded certificate authority to secure the docker connection. :type tls_ca_cert: str :param tls_client_cert: Path to the PEM-encoded certificate used to authenticate docker client. :type tls_client_cert: str :param tls_client_key: Path to the PEM-encoded key used to authenticate docker client. :type tls_client_key: str :param tls_hostname: Hostname to match against the docker server certificate or False to disable the check. :type tls_hostname: str or bool :param tls_ssl_version: Version of SSL to use when communicating with docker daemon. :type tls_ssl_version: str :param tmp_dir: Mount point inside the container to a temporary directory created on the host by the operator. The path is also made available via the environment variable ``AIRFLOW_TMP_DIR`` inside the container. :type tmp_dir: str :param user: Default user inside the docker container. :type user: int or str :param volumes: List of volumes to mount into the container, e.g. ``['/host/path:/container/path', '/host/path2:/container/path2:ro']``. :param working_dir: Working directory to set on the container (equivalent to the -w switch the docker client) :type working_dir: str :param xcom_push: Does the stdout will be pushed to the next step using XCom. The default is False. :type xcom_push: bool :param xcom_all: Push all the stdout or just the last line. The default is False (last line). :type xcom_all: bool :param auto_remove: Automatically remove the container when it exits :type auto_remove: bool """ template_fields = ('command',) template_ext = ('.sh', '.bash',) @apply_defaults def __init__( self, image, api_version=None, command=None, cpus=1.0, docker_url='unix://var/run/docker.sock', environment=None, force_pull=False, mem_limit=None, network_mode=None, tls_ca_cert=None, tls_client_cert=None, tls_client_key=None, tls_hostname=None, tls_ssl_version=None, tmp_dir='/tmp/airflow', user=None, volumes=None, working_dir=None, xcom_push=False, xcom_all=False, auto_remove=False, *args, **kwargs): super(DockerOperator, self).__init__(*args, **kwargs) self.api_version = api_version self.command = command self.cpus = cpus self.docker_url = docker_url self.environment = environment or {} self.force_pull = force_pull self.image = image self.mem_limit = mem_limit self.network_mode = network_mode self.tls_ca_cert = tls_ca_cert self.tls_client_cert = tls_client_cert self.tls_client_key = tls_client_key self.tls_hostname = tls_hostname self.tls_ssl_version = tls_ssl_version self.tmp_dir = tmp_dir self.user = user self.volumes = volumes or [] self.working_dir = working_dir self.xcom_push_flag = xcom_push self.xcom_all = xcom_all self.auto_remove = auto_remove self.cli = None self.container = None def execute(self, context): self.log.info('Starting docker container from image %s', self.image) tls_config = None if self.tls_ca_cert and self.tls_client_cert and self.tls_client_key: tls_config = tls.TLSConfig( ca_cert=self.tls_ca_cert, client_cert=(self.tls_client_cert, self.tls_client_key), verify=True, ssl_version=self.tls_ssl_version, assert_hostname=self.tls_hostname ) self.docker_url = self.docker_url.replace('tcp://', 'https://') self.cli = APIClient(base_url=self.docker_url, version=self.api_version, tls=tls_config) if ':' not in self.image: image = self.image + ':latest' else: image = self.image if self.force_pull or len(self.cli.images(name=image)) == 0: self.log.info('Pulling docker image %s', image) for l in self.cli.pull(image, stream=True): output = json.loads(l.decode('utf-8')) self.log.info("%s", output['status']) cpu_shares = int(round(self.cpus * 1024)) with TemporaryDirectory(prefix='airflowtmp') as host_tmp_dir: self.environment['AIRFLOW_TMP_DIR'] = self.tmp_dir self.volumes.append('{0}:{1}'.format(host_tmp_dir, self.tmp_dir)) self.container = self.cli.create_container( command=self.get_command(), cpu_shares=cpu_shares, environment=self.environment, host_config=self.cli.create_host_config( binds=self.volumes, network_mode=self.network_mode, auto_remove=self.auto_remove), image=image, mem_limit=self.mem_limit, user=self.user, working_dir=self.working_dir ) self.cli.start(self.container['Id']) line = '' for line in self.cli.logs(container=self.container['Id'], stream=True): line = line.strip() if hasattr(line, 'decode'): line = line.decode('utf-8') self.log.info(line) exit_code = self.cli.wait(self.container['Id']) if exit_code != 0: raise AirflowException('docker container failed') if self.xcom_push_flag: return self.cli.logs(container=self.container['Id']) if self.xcom_all else str(line) def get_command(self): if self.command is not None and self.command.strip().find('[') == 0: commands = ast.literal_eval(self.command) else: commands = self.command return commands def on_kill(self): if self.cli is not None: self.log.info('Stopping docker container') self.cli.stop(self.container['Id'])
class DockerController(object): def _load_config(self): config = os.environ.get('BROWSER_CONFIG', './config.yaml') with open(config) as fh: config = yaml.load(fh) config = config['browser_config'] for n, v in config.items(): new_v = os.environ.get(n) if not new_v: new_v = os.environ.get(n.upper()) if new_v: print('Setting Env Val: {0}={1}'.format(n, new_v)) config[n] = new_v return config def __init__(self): config = self._load_config() self.name = config['cluster_name'] self.label_name = config['label_name'] self.init_req_expire_secs = config['init_req_expire_secs'] self.queue_expire_secs = config['queue_expire_secs'] self.remove_expired_secs = config['remove_expired_secs'] self.api_version = config['api_version'] self.ports = config['ports'] self.port_bindings = dict((port, None) for port in self.ports.values()) self.max_containers = config['max_containers'] self.throttle_expire_secs = config['throttle_expire_secs'] self.browser_image_prefix = config['browser_image_prefix'] self.label_browser = config['label_browser'] self.label_prefix = config['label_prefix'] self.network_name = config['network_name'] self.volume_source = config['browser_volumes'] self.shm_size = config['shm_size'] self.default_browser = config['default_browser'] self._init_cli() while True: try: self._init_redis(config) break except BusyLoadingError: print('Waiting for Redis to Load...') time.sleep(5) def _init_cli(self): if os.path.exists('/var/run/docker.sock'): self.cli = APIClient(base_url='unix://var/run/docker.sock', version=self.api_version) else: kwargs = kwargs_from_env(assert_hostname=False) kwargs['version'] = self.api_version self.cli = APIClient(**kwargs) def _init_redis(self, config): redis_url = os.environ['REDIS_BROWSER_URL'] self.redis = redis.StrictRedis.from_url(redis_url, decode_responses=True) self.redis.setnx('next_client', '1') self.redis.setnx('max_containers', self.max_containers) self.redis.setnx('num_containers', '0') # TODO: support this #self.redis.set('cpu_auto_adjust', config['cpu_auto_adjust']) # if num_containers is invalid, reset to 0 try: assert (int(self.redis.get('num_containers') >= 0)) except: self.redis.set('num_containers', 0) self.redis.set('throttle_samples', config['throttle_samples']) self.redis.set('throttle_max_avg', config['throttle_max_avg']) self.duration = int(config['container_expire_secs']) self.redis.set('container_expire_secs', self.duration) def load_avail_browsers(self, params=None): filters = {"dangling": False} if params: all_filters = [] for k, v in params.items(): if k not in ('short'): all_filters.append(self.label_prefix + k + '=' + v) filters["label"] = all_filters else: filters["label"] = self.label_browser browsers = {} try: images = self.cli.images(filters=filters) for image in images: tags = image.get('RepoTags') id_ = self._get_primary_id(tags) if not id_: continue props = self._browser_info(image['Labels']) props['id'] = id_ browsers[id_] = props except: traceback.print_exc() return browsers def _get_primary_id(self, tags): if not tags: return None primary_tag = None for tag in tags: if not tag: continue if tag.endswith(':latest'): tag = tag.replace(':latest', '') if not tag.startswith(self.browser_image_prefix): continue # pick the longest tag as primary tag if not primary_tag or len(tag) > len(primary_tag): primary_tag = tag if primary_tag: return primary_tag[len(self.browser_image_prefix):] else: return None def load_browser(self, name, include_icon=False): tag = self.browser_image_prefix + name try: image = self.cli.inspect_image(tag) tags = image.get('RepoTags') props = self._browser_info(image['Config']['Labels'], include_icon=include_icon) props['id'] = self._get_primary_id(tags) props['tags'] = tags return props except: traceback.print_exc() return {} def _browser_info(self, labels, include_icon=False): props = {} caps = [] for n, v in labels.items(): wr_prop = n.split(self.label_prefix) if len(wr_prop) != 2: continue name = wr_prop[1] if not include_icon and name == 'icon': continue props[name] = v if name.startswith('caps.'): caps.append(name.split('.', 1)[1]) props['caps'] = ', '.join(caps) return props def _get_host_port(self, info, port, default_host): info = info['NetworkSettings']['Ports'][str(port) + '/tcp'] info = info[0] host = info['HostIp'] if host == '0.0.0.0' and default_host: host = default_host return host + ':' + info['HostPort'] def _get_port(self, info, port): info = info['NetworkSettings']['Ports'][str(port) + '/tcp'] info = info[0] return info['HostPort'] def sid(self, id): return id[:12] def timed_new_container(self, browser, env, host, reqid): start = time.time() info = self.new_container(browser, env, host) end = time.time() dur = end - start time_key = 't:' + reqid self.redis.setex(time_key, self.throttle_expire_secs, dur) throttle_samples = int(self.redis.get('throttle_samples')) print('INIT DUR: ' + str(dur)) self.redis.lpush('init_timings', time_key) self.redis.ltrim('init_timings', 0, throttle_samples - 1) return info def new_container(self, browser_id, env=None, default_host=None): #browser = self.browsers.get(browser_id) browser = self.load_browser(browser_id) # get default browser if not browser: browser = self.load_browser(browser_id) #browser = self.browsers.get(self.default_browser) if browser.get('req_width'): env['SCREEN_WIDTH'] = browser.get('req_width') if browser.get('req_height'): env['SCREEN_HEIGHT'] = browser.get('req_height') image = browser['tags'][0] print('Launching ' + image) short_id = None try: host_config = self.create_host_config() container = self.cli.create_container( image=image, ports=list(self.ports.values()), environment=env, runtime="nvidia", host_config=host_config, labels={self.label_name: self.name}, ) #container = self.cli.create_container(image=image, # ports=list(self.ports.values()), # environment=env, # host_config=host_config, # labels={self.label_name: self.name}, # ) id_ = container.get('Id') short_id = self.sid(id_) res = self.cli.start(container=id_) info = self.cli.inspect_container(id_) ip = info['NetworkSettings']['IPAddress'] if not ip: ip = info['NetworkSettings']['Networks'][ self.network_name]['IPAddress'] self.redis.hset('all_containers', short_id, ip) result = {} for port_name in self.ports: result[port_name + '_host'] = self._get_host_port( info, self.ports[port_name], default_host) result['id'] = short_id result['ip'] = ip result['audio'] = os.environ.get('AUDIO_TYPE', '') return result except Exception as e: traceback.print_exc() if short_id: print('EXCEPTION: ' + short_id) self.remove_container(short_id) return {} def create_host_config(self): if self.volume_source: volumes_from = [self.volume_source] else: volumes_from = None host_config = self.cli.create_host_config( binds={ '/tmp/.X11-unix/X0': { 'bind': '/tmp/.X11-unix/X0', 'ro': False }, }, port_bindings=self.port_bindings, volumes_from=volumes_from, network_mode=self.network_name, shm_size=self.shm_size, cap_add=['ALL'], security_opt=['apparmor=unconfined'], privileged=True, runtime="nvidia", ) return host_config def remove_container(self, short_id): print('REMOVING: ' + short_id) try: self.cli.remove_container(short_id, force=True) except Exception as e: print(e) reqid = None ip = self.redis.hget('all_containers', short_id) if ip: reqid = self.redis.hget('ip:' + ip, 'reqid') with redis.utils.pipeline(self.redis) as pi: pi.delete('ct:' + short_id) if not ip: return pi.hdel('all_containers', short_id) pi.delete('ip:' + ip) if reqid: pi.delete('req:' + reqid) def event_loop(self): for event in self.cli.events(decode=True): try: self.handle_docker_event(event) except Exception as e: print(e) def handle_docker_event(self, event): if event['Type'] != 'container': return if (event['status'] == 'die' and event['from'].startswith(self.browser_image_prefix) and event['Actor']['Attributes'].get( self.label_name) == self.name): short_id = self.sid(event['id']) print('EXITED: ' + short_id) self.remove_container(short_id) self.redis.decr('num_containers') return if (event['status'] == 'start' and event['from'].startswith(self.browser_image_prefix) and event['Actor']['Attributes'].get( self.label_name) == self.name): short_id = self.sid(event['id']) print('STARTED: ' + short_id) self.redis.incr('num_containers') self.redis.setex('ct:' + short_id, self.duration, 1) return def remove_expired_loop(self): while True: try: self.remove_expired() except Exception as e: print(e) time.sleep(self.remove_expired_secs) def remove_expired(self): all_known_ids = self.redis.hkeys('all_containers') all_containers = { self.sid(c['Id']) for c in self.cli.containers(quiet=True) } for short_id in all_known_ids: if not self.redis.get('ct:' + short_id): print('TIME EXPIRED: ' + short_id) self.remove_container(short_id) elif short_id not in all_containers: print('STALE ID: ' + short_id) self.remove_container(short_id) def auto_adjust_max(self): print('Auto-Adjust Max Loop') try: scale = self.redis.get('cpu_auto_adjust') if not scale: return info = self.cli.info() cpus = int(info.get('NCPU', 0)) if cpus <= 1: return total = int(float(scale) * cpus) self.redis.set('max_containers', total) except Exception as e: traceback.print_exc() def add_new_client(self, reqid): client_id = self.redis.incr('clients') #enc_id = base64.b64encode(os.urandom(27)).decode('utf-8') self.redis.setex('cm:' + reqid, self.queue_expire_secs, client_id) self.redis.setex('q:' + str(client_id), self.queue_expire_secs, 1) return client_id def _make_reqid(self): return base64.b32encode(os.urandom(15)).decode('utf-8') def _make_vnc_pass(self): return base64.b64encode(os.urandom(21)).decode('utf-8') def register_request(self, container_data): reqid = self._make_reqid() container_data['reqid'] = reqid self.redis.hmset('req:' + reqid, container_data) self.redis.expire('req:' + reqid, self.init_req_expire_secs) return reqid def am_i_next(self, reqid): client_id = self.redis.get('cm:' + reqid) if not client_id: client_id = self.add_new_client(reqid) else: self.redis.expire('cm:' + reqid, self.queue_expire_secs) client_id = int(client_id) next_client = int(self.redis.get('next_client')) # not next client if client_id != next_client: # if this client expired, delete it from queue if not self.redis.get('q:' + str(next_client)): print('skipping expired', next_client) self.redis.incr('next_client') # missed your number somehow, get a new one! if client_id < next_client: client_id = self.add_new_client(reqid) diff = client_id - next_client if self.throttle(): self.redis.expire('q:' + str(client_id), self.queue_expire_secs) return client_id - next_client #num_containers = self.redis.hlen('all_containers') num_containers = int(self.redis.get('num_containers')) max_containers = self.redis.get('max_containers') max_containers = int( max_containers) if max_containers else self.max_containers if diff <= (max_containers - num_containers): self.redis.incr('next_client') return -1 else: self.redis.expire('q:' + str(client_id), self.queue_expire_secs) return client_id - next_client def throttle(self): timings = self.redis.lrange('init_timings', 0, -1) if not timings: return False timings = self.redis.mget(*timings) avg = 0 count = 0 for val in timings: if val is not None: avg += float(val) count += 1 if count == 0: return False avg = avg / count print('AVG: ', avg) throttle_max_avg = float(self.redis.get('throttle_max_avg')) if avg >= throttle_max_avg: print('Throttling, too slow...') return True return False def _copy_env(self, env, name, override=None): env[name] = override or os.environ.get(name) def init_new_browser(self, reqid, host, width=None, height=None): req_key = 'req:' + reqid container_data = self.redis.hgetall(req_key) if not container_data: return None # already started, attempt to reconnect if 'queue' in container_data: container_data['ttl'] = self.redis.ttl('ct:' + container_data['id']) return container_data queue_pos = self.am_i_next(reqid) if queue_pos >= 0: return {'queue': queue_pos} browser = container_data['browser'] url = container_data.get('url', 'about:blank') ts = container_data.get('request_ts') env = {} env['URL'] = url env['TS'] = ts env['BROWSER'] = browser vnc_pass = self._make_vnc_pass() env['VNC_PASS'] = vnc_pass self._copy_env(env, 'PROXY_HOST') self._copy_env(env, 'PROXY_PORT') self._copy_env(env, 'PROXY_GET_CA') self._copy_env(env, 'SCREEN_WIDTH', width) self._copy_env(env, 'SCREEN_HEIGHT', height) self._copy_env(env, 'IDLE_TIMEOUT') self._copy_env(env, 'AUDIO_TYPE') info = self.timed_new_container(browser, env, host, reqid) info['queue'] = 0 info['vnc_pass'] = vnc_pass new_key = 'ip:' + info['ip'] # TODO: support different durations? self.duration = int(self.redis.get('container_expire_secs')) with redis.utils.pipeline(self.redis) as pi: pi.rename(req_key, new_key) pi.persist(new_key) pi.hmset(req_key, info) pi.expire(req_key, self.duration) info['ttl'] = self.duration return info def get_random_browser(self): browsers = self.load_avail_browsers() while True: id_ = random.choice(browsers.keys()) if browsers[id_].get('skip_random'): continue return id_
class DockerProxy: """ A wrapper over docker-py and some utility methods and classes. """ LOG_TAG = "Docker " shell_commands = ["source"] class ImageBuildException(Exception): def __init__(self, message=None): super("Something went wrong while building docker container image.\n{0}".format(message)) def __init__(self): self.client = Client(base_url=Constants.DOCKER_BASE_URL) self.build_count = 0 logging.basicConfig(level=logging.DEBUG) @staticmethod def get_container_volume_from_working_dir(working_directory): import os return os.path.join("/home/ubuntu/", os.path.basename(working_directory)) def create_container(self, image_str, working_directory=None, name=None, port_bindings={Constants.DEFAULT_PUBLIC_WEBSERVER_PORT: ('127.0.0.1', 8080), Constants.DEFAULT_PRIVATE_NOTEBOOK_PORT: ('127.0.0.1', 8081)}): """Creates a new container with elevated privileges. Returns the container ID. Maps port 80 of container to 8080 of locahost by default""" docker_image = DockerImage.from_string(image_str) volume_dir = DockerProxy.get_container_volume_from_working_dir(working_directory) if name is None: import uuid random_str = str(uuid.uuid4()) name = constants.Constants.MolnsDockerContainerNamePrefix + random_str[:8] image = docker_image.image_id if docker_image.image_id is not Constants.DockerNonExistentTag \ else docker_image.image_tag logging.info("Using image {0}".format(image)) import os if DockerProxy._verify_directory(working_directory) is False: if working_directory is not None: raise InvalidVolumeName("\n\nMOLNs uses certain reserved names for its configuration files in the " "controller environment, and unfortunately the provided name for working " "directory of the controller cannot be one of these. Please configure this " "controller again with a different volume name and retry. " "Here is the list of forbidden names: \n{0}" .format(Constants.ForbiddenVolumeNames)) logging.warning(DockerProxy.LOG_TAG + "Unable to verify provided directory to use to as volume. Volume will NOT " "be created.") hc = self.client.create_host_config(privileged=True, port_bindings=port_bindings) container = self.client.create_container(image=image, name=name, command="/bin/bash", tty=True, detach=True, ports=[Constants.DEFAULT_PUBLIC_WEBSERVER_PORT, Constants.DEFAULT_PRIVATE_NOTEBOOK_PORT], host_config=hc, environment={"PYTHONPATH": "/usr/local/"}) else: container_mount_point = '/home/ubuntu/{0}'.format(os.path.basename(working_directory)) hc = self.client.create_host_config(privileged=True, port_bindings=port_bindings, binds={working_directory: {'bind': container_mount_point, 'mode': 'rw'}}) container = self.client.create_container(image=image, name=name, command="/bin/bash", tty=True, detach=True, ports=[Constants.DEFAULT_PUBLIC_WEBSERVER_PORT, Constants.DEFAULT_PRIVATE_NOTEBOOK_PORT], volumes=container_mount_point, host_config=hc, working_dir=volume_dir, environment={"PYTHONPATH": "/usr/local/"}) container_id = container.get("Id") return container_id # noinspection PyBroadException @staticmethod def _verify_directory(working_directory): import os if working_directory is None or os.path.basename(working_directory) in Constants.ForbiddenVolumeNames: return False try: if not os.path.exists(working_directory): os.makedirs(working_directory) return True except: return False def stop_containers(self, container_ids): """Stops given containers.""" for container_id in container_ids: self.stop_container(container_id) def stop_container(self, container_id): """Stops the container with given ID.""" self.client.stop(container_id) def container_status(self, container_id): """Checks if container with given ID running.""" status = ProviderBase.STATUS_TERMINATED try: ret_val = str(self.client.inspect_container(container_id).get('State').get('Status')) if ret_val.startswith("running"): status = ProviderBase.STATUS_RUNNING else: status = ProviderBase.STATUS_STOPPED except NotFound: pass return status def start_containers(self, container_ids): """Starts each container in given list of container IDs.""" for container_id in container_ids: self.start_container(container_id) def start_container(self, container_id): """ Start the container with given ID.""" logging.info(DockerProxy.LOG_TAG + " Starting container " + container_id) try: self.client.start(container=container_id) except (NotFound, NullResource) as e: print (DockerProxy.LOG_TAG + "Something went wrong while starting container.", e) return False return True def execute_command(self, container_id, command): """Executes given command as a shell command in the given container. Returns None is anything goes wrong.""" run_command = "/bin/bash -c \"" + command + "\"" # print("CONTAINER: {0} COMMAND: {1}".format(container_id, run_command)) if self.start_container(container_id) is False: print (DockerProxy.LOG_TAG + "Could not start container.") return None try: exec_instance = self.client.exec_create(container_id, run_command) response = self.client.exec_start(exec_instance) return [self.client.exec_inspect(exec_instance), response] except (NotFound, APIError) as e: print (DockerProxy.LOG_TAG + " Could not execute command.", e) return None def build_image(self, dockerfile): """ Build image from given Dockerfile object and return ID of the image created. """ import uuid logging.info("Building image...") random_string = str(uuid.uuid4()) image_tag = Constants.DOCKER_IMAGE_PREFIX + "{0}".format(random_string[:]) last_line = "" try: for line in self.client.build(fileobj=dockerfile, rm=True, tag=image_tag): print(DockerProxy._decorate(line)) if "errorDetail" in line: raise DockerProxy.ImageBuildException() last_line = line # Return image ID. It's a hack around the fact that docker-py's build image command doesn't return an image # id. image_id = get_docker_image_id_from_string(str(last_line)) logging.info("Image ID: {0}".format(image_id)) return str(DockerImage(image_id, image_tag)) except (DockerProxy.ImageBuildException, IndexError) as e: raise DockerProxy.ImageBuildException(e) @staticmethod def _decorate(some_line): return some_line[11:-4].rstrip() def image_exists(self, image_str): """Checks if an image with the given ID/tag exists locally.""" docker_image = DockerImage.from_string(image_str) if docker_image.image_id is Constants.DockerNonExistentTag \ and docker_image.image_tag is Constants.DockerNonExistentTag: raise InvalidDockerImageException("Neither image_id nor image_tag provided.") for image in self.client.images(): some_id = image["Id"] some_tags = image["RepoTags"] or [None] if docker_image.image_id in \ some_id[:(Constants.DOCKER_PY_IMAGE_ID_PREFIX_LENGTH + Constants.DOKCER_IMAGE_ID_LENGTH)]: return True if docker_image.image_tag in some_tags: return True return False def terminate_containers(self, container_ids): """ Terminates containers with given container ids.""" for container_id in container_ids: try: if self.container_status(container_id) == ProviderBase.STATUS_RUNNING: self.stop_container(container_id) self.terminate_container(container_id) except NotFound: pass def terminate_container(self, container_id): self.client.remove_container(container_id) def get_mapped_ports(self, container_id): container_ins = self.client.inspect_container(container_id) mapped_ports = container_ins['HostConfig']['PortBindings'] ret_val = [] if mapped_ports is None: logging.info("No mapped ports for {0}".format(container_id)) return for k, v in mapped_ports.iteritems(): host_port = v[0]['HostPort'] ret_val.append(host_port) return ret_val def get_working_directory(self, container_id): return self.client.inspect_container(container_id)["Config"]["WorkingDir"] def get_home_directory(self, container_id): env_vars = self.client.inspect_container(container_id)["Config"]["Env"] home = [i for i in env_vars if i.startswith("HOME")] return home[0].split("=")[1] def put_archive(self, container_id, tar_file_bytes, target_path_in_container): """ Copies and unpacks a given tarfile in the container at specified location. Location must exist in container.""" if self.start_container(container_id) is False: raise Exception("Could not start container.") # Prepend file path with /home/ubuntu/. TODO Should be refined. if not target_path_in_container.startswith("/home/ubuntu/"): import os target_path_in_container = os.path.join("/home/ubuntu/", target_path_in_container) logging.info("target path in container: {0}".format(target_path_in_container)) if not self.client.put_archive(container_id, target_path_in_container, tar_file_bytes): logging.error(DockerProxy.LOG_TAG + "Failed to copy.") def get_container_ip_address(self, container_id): """ Returns the IP Address of given container.""" self.start_container(container_id) ins = self.client.inspect_container(container_id) ip_address = str(ins.get("NetworkSettings").get("IPAddress")) while True: ip_address = str(ins.get("NetworkSettings").get("IPAddress")) if ip_address == "": time.sleep(3) if ip_address.startswith("1") is True: break return ip_address
class ModifiedDockerOperator(DockerOperator): """ModifiedDockerOperator supports host temporary directories on OSX. Incorporates https://github.com/apache/airflow/pull/4315/ and an implementation of https://issues.apache.org/jira/browse/AIRFLOW-3825. :param host_tmp_dir: Specify the location of the temporary directory on the host which will be mapped to tmp_dir. If not provided defaults to using the standard system temp directory. :type host_tmp_dir: str """ def __init__(self, host_tmp_dir='/tmp', **kwargs): self.host_tmp_dir = host_tmp_dir kwargs['xcom_push'] = True super(ModifiedDockerOperator, self).__init__(**kwargs) @contextmanager def get_host_tmp_dir(self): '''Abstracts the tempdir context manager so that this can be overridden.''' with TemporaryDirectory(prefix='airflowtmp', dir=self.host_tmp_dir) as tmp_dir: yield tmp_dir def execute(self, context): '''Modified only to use the get_host_tmp_dir helper.''' self.log.info('Starting docker container from image %s', self.image) tls_config = self.__get_tls_config() if self.docker_conn_id: self.cli = self.get_hook().get_conn() else: self.cli = APIClient(base_url=self.docker_url, version=self.api_version, tls=tls_config) if self.force_pull or len(self.cli.images(name=self.image)) == 0: self.log.info('Pulling docker image %s', self.image) for l in self.cli.pull(self.image, stream=True): output = json.loads(l.decode('utf-8').strip()) if 'status' in output: self.log.info("%s", output['status']) with self.get_host_tmp_dir() as host_tmp_dir: self.environment['AIRFLOW_TMP_DIR'] = self.tmp_dir self.volumes.append('{0}:{1}'.format(host_tmp_dir, self.tmp_dir)) self.container = self.cli.create_container( command=self.get_command(), environment=self.environment, host_config=self.cli.create_host_config( auto_remove=self.auto_remove, binds=self.volumes, network_mode=self.network_mode, shm_size=self.shm_size, dns=self.dns, dns_search=self.dns_search, cpu_shares=int(round(self.cpus * 1024)), mem_limit=self.mem_limit, ), image=self.image, user=self.user, working_dir=self.working_dir, ) self.cli.start(self.container['Id']) res = [] line = '' for new_line in self.cli.logs(container=self.container['Id'], stream=True): line = new_line.strip() if hasattr(line, 'decode'): line = line.decode('utf-8') self.log.info(line) res.append(line) result = self.cli.wait(self.container['Id']) if result['StatusCode'] != 0: raise AirflowException('docker container failed: ' + repr(result)) if self.xcom_push_flag: # Try to avoid any kind of race condition? return '\n'.join(res) + '\n' if self.xcom_all else str(line) # This is a class-private name on DockerOperator for no good reason -- # all that the status quo does is inhibit extension of the class. # See https://issues.apache.org/jira/browse/AIRFLOW-3880 def __get_tls_config(self): # pylint: disable=no-member return super(ModifiedDockerOperator, self)._DockerOperator__get_tls_config()
class DagsterDockerOperator(DockerOperator): """Dagster operator for Apache Airflow. Wraps a modified DockerOperator incorporating https://github.com/apache/airflow/pull/4315. Additionally, if a Docker client can be initialized using docker.from_env, Unlike the standard DockerOperator, this operator also supports config using docker.from_env, so it isn't necessary to explicitly set docker_url, tls_config, or api_version. Incorporates https://github.com/apache/airflow/pull/4315/ and an implementation of https://issues.apache.org/jira/browse/AIRFLOW-3825. Parameters: host_tmp_dir (str): Specify the location of the temporary directory on the host which will be mapped to tmp_dir. If not provided defaults to using the standard system temp directory. """ def __init__(self, dagster_operator_parameters, *args): kwargs = dagster_operator_parameters.op_kwargs tmp_dir = kwargs.pop("tmp_dir", DOCKER_TEMPDIR) host_tmp_dir = kwargs.pop("host_tmp_dir", seven.get_system_temp_directory()) self.host_tmp_dir = host_tmp_dir run_config = dagster_operator_parameters.run_config if "filesystem" in run_config["intermediate_storage"]: if ( "config" in (run_config["intermediate_storage"].get("filesystem", {}) or {}) and "base_dir" in ( (run_config["intermediate_storage"].get("filesystem", {}) or {}).get( "config", {} ) or {} ) and run_config["intermediate_storage"]["filesystem"]["config"]["base_dir"] != tmp_dir ): warnings.warn( "Found base_dir '{base_dir}' set in filesystem storage config, which was not " "the tmp_dir we expected ('{tmp_dir}', mounting host_tmp_dir " "'{host_tmp_dir}' from the host). We assume you know what you are doing, but " "if you are having trouble executing containerized workloads, this may be the " "issue".format( base_dir=run_config["intermediate_storage"]["filesystem"]["config"][ "base_dir" ], tmp_dir=tmp_dir, host_tmp_dir=host_tmp_dir, ) ) else: run_config["intermediate_storage"]["filesystem"] = dict( run_config["intermediate_storage"]["filesystem"] or {}, **{ "config": dict( ( ( run_config["intermediate_storage"].get("filesystem", {}) or {} ).get("config", {}) or {} ), **{"base_dir": tmp_dir}, ) }, ) self.docker_conn_id_set = kwargs.get("docker_conn_id") is not None self.run_config = run_config self.pipeline_name = dagster_operator_parameters.pipeline_name self.pipeline_snapshot = dagster_operator_parameters.pipeline_snapshot self.execution_plan_snapshot = dagster_operator_parameters.execution_plan_snapshot self.parent_pipeline_snapshot = dagster_operator_parameters.parent_pipeline_snapshot self.mode = dagster_operator_parameters.mode self.step_keys = dagster_operator_parameters.step_keys self.recon_repo = dagster_operator_parameters.recon_repo self._run_id = None self.instance_ref = dagster_operator_parameters.instance_ref check.invariant(self.instance_ref) self.instance = DagsterInstance.from_ref(self.instance_ref) # These shenanigans are so we can override DockerOperator.get_hook in order to configure # a docker client using docker.from_env, rather than messing with the logic of # DockerOperator.execute if not self.docker_conn_id_set: try: from_env().version() except Exception: # pylint: disable=broad-except pass else: kwargs["docker_conn_id"] = True if "environment" not in kwargs: kwargs["environment"] = get_aws_environment() super(DagsterDockerOperator, self).__init__( task_id=dagster_operator_parameters.task_id, dag=dagster_operator_parameters.dag, tmp_dir=tmp_dir, host_tmp_dir=host_tmp_dir, xcom_push=True, # We do this because log lines won't necessarily be emitted in order (!) -- so we can't # just check the last log line to see if it's JSON. xcom_all=True, *args, **kwargs, ) @contextmanager def get_host_tmp_dir(self): yield self.host_tmp_dir def execute_raw(self, context): """Modified only to use the get_host_tmp_dir helper.""" self.log.info("Starting docker container from image %s", self.image) tls_config = self.__get_tls_config() if self.docker_conn_id: self.cli = self.get_hook().get_conn() else: self.cli = APIClient(base_url=self.docker_url, version=self.api_version, tls=tls_config) if self.force_pull or len(self.cli.images(name=self.image)) == 0: self.log.info("Pulling docker image %s", self.image) for l in self.cli.pull(self.image, stream=True): output = seven.json.loads(l.decode("utf-8").strip()) if "status" in output: self.log.info("%s", output["status"]) with self.get_host_tmp_dir() as host_tmp_dir: self.environment["AIRFLOW_TMP_DIR"] = self.tmp_dir self.volumes.append("{0}:{1}".format(host_tmp_dir, self.tmp_dir)) self.container = self.cli.create_container( command=self.get_docker_command(context.get("ts")), environment=self.environment, host_config=self.cli.create_host_config( auto_remove=self.auto_remove, binds=self.volumes, network_mode=self.network_mode, shm_size=self.shm_size, dns=self.dns, dns_search=self.dns_search, cpu_shares=int(round(self.cpus * 1024)), mem_limit=self.mem_limit, ), image=self.image, user=self.user, working_dir=self.working_dir, ) self.cli.start(self.container["Id"]) res = [] line = "" for new_line in self.cli.logs( container=self.container["Id"], stream=True, stdout=True, stderr=False ): line = new_line.strip() if hasattr(line, "decode"): line = line.decode("utf-8") self.log.info(line) res.append(line) result = self.cli.wait(self.container["Id"]) if result["StatusCode"] != 0: raise AirflowException( "docker container failed with result: {result} and logs: {logs}".format( result=repr(result), logs="\n".join(res) ) ) if self.xcom_push_flag: # Try to avoid any kind of race condition? return res if self.xcom_all else str(line) # This is a class-private name on DockerOperator for no good reason -- # all that the status quo does is inhibit extension of the class. # See https://issues.apache.org/jira/browse/AIRFLOW-3880 def __get_tls_config(self): # pylint: disable=no-member return super(DagsterDockerOperator, self)._DockerOperator__get_tls_config() @property def run_id(self): if self._run_id is None: return "" else: return self._run_id def query(self, airflow_ts): check.opt_str_param(airflow_ts, "airflow_ts") recon_pipeline = self.recon_repo.get_reconstructable_pipeline(self.pipeline_name) input_json = serialize_dagster_namedtuple( ExecuteStepArgs( pipeline_origin=recon_pipeline.get_python_origin(), pipeline_run_id=self.run_id, instance_ref=self.instance_ref, step_keys_to_execute=self.step_keys, ) ) command = "dagster api execute_step {}".format(json.dumps(input_json)) self.log.info("Executing: {command}\n".format(command=command)) return command def get_docker_command(self, airflow_ts): """Deliberately renamed from get_command to avoid shadoowing the method of the base class""" check.opt_str_param(airflow_ts, "airflow_ts") if self.command is not None and self.command.strip().find("[") == 0: commands = ast.literal_eval(self.command) elif self.command is not None: commands = self.command else: commands = self.query(airflow_ts) return commands def get_hook(self): if self.docker_conn_id_set: return super(DagsterDockerOperator, self).get_hook() class _DummyHook: def get_conn(self): return from_env().api return _DummyHook() def execute(self, context): if "run_id" in self.params: self._run_id = self.params["run_id"] elif "dag_run" in context and context["dag_run"] is not None: self._run_id = context["dag_run"].run_id try: tags = {AIRFLOW_EXECUTION_DATE_STR: context.get("ts")} if "ts" in context else {} self.instance.register_managed_run( pipeline_name=self.pipeline_name, run_id=self.run_id, run_config=self.run_config, mode=self.mode, solids_to_execute=None, step_keys_to_execute=None, tags=tags, root_run_id=None, parent_run_id=None, pipeline_snapshot=self.pipeline_snapshot, execution_plan_snapshot=self.execution_plan_snapshot, parent_pipeline_snapshot=self.parent_pipeline_snapshot, ) res = self.execute_raw(context) self.log.info("Finished executing container.") if not res: raise AirflowException("Missing query response") try: events = [deserialize_json_to_dagster_namedtuple(line) for line in res if line] except Exception: # pylint: disable=broad-except raise AirflowException( "Could not parse response {response}".format(response=repr(res)) ) if len(events) == 1 and isinstance(events[0], StepExecutionSkipped): raise AirflowSkipException( "Dagster emitted skip event, skipping execution in Airflow" ) check_events_for_failures(events) check_events_for_skips(events) return events finally: self._run_id = None
class DockerNode(CommonNode): """ An instance of this class will create a detached Docker container. This node binds the ``shared_dir_mount`` directory of the container to a local path in the host system defined in ``self.shared_dir``. :param str identifier: Node unique identifier in the topology being built. :param str image: The image to run on this node, in the form ``repository:tag``. :param str registry: Docker registry to pull image from. :param str command: The command to run when the container is brought up. :param str binds: Directories to bind for this container separated by a ``;`` in the form: :: '/tmp:/tmp;/dev/log:/dev/log;/sys/fs/cgroup:/sys/fs/cgroup' :param str network_mode: Network mode for this container. :param str hostname: Container hostname. :param environment: Environment variables to pass to the container. They can be set as a list of strings in the following format: :: ['environment_variable=value'] or as a dictionary in the following format: :: {'environment_variable': 'value'} :type environment: list or dict :param bool privileged: Run container in privileged mode or not. :param bool tty: Whether to allocate a TTY or not to the process. :param str shared_dir_base: Base path in the host where the shared directory will be created. The shared directory will always have the name of the container inside this directory. :param str shared_dir_mount: Mount point of the shared directory in the container. :param dict create_host_config_kwargs: Extra kwargs arguments to pass to docker-py's ``create_host_config()`` low-level API call. :param dict create_container_kwargs: Extra kwargs arguments to pass to docker-py's ``create_container()`` low-level API call. Read only public attributes: :var str image: Name of the Docker image being used by this node. Same as the ``image`` keyword argument. :var str container_id: Unique container identifier assigned by the Docker daemon in the form of a hash. :var str container_name: Unique container name assigned by the framework in the form ``{identifier}_{pid}_{timestamp}``. :var str shared_dir: Share directory in the host for this container. Always ``/tmp/topology/{container_name}``. :var str shared_dir_mount: Directory inside the container where the ``shared_dir`` is mounted. Same as the ``shared_dir_mount`` keyword .. automethod:: _get_network_config """ @abstractmethod def __init__(self, identifier, image='ubuntu:latest', registry=None, command='bash', binds=None, network_mode='none', hostname=None, environment=None, privileged=True, tty=True, shared_dir_base='/tmp/topology/docker/', shared_dir_mount='/var/topology', create_host_config_kwargs=None, create_container_kwargs=None, **kwargs): super(DockerNode, self).__init__(identifier, **kwargs) self._pid = None self._image = image self._registry = registry self._command = command self._hostname = hostname self._environment = environment self._client = APIClient(version='auto') self._container_name = '{identifier}_{pid}_{timestamp}'.format( identifier=identifier, pid=getpid(), timestamp=datetime.now().isoformat().replace(':', '-')) self._shared_dir_base = shared_dir_base self._shared_dir_mount = shared_dir_mount self._shared_dir = join(shared_dir_base, self._container_name) self._create_host_config_kwargs = create_host_config_kwargs or {} self._create_container_kwargs = create_container_kwargs or {} # Autopull docker image if necessary self._autopull() # Create shared directory ensure_dir(self._shared_dir) # Add binded directories container_binds = [ '{}:{}'.format(self._shared_dir, self._shared_dir_mount) ] if binds is not None: container_binds.extend(binds.split(';')) # Create host config create_host_config_call = { 'privileged': privileged, 'network_mode': network_mode, 'binds': container_binds, 'init': True } create_host_config_call.update(self._create_host_config_kwargs) self._host_config = self._client.create_host_config( **create_host_config_call) # Create container create_container_call = { 'image': self._image, 'command': self._command, 'name': self._container_name, 'detach': True, 'tty': tty, 'hostname': self._hostname, 'host_config': self._host_config, 'environment': self._environment, } create_container_call.update(self._create_container_kwargs) self._container_id = self._client.create_container( **create_container_call)['Id'] @property def image(self): return self._image @property def container_id(self): return self._container_id @property def container_name(self): return self._container_name @property def shared_dir(self): return self._shared_dir @property def shared_dir_mount(self): return self._shared_dir_mount def _get_network_config(self): """ Defines the network configuration for nodes of this type. This method should be overriden when implementing a new node type to return a dictionary with its network configuration by setting the following components: 'mapping' This is a dictionary of dictionaries, each parent-level key defines one network category, and each category *must* have these three keys: **netns**, **managed_by**, and **prefix**, and *can* (optionally) have a **connect_to** key). 'netns' Specifies the network namespace (inside the docker container) where all the ports belonging to this category will be moved after their creation. If set to None, then the ports will remain in the container's default network namespace. 'managed_by' Specifies who will manage different aspects of this network category depending on its value (which can be either **docker** or **platform**). 'docker' This network category will represent a network created by docker (identical to using the docker network create command) and will be visible to docker (right now all docker-managed networks are created using docker's 'bridge' built-in network plugin, this will likely change in the near future). 'platform' This network category will represent ports created by the Docker Platform Engine and is invisible to docker. 'prefix' Defines a prefix that will be used when a port/interface is moved into a namespace, its value can be set to '' (empty string) if no prefix is needed. In cases where the parent network category doesn't have a netns (i.e. 'netns' is set to None) this value will be ignored. 'connect_to' Specifies a Docker network this category will be connected to, if this network doesn't exists it will be created. If set to None, this category will be connected to a uniquely named Docker network that will be created by the platform. 'default_category' Every port that didn't explicitly set its category (using the "category" attribute in the SZN definition) will be set to this category. This is an example of a network configuration dictionary as expected to be returned by this funcition:: { 'default_category': 'front_panel', 'mapping': { 'oobm': { 'netns': 'oobmns', 'managed_by': 'docker', 'connect_to': 'oobm' 'prefix': '' }, 'back_panel': { 'netns': None, 'managed_by': 'docker', 'prefix': '' }, 'front_panel': { 'netns': 'front', 'managed_by': 'platform', 'prefix': 'f_' } } } :returns: The dictionary defining the network configuration. :rtype: dict """ return { 'default_category': 'front_panel', 'mapping': { 'oobm': { 'netns': None, 'managed_by': 'docker', 'prefix': '' }, 'front_panel': { 'netns': 'front_panel', 'managed_by': 'platform', 'prefix': '' } } } def _autopull(self): """ Autopulls the docker image of the node, if necessary. """ # Search for image in available images for tags in [img['RepoTags'] for img in self._client.images()]: # Docker py can return repo tags as None if tags and self._image in tags: return # Determine image parts registry = self._registry image = self._image tag = 'latest' if ':' in image: image, tag = image.split(':') # Pull image pull_uri = image if registry: pull_uri = '{}/{}'.format(registry, image) pull_name = '{}:{}'.format(pull_uri, tag) log.info('Trying to pull image {} ...'.format(pull_name)) last = '' for line in self._client.pull(pull_uri, tag=tag, stream=True): last = line status = loads(last.decode('utf8')) log.debug('Pulling result :: {}'.format(status)) if 'error' in status: raise Exception(status['error']) # Retag if required if pull_name != self._image: if not self._client.tag(pull_name, image, tag): raise Exception( 'Error when tagging image {} with tag {}:{}'.format( pull_name, image, tag)) log.info('Tagged image {} with tag {}:{}'.format( pull_name, image, tag)) def _docker_exec(self, command): """ Execute a command inside the docker. :param str command: The command to execute. """ log.debug('[{}]._docker_exec(\'{}\') ::'.format( self._container_id, command)) response = check_output( shsplit('docker exec {container_id} {command}'.format( container_id=self._container_id, command=command.strip()))).decode('utf8') log.debug(response) return response def _get_services_address(self): """ Get the service address of the node using Docker's inspect mechanism to grab OOBM interface address. :return: The address (IP or FQDN) of the services interface (oobm). :rtype: str """ network_name = self._container_name + '_oobm' address = self._client.inspect_container( self.container_id )['NetworkSettings']['Networks'][network_name]['IPAddress'] return address def notify_add_biport(self, node, biport): """ Get notified that a new biport was added to this engine node. :param node: The specification node that spawn this engine node. :type node: pynml.nml.Node :param biport: The specification bidirectional port added. :type biport: pynml.nml.BidirectionalPort :rtype: str :return: The assigned interface name of the port. """ network_config = self._get_network_config() category = biport.metadata.get('category', network_config['default_category']) category_config = network_config['mapping'][category] if category_config['managed_by'] == 'docker': netname = category_config.get( 'connect_to', '{}_{}'.format(self._container_name, category)) return get_iface_name(self, netname) else: return biport.metadata.get('label', biport.identifier) def notify_add_bilink(self, nodeport, bilink): """ Get notified that a new bilink was added to a port of this engine node. :param nodeport: A tuple with the specification node and port being linked. :type nodeport: (pynml.nml.Node, pynml.nml.BidirectionalPort) :param bilink: The specification bidirectional link added. :type bilink: pynml.nml.BidirectionalLink """ def notify_post_build(self): """ Get notified that the post build stage of the topology build was reached. """ # Log container data image_data = self._client.inspect_image(image=self._image) log.info('Started container {}:\n' ' Image name: {}\n' ' Image id: {}\n' ' Image creation date: {}' ' Image tags: {}'.format( self._container_name, self._image, image_data.get('Id', '????'), image_data.get('Created', '????'), ', '.join(image_data.get('RepoTags', [])))) container_data = self._client.inspect_container( container=self._container_id) log.debug(container_data) def start(self): """ Start the docker node and configures a netns for it. """ self._client.start(self._container_id) self._pid = self._client.inspect_container( self._container_id)['State']['Pid'] def stop(self): """ Request container to stop. """ self._client.stop(self._container_id) self._client.wait(self._container_id) self._client.remove_container(self._container_id) def disable(self): """ Disable the node. In Docker implementation this pauses the container. """ for portlbl in self.ports: self.set_port_state(portlbl, False) self._client.pause(self._container_id) def enable(self): """ Enable the node. In Docker implementation this unpauses the container. """ self._client.unpause(self._container_id) for portlbl in self.ports: self.set_port_state(portlbl, True) def set_port_state(self, portlbl, state): """ Set the given port label to the given state. :param str portlbl: The label of the port. :param bool state: True for up, False for down. """ iface = self.ports[portlbl] state = 'up' if state else 'down' command = ('ip netns exec front_panel ' 'ip link set dev {iface} {state}'.format(**locals())) self._docker_exec(command)
class DockerCli: def __init__(self): self.client = APIClient('unix://var/run/docker.sock') self.filtered_statuses = ('running', 'restarting', 'paused', 'exited') self.config = Config() def _get_containers(self, filters=None): filters = filters if filters else dict() for status in self.filtered_statuses: filters.update({'status': status}) for container in self.client.containers( all=True, filters=filters ): img_name, _, img_version = container['Image'].partition(':') service = self.config.get_service_by_name(img_name) if service: instance = dict() instance['created'] = container['Created'] instance['id'] = container['Id'] instance['image'] = img_name for con_port in container['Ports']: if service['port'] is con_port['PrivatePort']: instance['port'] = con_port.get('PublicPort') else: instance['port'] = None instance['state'] = container['State'] instance['status'] = container['Status'] instance['version'] = img_version yield instance return def get_all_containers(self): containers = [] for container in self._get_containers(): if container: containers.append(container) return containers def get_container(self, by_id): for container in self._get_containers({'id': by_id}): return container raise NotFoundContainerException( 'Container was not found: {}'.format(by_id) ) def create_container(self, image): service = self.config.get_service_by_name(image) if service: container = self.client.create_container( image='{0}:{1}'.format(image, service['version']), ports=[service['port']], detach=True, host_config=self.client.create_host_config( port_bindings={service['port']: None} ) ) self.client.start(container=container['Id']) return self.get_container(container['Id']) raise NotFoundImageException('Image was not found: {}'.format(image)) def remove_container(self, by_id): try: self.client.remove_container( container=by_id, force=True, v=True ) except errors.NotFound as e: raise NotFoundContainerException(e) return {'status': 'OK'}
class DockerOperator(BaseOperator): """ Execute a command inside a docker container. A temporary directory is created on the host and mounted into a container to allow storing files that together exceed the default disk size of 10GB in a container. The path to the mounted directory can be accessed via the environment variable ``AIRFLOW_TMP_DIR``. If a login to a private registry is required prior to pulling the image, a Docker connection needs to be configured in Airflow and the connection ID be provided with the parameter ``docker_conn_id``. :param image: Docker image from which to create the container. If image tag is omitted, "latest" will be used. :type image: str :param api_version: Remote API version. Set to ``auto`` to automatically detect the server's version. :type api_version: str :param command: Command to be run in the container. (templated) :type command: str or list :param container_name: Name of the container. :type container_name: str :param cpus: Number of CPUs to assign to the container. This value gets multiplied with 1024. See https://docs.docker.com/engine/reference/run/#cpu-share-constraint :type cpus: float :param docker_url: URL of the host running the docker daemon. Default is unix://var/run/docker.sock :type docker_url: str :param environment: Environment variables to set in the container. (templated) :type environment: dict :param force_pull: Pull the docker image on every run. Default is False. :type force_pull: bool :param mem_limit: Maximum amount of memory the container can use. Either a float value, which represents the limit in bytes, or a string like ``128m`` or ``1g``. :type mem_limit: float or str :param host_tmp_dir: Specify the location of the temporary directory on the host which will be mapped to tmp_dir. If not provided defaults to using the standard system temp directory. :type host_tmp_dir: str :param network_mode: Network mode for the container. :type network_mode: str :param tls_ca_cert: Path to a PEM-encoded certificate authority to secure the docker connection. :type tls_ca_cert: str :param tls_client_cert: Path to the PEM-encoded certificate used to authenticate docker client. :type tls_client_cert: str :param tls_client_key: Path to the PEM-encoded key used to authenticate docker client. :type tls_client_key: str :param tls_hostname: Hostname to match against the docker server certificate or False to disable the check. :type tls_hostname: str or bool :param tls_ssl_version: Version of SSL to use when communicating with docker daemon. :type tls_ssl_version: str :param tmp_dir: Mount point inside the container to a temporary directory created on the host by the operator. The path is also made available via the environment variable ``AIRFLOW_TMP_DIR`` inside the container. :type tmp_dir: str :param user: Default user inside the docker container. :type user: int or str :param volumes: List of volumes to mount into the container, e.g. ``['/host/path:/container/path', '/host/path2:/container/path2:ro']``. :type volumes: list :param working_dir: Working directory to set on the container (equivalent to the -w switch the docker client) :type working_dir: str :param xcom_all: Push all the stdout or just the last line. The default is False (last line). :type xcom_all: bool :param docker_conn_id: ID of the Airflow connection to use :type docker_conn_id: str :param dns: Docker custom DNS servers :type dns: list[str] :param dns_search: Docker custom DNS search domain :type dns_search: list[str] :param auto_remove: Auto-removal of the container on daemon side when the container's process exits. The default is False. :type auto_remove: bool :param shm_size: Size of ``/dev/shm`` in bytes. The size must be greater than 0. If omitted uses system default. :type shm_size: int """ template_fields = ('command', 'environment',) template_ext = ('.sh', '.bash',) @apply_defaults def __init__( self, image: str, api_version: str = None, command: Union[str, List[str]] = None, container_name: str = None, cpus: float = 1.0, docker_url: str = 'unix://var/run/docker.sock', environment: Dict = None, force_pull: bool = False, mem_limit: Union[float, str] = None, host_tmp_dir: str = None, network_mode: str = None, tls_ca_cert: str = None, tls_client_cert: str = None, tls_client_key: str = None, tls_hostname: Union[str, bool] = None, tls_ssl_version: str = None, tmp_dir: str = '/tmp/airflow', user: Union[str, int] = None, volumes: Iterable[str] = None, working_dir: str = None, xcom_all: bool = False, docker_conn_id: str = None, dns: List[str] = None, dns_search: List[str] = None, auto_remove: bool = False, shm_size: int = None, *args, **kwargs) -> None: super().__init__(*args, **kwargs) self.api_version = api_version self.auto_remove = auto_remove self.command = command self.container_name = container_name self.cpus = cpus self.dns = dns self.dns_search = dns_search self.docker_url = docker_url self.environment = environment or {} self.force_pull = force_pull self.image = image self.mem_limit = mem_limit self.host_tmp_dir = host_tmp_dir self.network_mode = network_mode self.tls_ca_cert = tls_ca_cert self.tls_client_cert = tls_client_cert self.tls_client_key = tls_client_key self.tls_hostname = tls_hostname self.tls_ssl_version = tls_ssl_version self.tmp_dir = tmp_dir self.user = user self.volumes = volumes or [] self.working_dir = working_dir self.xcom_all = xcom_all self.docker_conn_id = docker_conn_id self.shm_size = shm_size if kwargs.get('xcom_push') is not None: raise AirflowException("'xcom_push' was deprecated, use 'BaseOperator.do_xcom_push' instead") self.cli = None self.container = None def get_hook(self): return DockerHook( docker_conn_id=self.docker_conn_id, base_url=self.docker_url, version=self.api_version, tls=self.__get_tls_config() ) def execute(self, context): self.log.info('Starting docker container from image %s', self.image) tls_config = self.__get_tls_config() if self.docker_conn_id: self.cli = self.get_hook().get_conn() else: self.cli = APIClient( base_url=self.docker_url, version=self.api_version, tls=tls_config ) if self.force_pull or len(self.cli.images(name=self.image)) == 0: self.log.info('Pulling docker image %s', self.image) for l in self.cli.pull(self.image, stream=True): output = json.loads(l.decode('utf-8').strip()) if 'status' in output: self.log.info("%s", output['status']) with TemporaryDirectory(prefix='airflowtmp', dir=self.host_tmp_dir) as host_tmp_dir: self.environment['AIRFLOW_TMP_DIR'] = self.tmp_dir self.volumes.append('{0}:{1}'.format(host_tmp_dir, self.tmp_dir)) self.container = self.cli.create_container( command=self.get_command(), name=self.container_name, environment=self.environment, host_config=self.cli.create_host_config( auto_remove=self.auto_remove, binds=self.volumes, network_mode=self.network_mode, shm_size=self.shm_size, dns=self.dns, dns_search=self.dns_search, cpu_shares=int(round(self.cpus * 1024)), mem_limit=self.mem_limit), image=self.image, user=self.user, working_dir=self.working_dir ) self.cli.start(self.container['Id']) line = '' for line in self.cli.attach(container=self.container['Id'], stdout=True, stderr=True, stream=True): line = line.strip() if hasattr(line, 'decode'): line = line.decode('utf-8') self.log.info(line) result = self.cli.wait(self.container['Id']) if result['StatusCode'] != 0: raise AirflowException('docker container failed: ' + repr(result)) # duplicated conditional logic because of expensive operation if self.do_xcom_push: return self.cli.logs(container=self.container['Id']) \ if self.xcom_all else line.encode('utf-8') def get_command(self): if isinstance(self.command, str) and self.command.strip().find('[') == 0: commands = ast.literal_eval(self.command) else: commands = self.command return commands def on_kill(self): if self.cli is not None: self.log.info('Stopping docker container') self.cli.stop(self.container['Id']) def __get_tls_config(self): tls_config = None if self.tls_ca_cert and self.tls_client_cert and self.tls_client_key: tls_config = tls.TLSConfig( ca_cert=self.tls_ca_cert, client_cert=(self.tls_client_cert, self.tls_client_key), verify=True, ssl_version=self.tls_ssl_version, assert_hostname=self.tls_hostname ) self.docker_url = self.docker_url.replace('tcp://', 'https://') return tls_config
class Portal(object): def __init__(self): self._docker_client = APIClient() self._kill_now = False self._container_id = None self._std_in = None signal.signal(signal.SIGINT, self._exit_gracefully) signal.signal(signal.SIGTERM, self._exit_gracefully) def _cleanup(self): self._kill_now = True if (self._container_id is not None): self._docker_client.stop(self._container_id) self._docker_client.remove_container(self._container_id, v=True, force=True) def _exit_gracefully(self, signum, frame): self._cleanup() # Bad code to capture whether stdin is set or not def _capture_stdin(self): if select.select([sys.stdin, ], [], [], 0.0)[0]: self._std_in = sys.stdin.buffer.read() elif not sys.stdin.isatty(): self._std_in = sys.stdin.buffer.read() def _download_docker_image(self, command, docker_spec): docker_image_name = None if (docker_spec['image'] == 'Dockerfile'): docker_image_name = "portal/" + command try: image_data = self._docker_client.inspect_image(docker_image_name) return image_data except ImageNotFound: dockerfile = pkgutil.get_data( __name__, "commands/%s/Dockerfile" % command ).decode('utf-8') f = BytesIO(dockerfile.encode('utf-8')) for progress_dict in self._docker_client.build(fileobj=f, quiet=True, tag=docker_image_name, decode=True, rm=True): print(progress_dict) # if ('progress' in progress_dict): # print(progress_dict['progress']) else: docker_image_name = docker_spec['image'] try: image_data = self._docker_client.inspect_image(docker_image_name) return image_data except ImageNotFound: print('Pulling Docker Image...') for progress_dict in self._docker_client.pull(docker_spec['image'], stream=True, decode=True): print(progress_dict['status']) if ('progress' in progress_dict): print(progress_dict['progress']) return self._docker_client.inspect_image(docker_image_name) def _parse_args(self, spec_data, argv): parser = generate_argparse(spec_data['command'], spec_data['arguments']) cmd_options = vars(parser.parse_args(argv)) cmd_args = cmd_options['cmdargs'] for argkey in spec_data['arguments'].keys(): if (spec_data['arguments'][argkey]['shorthand'] == '*'): if (len(cmd_args) > 0): spec_data['arguments'][argkey]['value'] = cmd_args[0] if ('File' in spec_data['arguments'][argkey]['docker']): cmd_args = [os.path.join(spec_data['docker']['working_dir'], cmd_options['cmdargs'][0])] continue spec_data['arguments'][argkey]['value'] = cmd_options[spec_data['arguments'] [argkey]['shorthand']] cmd_args += merge_passthrough_vars(spec_data) return spec_data, cmd_args def _validate_spec(self, spec_data): for _, vargs in spec_data['arguments'].items(): if (vargs['argType'] == 'path' and vargs['docker'] == 'volumeBinding'): # Check if path exists # if (not os.path.isfile(vargs['value'])): #TODO: Fix! # print('Error: Path %s does not exist!' % vargs['value']) # exit(101) pass def _create_container(self, cinfo, attach_stdin): host_config = self._docker_client.create_host_config( port_bindings=cinfo.port_bindings, binds=cinfo.vol_bindings ) return self._docker_client.create_container( cinfo.container_id, command=cinfo.command, ports=cinfo.ports, environment=cinfo.environment_vars, stdin_open=attach_stdin, volumes=cinfo.volumes, # tty=True, host_config=host_config ) def _copy_artefacts_to_container(self, container_id, command_spec): def copy_file(input_path, input_name, output_path): tar_name = str(uuid.uuid4()) + '.tar' tf = tarfile.open(tar_name, mode='w') if (os.path.isfile(input_path)): tf.add(input_path, arcname=input_name) else: print("Could not find file %s " % input_path) tf.close() os.remove(tar_name) return False tf.close() with open(tar_name, 'rb') as tar_file: data = tar_file.read() self._docker_client.put_archive(container_id, output_path, data) os.remove(tar_name) for file in get_input_files(command_spec): copy_file(file['value'], file['value'], command_spec['docker']['working_dir']) home = str(Path.home()) for file in get_input_env_files(command_spec): copy_file(os.path.join(home, file['name']), file['name'], '/root') return True def _copy_artefacts_from_container(self, container_id, command_spec): def copy_file(input_file, output_path): tar_name = str(uuid.uuid4()) + '.tar' f = open(tar_name, 'wb') bits, _ = self._docker_client.get_archive( container_id, input_file) for chunk in bits: f.write(chunk) f.close() tar = tarfile.open(tar_name) tar.extractall() tar.close() os.remove(tar_name) for file in get_output_files(command_spec): copy_file(os.path.join(command_spec['docker']['working_dir'], file['value']), None) for file in get_output_env_files(command_spec): copy_file(os.path.join('/root/', file['name']), None) def run_command(self, command, argv): command_spec = None try: spec_data = pkgutil.get_data( __name__, "commands/%s/spec.toml" % command ).decode('utf-8') command_spec = toml.loads(spec_data) except FileNotFoundError: print('Command not found') return 101 self._capture_stdin() command_spec, cmd_argv = self._parse_args(command_spec, argv) self._validate_spec(command_spec) image_info = self._download_docker_image(command, command_spec['docker']) cinfo = construct_container(image_info, cmd_argv, command_spec) docker_container = self._create_container(cinfo, (self._std_in is not None)) if (len(docker_container.get('Warnings')) > 0): print('Could not start container. Warnings: %s', ' '.join(docker_container.get('Warnings'))) return 101 self._container_id = docker_container.get('Id') print('Process created in container: %s' % self._container_id) if (not self._copy_artefacts_to_container(self._container_id, command_spec)): self._cleanup() return 101 if (self._std_in is not None): s = self._docker_client.attach_socket(self._container_id, params={'stdin': 1, 'stream': 1}) os.write(s.fileno(), self._std_in) # s._sock.sendall(self._std_in) s.close() ## Attaching stdin self._docker_client.start(container=self._container_id) for log in self._docker_client.logs( container=self._container_id, stream=True, follow=True): sys.stdout.buffer.write(log) self._docker_client.wait(container=self._container_id) self._copy_artefacts_from_container(self._container_id, command_spec) self._docker_client.remove_container(container=self._container_id) return 0
class CurwDockerOperator(BaseOperator): """ Execute a command inside a docker container. Additional functionality - container auto remove - container privileged A temporary directory is created on the host and mounted into a container to allow storing files that together exceed the default disk size of 10GB in a container. The path to the mounted directory can be accessed via the environment variable ``AIRFLOW_TMP_DIR``. :param image: Docker image from which to create the container. :type image: str :param api_version: Remote API version. :type api_version: str :param command: Command to be run in the container. :type command: str or list :param cpus: Number of CPUs to assign to the container. This value gets multiplied with 1024. See https://docs.docker.com/engine/reference/run/#cpu-share-constraint :type cpus: float :param docker_url: URL of the host running the docker daemon. :type docker_url: str :param environment: Environment variables to set in the container. :type environment: dict :param force_pull: Pull the docker image on every run. :type force_pull: bool :param mem_limit: Maximum amount of memory the container can use. Either a float value, which represents the limit in bytes, or a string like ``128m`` or ``1g``. :type mem_limit: float or str :param network_mode: Network mode for the container. :type network_mode: str :param tls_ca_cert: Path to a PEM-encoded certificate authority to secure the docker connection. :type tls_ca_cert: str :param tls_client_cert: Path to the PEM-encoded certificate used to authenticate docker client. :type tls_client_cert: str :param tls_client_key: Path to the PEM-encoded key used to authenticate docker client. :type tls_client_key: str :param tls_hostname: Hostname to match against the docker server certificate or False to disable the check. :type tls_hostname: str or bool :param tls_ssl_version: Version of SSL to use when communicating with docker daemon. :type tls_ssl_version: str :param tmp_dir: Mount point inside the container to a temporary directory created on the host by the operator. The path is also made available via the environment variable ``AIRFLOW_TMP_DIR`` inside the container. :type tmp_dir: str :param user: Default user inside the docker container. :type user: int or str :param volumes: List of volumes to mount into the container, e.g. ``['/host/path:/container/path', '/host/path2:/container/path2:ro']``. :param xcom_push: Does the stdout will be pushed to the next step using XCom. The default is False. :type xcom_push: bool :param xcom_all: Push all the stdout or just the last line. The default is False (last line). :type xcom_all: bool """ template_fields = ('command', ) template_ext = ( '.sh', '.bash', ) @apply_defaults def __init__(self, image, api_version=None, command=None, cpus=1.0, docker_url='unix://var/run/docker.sock', environment=None, force_pull=False, mem_limit=None, network_mode=None, tls_ca_cert=None, tls_client_cert=None, tls_client_key=None, tls_hostname=None, tls_ssl_version=None, tmp_dir='/tmp/airflow', user=None, volumes=None, xcom_push=False, xcom_all=False, auto_remove=False, privileged=False, *args, **kwargs): super(CurwDockerOperator, self).__init__(*args, **kwargs) self.api_version = api_version self.command = command self.cpus = cpus self.docker_url = docker_url self.environment = environment or {} self.force_pull = force_pull self.image = image self.mem_limit = mem_limit self.network_mode = network_mode self.tls_ca_cert = tls_ca_cert self.tls_client_cert = tls_client_cert self.tls_client_key = tls_client_key self.tls_hostname = tls_hostname self.tls_ssl_version = tls_ssl_version self.tmp_dir = tmp_dir self.user = user self.volumes = volumes or [] self.xcom_push = xcom_push self.xcom_all = xcom_all self.auto_remove = auto_remove self.priviledged = privileged self.cli = None self.container = None def execute(self, context): logging.info('Starting docker container from image ' + self.image) tls_config = None if self.tls_ca_cert and self.tls_client_cert and self.tls_client_key: tls_config = tls.TLSConfig(ca_cert=self.tls_ca_cert, client_cert=(self.tls_client_cert, self.tls_client_key), verify=True, ssl_version=self.tls_ssl_version, assert_hostname=self.tls_hostname) self.docker_url = self.docker_url.replace('tcp://', 'https://') self.cli = Client(base_url=self.docker_url, version=self.api_version, tls=tls_config) if ':' not in self.image: image = self.image + ':latest' else: image = self.image if self.force_pull or len(self.cli.images(name=image)) == 0: logging.info('Pulling docker image ' + image) for l in self.cli.pull(image, stream=True): output = json.loads(l.decode('utf-8')) logging.info("{}".format(output['status'])) cpu_shares = int(round(self.cpus * 1024)) with TemporaryDirectory(prefix='airflowtmp') as host_tmp_dir: self.environment['AIRFLOW_TMP_DIR'] = self.tmp_dir self.volumes.append('{0}:{1}'.format(host_tmp_dir, self.tmp_dir)) cmd = self.get_command() logging.info('Creating container and running cmd:\n' + cmd) self.container = self.cli.create_container( command=cmd, cpu_shares=cpu_shares, environment=self.environment, host_config=self.cli.create_host_config( binds=self.volumes, network_mode=self.network_mode, auto_remove=self.auto_remove, privileged=self.priviledged), image=image, mem_limit=self.mem_limit, user=self.user) self.cli.start(self.container['Id']) line = '' for line in self.cli.logs(container=self.container['Id'], stream=True): logging.info("{}".format(line.strip())) exit_code = self.cli.wait(self.container['Id']) if exit_code != 0: raise AirflowException('docker container failed') if self.xcom_push: return self.cli.logs( container=self.container['Id']) if self.xcom_all else str( line.strip()) def get_command(self): if self.command is not None and self.command.strip().find('[') == 0: commands = ast.literal_eval(self.command) else: commands = self.command return commands def on_kill(self): if self.cli is not None: logging.info('Stopping docker container') self.cli.stop(self.container['Id'])