class AzureContainerInstancesOperator(BaseOperator): """ Start a container on Azure Container Instances :param ci_conn_id: connection id of a service principal which will be used to start the container instance :type ci_conn_id: str :param registry_conn_id: connection id of a user which can login to a private docker registry. If None, we assume a public registry :type registry_conn_id: str :param resource_group: name of the resource group wherein this container instance should be started :type resource_group: str :param name: name of this container instance. Please note this name has to be unique in order to run containers in parallel. :type name: str :param image: the docker image to be used :type image: str :param region: the region wherein this container instance should be started :type region: str :param environment_variables: key,value pairs containing environment variables which will be passed to the running container :type environment_variables: dict :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] :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": "******", "POSTGRES_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', 'image', 'command', 'environment_variables') @apply_defaults 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().__init__(*args, **kwargs) self.ci_conn_id = ci_conn_id self.resource_group = resource_group 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): # 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 self._ci_hook.exists(self.resource_group, self.name): raise AirflowException("Container group exists") if self.registry_conn_id: registry_hook = AzureContainerRegistryHook(self.registry_conn_id) image_registry_credentials = [ registry_hook.connection, ] else: image_registry_credentials = None environment_variables = [] for key, value in self.environment_variables.items(): 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 = [] for conn_id, account_name, share_name, mount_path, read_only in self.volumes: hook = AzureContainerVolumeHook(conn_id) mount_name = "mount-%d" % len(volumes) volumes.append( hook.get_file_volume(mount_name, share_name, account_name, 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, gpu=self.gpu)) container = Container(name=self.name, image=self.image, resources=resources, command=self.command, environment_variables=environment_variables, volume_mounts=volume_mounts) container_group = ContainerGroup( location=self.region, containers=[ container, ], image_registry_credentials=image_registry_credentials, volumes=volumes, restart_policy='Never', os_type='Linux') 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(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.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 while True: try: 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 if state in ["Running", "Terminated"]: try: 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 " "container instance, retrying...") if state == "Terminated": 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 " "by another process " "(make sure that the name is unique).") return 1 else: self.log.exception( "Exception while getting container groups") except Exception: self.log.exception("Exception while getting container groups") sleep(1) def _log_last(self, logs, last_line_logged): if logs: # determine the last line which was logged before last_line_index = 0 for i in range(len(logs) - 1, -1, -1): if logs[i] == last_line_logged: # this line is the same, hence print from i+1 last_line_index = i + 1 break # log all new ones for line in logs[last_line_index:]: 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
class TestAzureContainerInstanceHook(unittest.TestCase): def setUp(self): db.merge_conn( Connection( conn_id='azure_container_instance_test', conn_type='azure_container_instances', login='******', password='******', extra=json.dumps({'tenantId': 'tenant_id', 'subscriptionId': 'subscription_id'}) ) ) self.resources = ResourceRequirements(requests=ResourceRequests( memory_in_gb='4', cpu='1')) with patch('azure.common.credentials.ServicePrincipalCredentials.__init__', autospec=True, return_value=None): with patch('azure.mgmt.containerinstance.ContainerInstanceManagementClient'): self.testHook = AzureContainerInstanceHook(conn_id='azure_container_instance_test') @patch('azure.mgmt.containerinstance.models.ContainerGroup') @patch('azure.mgmt.containerinstance.operations.ContainerGroupsOperations.create_or_update') def test_create_or_update(self, create_or_update_mock, container_group_mock): self.testHook.create_or_update('resource_group', 'aci-test', container_group_mock) create_or_update_mock.assert_called_once_with('resource_group', 'aci-test', container_group_mock) @patch('azure.mgmt.containerinstance.operations.ContainerGroupsOperations.get') def test_get_state(self, get_state_mock): self.testHook.get_state('resource_group', 'aci-test') get_state_mock.assert_called_once_with('resource_group', 'aci-test', raw=False) @patch('azure.mgmt.containerinstance.operations.ContainerOperations.list_logs') def test_get_logs(self, list_logs_mock): expected_messages = ['log line 1\n', 'log line 2\n', 'log line 3\n'] logs = Logs(content=''.join(expected_messages)) list_logs_mock.return_value = logs logs = self.testHook.get_logs('resource_group', 'name', 'name') self.assertSequenceEqual(logs, expected_messages) @patch('azure.mgmt.containerinstance.operations.ContainerGroupsOperations.delete') def test_delete(self, delete_mock): self.testHook.delete('resource_group', 'aci-test') delete_mock.assert_called_once_with('resource_group', 'aci-test') @patch('azure.mgmt.containerinstance.operations.ContainerGroupsOperations.list_by_resource_group') def test_exists_with_existing(self, list_mock): list_mock.return_value = [ContainerGroup(os_type='Linux', containers=[Container(name='test1', image='hello-world', resources=self.resources)])] self.assertFalse(self.testHook.exists('test', 'test1')) @patch('azure.mgmt.containerinstance.operations.ContainerGroupsOperations.list_by_resource_group') def test_exists_with_not_existing(self, list_mock): list_mock.return_value = [ContainerGroup(os_type='Linux', containers=[Container(name='test1', image='hello-world', resources=self.resources)])] self.assertFalse(self.testHook.exists('test', 'not found'))