def test_reconcile_pods_empty_mutator_pod(self, mock_uuid): mock_uuid.return_value = self.static_uuid base_pod = PodGenerator( image='image1', name='name1', envs={ 'key1': 'val1' }, cmds=['/bin/command1.sh', 'arg1'], ports=[k8s.V1ContainerPort(name='port', container_port=2118)], volumes=[{ 'hostPath': { 'path': '/tmp/' }, 'name': 'example-kubernetes-test-volume1' }], volume_mounts=[{ 'mountPath': '/foo/', 'name': 'example-kubernetes-test-volume1' }], ).gen_pod() mutator_pod = None name = 'name1-' + self.static_uuid.hex base_pod.metadata.name = name result = PodGenerator.reconcile_pods(base_pod, mutator_pod) self.assertEqual(base_pod, result) mutator_pod = k8s.V1Pod() result = PodGenerator.reconcile_pods(base_pod, mutator_pod) self.assertEqual(base_pod, result)
def run_next(self, next_job): """ The run_next command will check the task_queue for any un-run jobs. It will then create a unique job-id, launch that job in the cluster, and store relevant info in the current_jobs map so we can track the job's status """ self.log.info('Kubernetes job is %s', str(next_job)) key, command, kube_executor_config = next_job dag_id, task_id, execution_date, try_number = key config_pod = self.worker_configuration.make_pod( namespace=self.namespace, worker_uuid=self.worker_uuid, pod_id=self._create_pod_id(dag_id, task_id), dag_id=self._make_safe_label_value(dag_id), task_id=self._make_safe_label_value(task_id), try_number=try_number, execution_date=self._datetime_to_label_safe_datestring( execution_date), airflow_command=command) # Reconcile the pod generated by the Operator and the Pod # generated by the .cfg file pod = PodGenerator.reconcile_pods(config_pod, kube_executor_config) self.log.debug("Kubernetes running for command %s", command) self.log.debug("Kubernetes launching image %s", pod.spec.containers[0].image) # the watcher will monitor pods, so we do not block. self.launcher.run_pod_async( pod, **self.kube_config.kube_client_request_args) self.log.debug("Kubernetes Job created!")
def create_pod_request_obj(self) -> k8s.V1Pod: """ Creates a V1Pod based on user parameters. Note that a `pod` or `pod_template_file` will supersede all other values. """ self.log.debug("Creating pod for K8sPodOperator task %s", self.task_id) if self.pod_template_file: self.log.debug("Pod template file found, will parse for base pod") pod_template = pod_generator.PodGenerator.deserialize_model_file( self.pod_template_file) else: pod_template = k8s.V1Pod(metadata=k8s.V1ObjectMeta(name="name")) pod = k8s.V1Pod(api_version="v1", kind="Pod", metadata=k8s.V1ObjectMeta( namespace=self.namespace, labels=self.labels, name=self.name, annotations=self.annotations, ), spec=k8s.V1PodSpec( node_selector=self.node_selectors, affinity=self.affinity, tolerations=self.tolerations, init_containers=self.init_containers, containers=[ k8s.V1Container( image=self.image, name="base", command=self.cmds, ports=self.ports, resources=self.k8s_resources, volume_mounts=self.volume_mounts, args=self.arguments, env=self.env_vars, env_from=self.env_from, ) ], image_pull_secrets=self.image_pull_secrets, service_account_name=self.service_account_name, host_network=self.hostnetwork, security_context=self.security_context, dns_policy=self.dnspolicy, scheduler_name=self.schedulername, restart_policy='Never', priority_class_name=self.priority_class_name, volumes=self.volumes, )) pod = PodGenerator.reconcile_pods(pod_template, pod) for secret in self.secrets: self.log.debug("Adding secret to task %s", self.task_id) pod = secret.attach_to_pod(pod) if self.do_xcom_push: self.log.debug("Adding xcom sidecar to task %s", self.task_id) pod = PodGenerator.add_xcom_sidecar(pod) return pod
def test_reconcile_pods_empty_mutator_pod(self, mock_uuid): mock_uuid.return_value = self.static_uuid path = sys.path[0] + '/tests/kubernetes/pod_generator_base_with_secrets.yaml' pod_generator = PodGenerator(pod_template_file=path, extract_xcom=True) base_pod = pod_generator.gen_pod() mutator_pod = None name = 'name1-' + self.static_uuid.hex base_pod.metadata.name = name result = PodGenerator.reconcile_pods(base_pod, mutator_pod) assert base_pod == result mutator_pod = k8s.V1Pod() result = PodGenerator.reconcile_pods(base_pod, mutator_pod) assert base_pod == result
def test_reconcile_pods(self, mock_uuid): mock_uuid.return_value = self.static_uuid path = sys.path[ 0] + '/tests/kubernetes/pod_generator_base_with_secrets.yaml' base_pod = PodGenerator(pod_template_file=path, extract_xcom=False).gen_pod() mutator_pod = k8s.V1Pod( metadata=k8s.V1ObjectMeta( name="name2", labels={"bar": "baz"}, ), spec=k8s.V1PodSpec( containers=[ k8s.V1Container( image='', name='name', command=['/bin/command2.sh', 'arg2'], volume_mounts=[ k8s.V1VolumeMount( mount_path="/foo/", name="example-kubernetes-test-volume2") ], ) ], volumes=[ k8s.V1Volume( host_path=k8s.V1HostPathVolumeSource(path="/tmp/"), name="example-kubernetes-test-volume2", ) ], ), ) result = PodGenerator.reconcile_pods(base_pod, mutator_pod) expected: k8s.V1Pod = self.expected expected.metadata.name = "name2" expected.metadata.labels['bar'] = 'baz' expected.spec.volumes = expected.spec.volumes or [] expected.spec.volumes.append( k8s.V1Volume(host_path=k8s.V1HostPathVolumeSource(path="/tmp/"), name="example-kubernetes-test-volume2")) base_container: k8s.V1Container = expected.spec.containers[0] base_container.command = ['/bin/command2.sh', 'arg2'] base_container.volume_mounts = [ k8s.V1VolumeMount(mount_path="/foo/", name="example-kubernetes-test-volume2") ] base_container.name = "name" expected.spec.containers[0] = base_container result_dict = self.k8s_client.sanitize_for_serialization(result) expected_dict = self.k8s_client.sanitize_for_serialization(expected) assert result_dict == expected_dict
def _mutate_pod_backcompat(pod): """Backwards compatible Pod Mutation Hook""" try: dummy_pod = _convert_to_airflow_pod(pod) settings.pod_mutation_hook(dummy_pod) warnings.warn( "Using `airflow.contrib.kubernetes.pod.Pod` is deprecated. " "Please use `k8s.V1Pod` instead.", DeprecationWarning, stacklevel=2 ) dummy_pod = dummy_pod.to_v1_kubernetes_pod() new_pod = PodGenerator.reconcile_pods(pod, dummy_pod) except AttributeError as e: try: settings.pod_mutation_hook(pod) return pod except AttributeError as e2: raise Exception([e, e2]) return new_pod
def test_make_pod_with_executor_config(self): self.kube_config.dags_folder = 'dags' worker_config = WorkerConfiguration(self.kube_config) config_pod = PodGenerator( image='', affinity=self.affinity_config, tolerations=self.tolerations_config, ).gen_pod() pod = worker_config.as_pod() result = PodGenerator.reconcile_pods(pod, config_pod) self.assertTrue(result.spec.affinity['podAntiAffinity'] is not None) self.assertEqual( 'app', result.spec.affinity['podAntiAffinity'] ['requiredDuringSchedulingIgnoredDuringExecution'][0] ['labelSelector']['matchExpressions'][0]['key']) self.assertEqual(2, len(result.spec.tolerations)) self.assertEqual('prod', result.spec.tolerations[1]['key'])
def build_pod_request_obj(self, context=None): """ Returns V1Pod object based on pod template file, full pod spec, and other operator parameters. The V1Pod attributes are derived (in order of precedence) from operator params, full pod spec, pod template file. """ self.log.debug("Creating pod for KubernetesPodOperator task %s", self.task_id) if self.pod_template_file: self.log.debug("Pod template file found, will parse for base pod") pod_template = pod_generator.PodGenerator.deserialize_model_file( self.pod_template_file) if self.full_pod_spec: pod_template = PodGenerator.reconcile_pods( pod_template, self.full_pod_spec) elif self.full_pod_spec: pod_template = self.full_pod_spec else: pod_template = k8s.V1Pod(metadata=k8s.V1ObjectMeta(name="name")) pod = k8s.V1Pod( api_version="v1", kind="Pod", metadata=k8s.V1ObjectMeta( namespace=self.namespace, labels=self.labels, name=self.name, annotations=self.annotations, ), spec=k8s.V1PodSpec( node_selector=self.node_selector, affinity=self.affinity, tolerations=self.tolerations, init_containers=self.init_containers, containers=[ k8s.V1Container( image=self.image, name=self.BASE_CONTAINER_NAME, command=self.cmds, ports=self.ports, image_pull_policy=self.image_pull_policy, resources=self.k8s_resources, volume_mounts=self.volume_mounts, args=self.arguments, env=self.env_vars, env_from=self.env_from, ) ], image_pull_secrets=self.image_pull_secrets, service_account_name=self.service_account_name, host_network=self.hostnetwork, security_context=self.security_context, dns_policy=self.dnspolicy, scheduler_name=self.schedulername, restart_policy='Never', priority_class_name=self.priority_class_name, volumes=self.volumes, ), ) pod = PodGenerator.reconcile_pods(pod_template, pod) if self.random_name_suffix: pod.metadata.name = PodGenerator.make_unique_pod_id( pod.metadata.name) for secret in self.secrets: self.log.debug("Adding secret to task %s", self.task_id) pod = secret.attach_to_pod(pod) if self.do_xcom_push: self.log.debug("Adding xcom sidecar to task %s", self.task_id) pod = xcom_sidecar.add_xcom_sidecar(pod) labels = self._get_ti_pod_labels(context) self.log.info("Creating pod %s with labels: %s", pod.metadata.name, labels) # Merge Pod Identifying labels with labels passed to operator pod.metadata.labels.update(labels) # Add Airflow Version to the label # And a label to identify that pod is launched by KubernetesPodOperator pod.metadata.labels.update({ 'airflow_version': airflow_version.replace('+', '-'), 'kubernetes_pod_operator': 'True', }) pod_mutation_hook(pod) return pod
def test_reconcile_pods(self): with mock.patch('uuid.uuid4') as mock_uuid: mock_uuid.return_value = '0' base_pod = PodGenerator( image='image1', name='name1', envs={ 'key1': 'val1' }, cmds=['/bin/command1.sh', 'arg1'], ports=k8s.V1ContainerPort(name='port', container_port=2118), volumes=[{ 'hostPath': { 'path': '/tmp/' }, 'name': 'example-kubernetes-test-volume1' }], volume_mounts=[{ 'mountPath': '/foo/', 'name': 'example-kubernetes-test-volume1' }], ).gen_pod() mutator_pod = PodGenerator(envs={ 'key2': 'val2' }, image='', name='name2', cmds=['/bin/command2.sh', 'arg2'], volumes=[{ 'hostPath': { 'path': '/tmp/' }, 'name': 'example-kubernetes-test-volume2' }], volume_mounts=[{ 'mountPath': '/foo/', 'name': 'example-kubernetes-test-volume2' }]).gen_pod() result = PodGenerator.reconcile_pods(base_pod, mutator_pod) result = self.k8s_client.sanitize_for_serialization(result) self.assertEqual( result, { 'apiVersion': 'v1', 'kind': 'Pod', 'metadata': { 'name': 'name2-0' }, 'spec': { 'containers': [{ 'args': [], 'command': ['/bin/command1.sh', 'arg1'], 'env': [{ 'name': 'key1', 'value': 'val1' }, { 'name': 'key2', 'value': 'val2' }], 'envFrom': [], 'image': 'image1', 'imagePullPolicy': 'IfNotPresent', 'name': 'base', 'ports': { 'containerPort': 2118, 'name': 'port', }, 'volumeMounts': [ { 'mountPath': '/foo/', 'name': 'example-kubernetes-test-volume1' }, { 'mountPath': '/foo/', 'name': 'example-kubernetes-test-volume2' } ] }], 'hostNetwork': False, 'imagePullSecrets': [], 'restartPolicy': 'Never', 'volumes': [{ 'hostPath': { 'path': '/tmp/' }, 'name': 'example-kubernetes-test-volume1' }, { 'hostPath': { 'path': '/tmp/' }, 'name': 'example-kubernetes-test-volume2' }] } })
def test_reconcile_pods(self, mock_uuid): mock_uuid.return_value = self.static_uuid base_pod = PodGenerator( image='image1', name='name1', envs={ 'key1': 'val1' }, cmds=['/bin/command1.sh', 'arg1'], ports=[k8s.V1ContainerPort(name='port', container_port=2118)], volumes=[{ 'hostPath': { 'path': '/tmp/' }, 'name': 'example-kubernetes-test-volume1' }], labels={ "foo": "bar" }, volume_mounts=[{ 'mountPath': '/foo/', 'name': 'example-kubernetes-test-volume1' }], ).gen_pod() mutator_pod = PodGenerator(envs={ 'key2': 'val2' }, image='', name='name2', labels={ "bar": "baz" }, cmds=['/bin/command2.sh', 'arg2'], volumes=[{ 'hostPath': { 'path': '/tmp/' }, 'name': 'example-kubernetes-test-volume2' }], volume_mounts=[{ 'mountPath': '/foo/', 'name': 'example-kubernetes-test-volume2' }]).gen_pod() result = PodGenerator.reconcile_pods(base_pod, mutator_pod) result = self.k8s_client.sanitize_for_serialization(result) self.assertEqual( { 'apiVersion': 'v1', 'kind': 'Pod', 'metadata': { 'name': 'name2-' + self.static_uuid.hex, 'labels': { 'foo': 'bar', "bar": "baz" } }, 'spec': { 'containers': [{ 'args': [], 'command': ['/bin/command2.sh', 'arg2'], 'env': [{ 'name': 'key1', 'value': 'val1' }, { 'name': 'key2', 'value': 'val2' }], 'envFrom': [], 'image': 'image1', 'name': 'base', 'ports': [{ 'containerPort': 2118, 'name': 'port', }], 'volumeMounts': [ { 'mountPath': '/foo/', 'name': 'example-kubernetes-test-volume1' }, { 'mountPath': '/foo/', 'name': 'example-kubernetes-test-volume2' } ] }], 'hostNetwork': False, 'imagePullSecrets': [], 'volumes': [{ 'hostPath': { 'path': '/tmp/' }, 'name': 'example-kubernetes-test-volume1' }, { 'hostPath': { 'path': '/tmp/' }, 'name': 'example-kubernetes-test-volume2' }] } }, result)
def create_pod_request_obj(self): """ Creates a V1Pod based on user parameters. Note that a `pod` or `pod_template_file` will supersede all other values. """ if self.pod_template_file: pod_template = pod_generator.PodGenerator.deserialize_model_file( self.pod_template_file) else: pod_template = k8s.V1Pod(metadata=k8s.V1ObjectMeta(name="name")) pod = pod_generator.PodGenerator( image=self.image, namespace=self.namespace, cmds=self.cmds, args=self.arguments, labels=self.labels, name=self.name, envs=self.env_vars, extract_xcom=self.do_xcom_push, image_pull_policy=self.image_pull_policy, node_selectors=self.node_selectors, annotations=self.annotations, affinity=self.affinity, image_pull_secrets=self.image_pull_secrets, service_account_name=self.service_account_name, hostnetwork=self.hostnetwork, tolerations=self.tolerations, security_context=self.security_context, dnspolicy=self.dnspolicy, init_containers=self.init_containers, restart_policy='Never', schedulername=self.schedulername, priority_class_name=self.priority_class_name, ).gen_pod() # noinspection PyTypeChecker pod = append_to_pod( pod, self.pod_runtime_info_envs + # type: ignore self.ports + # type: ignore self.resources + # type: ignore self.secrets + # type: ignore self.volumes + # type: ignore self.volume_mounts # type: ignore ) env_from = pod.spec.containers[0].env_from or [] for configmap in self.configmaps: env_from.append( k8s.V1EnvFromSource(config_map_ref=k8s.V1ConfigMapEnvSource( name=configmap))) pod.spec.containers[0].env_from = env_from if self.full_pod_spec: pod_template = PodGenerator.reconcile_pods(pod_template, self.full_pod_spec) pod = PodGenerator.reconcile_pods(pod_template, pod) # if self.do_xcom_push: # pod = PodGenerator.add_sidecar(pod) return pod