Skip to content

Commit

Permalink
[AIRFLOW-4564] AzureContainerInstance bugfixes and improvements (apac…
Browse files Browse the repository at this point in the history 
…he#5319)

(cherry picked from commit 722379a)
  • Loading branch information
@Trollgeir @ashb
Trollgeir authored and ashb committed Jul 1, 2019
1 parent d06a501 commit bd8b1e1
Show file tree
Hide file tree
Showing 5 changed files with 220 additions and 103 deletions.
35 changes: 23 additions & 12 deletions airflow/contrib/hooks/azure_container_instance_hook.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
from azure.common.credentials import ServicePrincipalCredentials

from azure.mgmt.containerinstance import ContainerInstanceManagementClient
from zope.deprecation import deprecation


class AzureContainerInstanceHook(BaseHook):
Expand Down Expand Up @@ -87,6 +88,7 @@ def create_or_update(self, resource_group, name, container_group):
name,
container_group)

@deprecation.deprecate("get_state_exitcode_details() is deprecated. Related method is get_state()")
def get_state_exitcode_details(self, resource_group, name):
"""
Get the state and exitcode of a container group
Expand All @@ -99,17 +101,11 @@ def get_state_exitcode_details(self, resource_group, name):
If the exitcode is unknown 0 is returned.
:rtype: tuple(state,exitcode,details)
"""
current_state = self._get_instance_view(resource_group, name).current_state
return (current_state.state,
current_state.exit_code,
current_state.detail_status)

def _get_instance_view(self, resource_group, name):
response = self.connection.container_groups.get(resource_group,
name,
raw=False)
return response.containers[0].instance_view
cg_state = self.get_state(resource_group, name)
c_state = cg_state.containers[0].instance_view.current_state
return (c_state.state, c_state.exit_code, c_state.detail_status)

@deprecation.deprecate("get_messages() is deprecated. Related method is get_state()")
def get_messages(self, resource_group, name):
"""
Get the messages of a container group
Expand All @@ -121,10 +117,25 @@ def get_messages(self, resource_group, name):
:return: A list of the event messages
:rtype: list[str]
"""
instance_view = self._get_instance_view(resource_group, name)

cg_state = self.get_state(resource_group, name)
instance_view = cg_state.containers[0].instance_view
return [event.message for event in instance_view.events]

def get_state(self, resource_group, name):
"""
Get the state of a container group
:param resource_group: the name of the resource group
:type resource_group: str
:param name: the name of the container group
:type name: str
:return: ContainerGroup
:rtype: ~azure.mgmt.containerinstance.models.ContainerGroup
"""
return self.connection.container_groups.get(resource_group,
name,
raw=False)

def get_logs(self, resource_group, name, tail=1000):
"""
Get the tail from logs of a container group
Expand Down
184 changes: 133 additions & 51 deletions airflow/contrib/operators/azure_container_instances_operator.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,8 +18,9 @@
# under the License.

from collections import namedtuple
from typing import Sequence, Dict
from time import sleep
from typing import Dict, Sequence
import re

from airflow.contrib.hooks.azure_container_instance_hook import AzureContainerInstanceHook
from airflow.contrib.hooks.azure_container_registry_hook import AzureContainerRegistryHook
Expand All @@ -43,7 +44,9 @@
['conn_id', 'account_name', 'share_name', 'mount_path', 'read_only'],
)


DEFAULT_ENVIRONMENT_VARIABLES = {} # type: Dict[str, str]
DEFAULT_SECURED_VARIABLES = [] # type: Sequence[str]
DEFAULT_VOLUMES = [] # type: Sequence[Volume]
DEFAULT_MEMORY_IN_GB = 2.0
DEFAULT_CPU = 1.0
Expand Down Expand Up @@ -72,66 +75,101 @@ class AzureContainerInstancesOperator(BaseOperator):
:param environment_variables: key,value pairs containing environment
variables which will be passed to the running container
:type environment_variables: dict
:param volumes: list of volumes to be mounted to the container.
:param secured_variables: names of environmental variables that should not
be exposed outside the container (typically passwords).
:type secured_variables: [str]
:param volumes: list of ``Volume`` tuples to be mounted to the container.
Currently only Azure Fileshares are supported.
:type volumes: list[<conn_id, account_name, share_name, mount_path, read_only>]
:param memory_in_gb: the amount of memory to allocate to this container
:type memory_in_gb: double
:param cpu: the number of cpus to allocate to this container
:type cpu: double
:param gpu: GPU Resource for the container.
:type gpu: azure.mgmt.containerinstance.models.GpuResource
:param command: the command to run inside the container
:type command: str
:Example:
>>> a = AzureContainerInstancesOperator(
'azure_service_principal',
'azure_registry_user',
'my-resource-group',
'my-container-name-{{ ds }}',
'myprivateregistry.azurecr.io/my_container:latest',
'westeurope',
{'EXECUTION_DATE': '{{ ds }}'},
[('azure_wasb_conn_id',
'my_storage_container',
'my_fileshare',
'/input-data',
True),],
memory_in_gb=14.0,
cpu=4.0,
command='python /app/myfile.py',
task_id='start_container'
)
:type command: [str]
:param container_timeout: max time allowed for the execution of
the container instance.
:type container_timeout: datetime.timedelta
**Example**::
AzureContainerInstancesOperator(
"azure_service_principal",
"azure_registry_user",
"my-resource-group",
"my-container-name-{{ ds }}",
"myprivateregistry.azurecr.io/my_container:latest",
"westeurope",
{"MODEL_PATH": "my_value",
"POSTGRES_LOGIN": "{{ macros.connection('postgres_default').login }}"
"POSTGRES_PASSWORD": "{{ macros.connection('postgres_default').password }}",
"JOB_GUID": "{{ ti.xcom_pull(task_ids='task1', key='guid') }}" },
['POSTGRES_PASSWORD'],
[("azure_wasb_conn_id",
"my_storage_container",
"my_fileshare",
"/input-data",
True),],
memory_in_gb=14.0,
cpu=4.0,
gpu=GpuResource(count=1, sku='K80'),
command=["/bin/echo", "world"],
container_timeout=timedelta(hours=2),
task_id="start_container"
)
"""

