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'])