template_fields = ('name', 'environment_variables')
template_fields = ('name', 'image', 'command', 'environment_variables')

@apply_defaults
def __init__(self, ci_conn_id, registry_conn_id, resource_group, name, image, region,
environment_variables=None, volumes=None, memory_in_gb=None, cpu=None,
command=None, remove_on_error=True, fail_if_exists=True, *args, **kwargs):
def __init__(self,
ci_conn_id,
registry_conn_id,
resource_group,
name,
image,
region,
environment_variables=None,
secured_variables=None,
volumes=None,
memory_in_gb=None,
cpu=None,
gpu=None,
command=None,
remove_on_error=True,
fail_if_exists=True,
*args,
**kwargs):
super(AzureContainerInstancesOperator, self).__init__(*args, **kwargs)

self.ci_conn_id = ci_conn_id
self.resource_group = resource_group
self.name = name
self.name = self._check_name(name)
self.image = image
self.region = region
self.registry_conn_id = registry_conn_id
self.environment_variables = environment_variables or DEFAULT_ENVIRONMENT_VARIABLES
self.secured_variables = secured_variables or DEFAULT_SECURED_VARIABLES
self.volumes = volumes or DEFAULT_VOLUMES
self.memory_in_gb = memory_in_gb or DEFAULT_MEMORY_IN_GB
self.cpu = cpu or DEFAULT_CPU
self.gpu = gpu
self.command = command
self.remove_on_error = remove_on_error
self.fail_if_exists = fail_if_exists
self._ci_hook = None

def execute(self, context):
ci_hook = AzureContainerInstanceHook(self.ci_conn_id)
# Check name again in case it was templated.
self._check_name(self.name)

self._ci_hook = AzureContainerInstanceHook(self.ci_conn_id)

if self.fail_if_exists:
self.log.info("Testing if container group already exists")
if ci_hook.exists(self.resource_group, self.name):
if self._ci_hook.exists(self.resource_group, self.name):
raise AirflowException("Container group exists")

if self.registry_conn_id:
Expand All @@ -142,7 +180,11 @@ def execute(self, context):

environment_variables = []
for key, value in self.environment_variables.items():
environment_variables.append(EnvironmentVariable(key, value))
if key in self.secured_variables:
e = EnvironmentVariable(name=key, secure_value=value)
else:
e = EnvironmentVariable(name=key, value=value)
environment_variables.append(e)

volumes = []
volume_mounts = []
Expand All @@ -154,16 +196,22 @@ def execute(self, context):
share_name,
account_name,
read_only))
volume_mounts.append(VolumeMount(mount_name, mount_path, read_only))
volume_mounts.append(VolumeMount(name=mount_name,
mount_path=mount_path,
read_only=read_only))

exit_code = 1
try:
self.log.info("Starting container group with %.1f cpu %.1f mem",
self.cpu, self.memory_in_gb)
if self.gpu:
self.log.info("GPU count: %.1f, GPU SKU: %s",
self.gpu.count, self.gpu.sku)

resources = ResourceRequirements(requests=ResourceRequests(
memory_in_gb=self.memory_in_gb,
cpu=self.cpu))
cpu=self.cpu,
gpu=self.gpu))

container = Container(
name=self.name,
Expand All @@ -181,46 +229,65 @@ def execute(self, context):
restart_policy='Never',
os_type='Linux')

ci_hook.create_or_update(self.resource_group, self.name, container_group)
self._ci_hook.create_or_update(self.resource_group, self.name, container_group)

self.log.info("Container group started %s/%s", self.resource_group, self.name)

exit_code = self._monitor_logging(ci_hook, self.resource_group, self.name)
exit_code = self._monitor_logging(self._ci_hook, self.resource_group, self.name)

self.log.info("Container had exit code: %s", exit_code)
if exit_code != 0:
raise AirflowException("Container had a non-zero exit code, %s"
% exit_code)
return exit_code

except CloudError:
self.log.exception("Could not start container group")
raise AirflowException("Could not start container group")

finally:
if exit_code == 0 or self.remove_on_error:
self.log.info("Deleting container group")
try:
ci_hook.delete(self.resource_group, self.name)
except Exception:
self.log.exception("Could not delete container group")
self.on_kill()

def on_kill(self):
if self.remove_on_error:
self.log.info("Deleting container group")
try:
self._ci_hook.delete(self.resource_group, self.name)
except Exception:
self.log.exception("Could not delete container group")

def _monitor_logging(self, ci_hook, resource_group, name):
last_state = None
last_message_logged = None
last_line_logged = None
for _ in range(43200): # roughly 12 hours

while True:
try:
state, exit_code, detail_status = ci_hook.get_state_exitcode_details(resource_group, name)
cg_state = self._ci_hook.get_state(resource_group, name)
instance_view = cg_state.containers[0].instance_view

# If there is no instance view, we show the provisioning state
if instance_view is not None:
c_state = instance_view.current_state
state, exit_code, detail_status = (c_state.state,
c_state.exit_code,
c_state.detail_status)

messages = [event.message for event in instance_view.events]
last_message_logged = self._log_last(messages, last_message_logged)
else:
state = cg_state.provisioning_state
exit_code = 0
detail_status = "Provisioning"

if state != last_state:
self.log.info("Container group state changed to %s", state)
last_state = state

messages = ci_hook.get_messages(resource_group, name)
last_message_logged = self._log_last(messages, last_message_logged)

if state in ["Running", "Terminated"]:
try:
logs = ci_hook.get_logs(resource_group, name)
logs = self._ci_hook.get_logs(resource_group, name)
last_line_logged = self._log_last(logs, last_line_logged)
except CloudError:
self.log.exception("Exception while getting logs from "
Expand All @@ -230,6 +297,12 @@ def _monitor_logging(self, ci_hook, resource_group, name):
self.log.info("Container exited with detail_status %s", detail_status)
return exit_code

if state == "Failed":
self.log.info("Azure provision failure")
return 1

except AirflowTaskTimeout:
raise
except CloudError as err:
if 'ResourceNotFound' in str(err):
self.log.warning("ResourceNotFound, container is probably removed "
Expand All @@ -241,10 +314,7 @@ def _monitor_logging(self, ci_hook, resource_group, name):
except Exception:
self.log.exception("Exception while getting container groups")

sleep(1)

# no return -> hence still running
raise AirflowTaskTimeout("Did not complete on time")
sleep(1)

def _log_last(self, logs, last_line_logged):
if logs:
Expand All @@ -261,3 +331,15 @@ def _log_last(self, logs, last_line_logged):
self.log.info(line.rstrip())

return logs[-1]

@staticmethod
def _check_name(name):
if '{{' in name:
# Let macros pass as they cannot be checked at construction time
return name
regex_check = re.match("[a-z0-9]([-a-z0-9]*[a-z0-9])?", name)
if regex_check is None or regex_check.group() != name:
raise AirflowException('ACI name must match regex [a-z0-9]([-a-z0-9]*[a-z0-9])? (like "my-name")')
if len(name) > 63:
raise AirflowException('ACI name cannot be longer than 63 characters')
return name
8 changes: 4 additions & 4 deletions setup.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -141,12 +141,12 @@ def write_version(filename=os.path.join(*['airflow',
atlas = ['atlasclient>=0.1.2']
azure_blob_storage = ['azure-storage>=0.34.0']
azure_data_lake = [
'azure-mgmt-resource==1.2.2',
'azure-mgmt-datalake-store==0.4.0',
'azure-datalake-store==0.0.19'
'azure-mgmt-resource>=2.2.0',
'azure-mgmt-datalake-store>=0.5.0',
'azure-datalake-store>=0.0.45'
]
azure_cosmos = ['azure-cosmos>=3.0.1']
azure_container_instances = ['azure-mgmt-containerinstance']
azure_container_instances = ['azure-mgmt-containerinstance>=1.5.0']
cassandra = ['cassandra-driver>=3.13.0']
celery = [
'celery~=4.3',
Expand Down
Loading

0 comments on commit bd8b1e1

Please sign in to comment.