def test_pending_pod_timeout(self, mock_kubescheduler, mock_get_kube_client, mock_kubernetes_job_watcher): mock_delete_pod = mock_kubescheduler.return_value.delete_pod mock_kube_client = mock_get_kube_client.return_value now = timezone.utcnow() pending_pods = [ k8s.V1Pod(metadata=k8s.V1ObjectMeta( name="foo60", labels={"airflow-worker": "123"}, creation_timestamp=now - timedelta(seconds=60), namespace="mynamespace", )), k8s.V1Pod(metadata=k8s.V1ObjectMeta( name="foo90", labels={"airflow-worker": "123"}, creation_timestamp=now - timedelta(seconds=90), namespace="mynamespace", )), ] mock_kube_client.list_namespaced_pod.return_value.items = pending_pods config = { ('kubernetes', 'namespace'): 'mynamespace', ('kubernetes', 'worker_pods_pending_timeout'): '75', ('kubernetes', 'worker_pods_pending_timeout_batch_size'): '5', ('kubernetes', 'kube_client_request_args'): '{"sentinel": "foo"}', } with conf_vars(config): executor = KubernetesExecutor() executor.job_id = "123" executor.start() assert 1 == len(executor.event_scheduler.queue) executor._check_worker_pods_pending_timeout() mock_kube_client.list_namespaced_pod.assert_called_once_with( 'mynamespace', field_selector='status.phase=Pending', label_selector='airflow-worker=123', limit=5, sentinel='foo', ) mock_delete_pod.assert_called_once_with('foo90', 'mynamespace')
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. """ if self.pod_template_file: pod = pod_generator.PodGenerator.deserialize_model_file( self.pod_template_file) elif not self.pod: 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, )) else: pod = self.pod for secret in self.secrets: pod = secret.attach_to_pod(pod) if self.do_xcom_push: from airflow.kubernetes.pod_generator import PodGenerator pod = PodGenerator.add_xcom_sidecar(pod) return pod
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 test_construct_pod_empty_execuctor_config(self, mock_uuid): mock_uuid.return_value = self.static_uuid worker_config = k8s.V1Pod(spec=k8s.V1PodSpec(containers=[ k8s.V1Container(name='', resources=k8s.V1ResourceRequirements(limits={ 'cpu': '1m', 'memory': '1G' })) ])) executor_config = None result = PodGenerator.construct_pod( self.dag_id, self.task_id, 'pod_id', self.try_number, "kube_image", self.execution_date, ['command'], executor_config, worker_config, 'namespace', 'uuid', ) sanitized_result = self.k8s_client.sanitize_for_serialization(result) self.assertEqual( { 'apiVersion': 'v1', 'kind': 'Pod', 'metadata': self.metadata, 'spec': { 'containers': [{ 'args': [], 'command': ['command'], 'env': [], 'envFrom': [], 'name': 'base', 'image': 'kube_image', 'ports': [], 'resources': { 'limits': { 'cpu': '1m', 'memory': '1G' } }, 'volumeMounts': [] }], 'hostNetwork': False, 'imagePullSecrets': [], 'volumes': [] } }, sanitized_result)
def construct_pod( # pylint: disable=too-many-arguments dag_id: str, task_id: str, pod_id: str, try_number: int, kube_image: str, date: datetime.datetime, command: List[str], pod_override_object: Optional[k8s.V1Pod], base_worker_pod: k8s.V1Pod, namespace: str, worker_uuid: str) -> k8s.V1Pod: """ Construct a pod by gathering and consolidating the configuration from 3 places: - airflow.cfg - executor_config - dynamic arguments """ try: image = pod_override_object.spec.containers[ 0].image # type: ignore if not image: image = kube_image except Exception: # pylint: disable=W0703 image = kube_image dynamic_pod = k8s.V1Pod(metadata=k8s.V1ObjectMeta( namespace=namespace, annotations={ 'dag_id': dag_id, 'task_id': task_id, 'execution_date': date.isoformat(), 'try_number': str(try_number), }, name=PodGenerator.make_unique_pod_id(pod_id), labels={ 'airflow-worker': worker_uuid, 'dag_id': dag_id, 'task_id': task_id, 'execution_date': datetime_to_label_safe_datestring(date), 'try_number': str(try_number), 'airflow_version': airflow_version.replace('+', '-'), 'kubernetes_executor': 'True', }), spec=k8s.V1PodSpec(containers=[ k8s.V1Container( name="base", command=command, image=image, ) ])) # Reconcile the pods starting with the first chronologically, # Pod from the pod_template_File -> Pod from executor_config arg -> Pod from the K8s executor pod_list = [base_worker_pod, pod_override_object, dynamic_pod] return reduce(PodGenerator.reconcile_pods, pod_list)
def test_try_adopt_task_instances(self, mock_adopt_completed_pods, mock_adopt_launched_task): executor = self.kubernetes_executor executor.scheduler_job_id = "10" ti_key = annotations_to_key({ 'dag_id': 'dag', 'execution_date': datetime.utcnow().isoformat(), 'task_id': 'task', 'try_number': '1', }) mock_ti = mock.MagicMock(queued_by_job_id="1", external_executor_id="1", key=ti_key) pod = k8s.V1Pod(metadata=k8s.V1ObjectMeta(name="foo")) mock_kube_client = mock.MagicMock() mock_kube_client.list_namespaced_pod.return_value.items = [pod] executor.kube_client = mock_kube_client # First adoption reset_tis = executor.try_adopt_task_instances([mock_ti]) mock_kube_client.list_namespaced_pod.assert_called_once_with( namespace='default', label_selector='airflow-worker=1') mock_adopt_launched_task.assert_called_once_with( mock_kube_client, pod, {ti_key: mock_ti}) mock_adopt_completed_pods.assert_called_once() assert reset_tis == [mock_ti ] # assume failure adopting when checking return # Second adoption (queued_by_job_id and external_executor_id no longer match) mock_kube_client.reset_mock() mock_adopt_launched_task.reset_mock() mock_adopt_completed_pods.reset_mock() mock_ti.queued_by_job_id = "10" # scheduler_job would have updated this after the first adoption executor.scheduler_job_id = "20" # assume success adopting when checking return, `adopt_launched_task` pops `ti_key` from `pod_ids` mock_adopt_launched_task.side_effect = lambda client, pod, pod_ids: pod_ids.pop( ti_key) reset_tis = executor.try_adopt_task_instances([mock_ti]) mock_kube_client.list_namespaced_pod.assert_called_once_with( namespace='default', label_selector='airflow-worker=10') mock_adopt_launched_task.assert_called_once( ) # Won't check args this time around as they get mutated mock_adopt_completed_pods.assert_called_once() assert reset_tis == [ ] # This time our return is empty - no TIs to reset
def __init__(self, dag, version, release_stream, latest_release, platform, profile, default_args): self.exec_config = { "pod_override": k8s.V1Pod( spec=k8s.V1PodSpec( containers=[ k8s.V1Container( name="base", image="quay.io/keithwhitley4/airflow-ansible:2.1.0", image_pull_policy="Always", env=[ kubeconfig.get_kubeadmin_password( version, platform, profile) ], volume_mounts=[ kubeconfig.get_kubeconfig_volume_mount()] ) ], volumes=[kubeconfig.get_kubeconfig_volume( version, platform, profile)] ) ) } # General DAG Configuration self.dag = dag self.platform = platform # e.g. aws self.version = version # e.g. stable/.next/.future self.release_stream = release_stream self.latest_release = latest_release # latest relase from the release stream self.profile = profile # e.g. default/ovn self.default_args = default_args # Airflow Variables self.SNAPPY_DATA_SERVER_URL = Variable.get("SNAPPY_DATA_SERVER_URL") self.SNAPPY_DATA_SERVER_USERNAME = Variable.get("SNAPPY_DATA_SERVER_USERNAME") self.SNAPPY_DATA_SERVER_PASSWORD = Variable.get("SNAPPY_DATA_SERVER_PASSWORD") # Specific Task Configuration self.vars = var_loader.build_task_vars( task="utils", version=version, platform=platform, profile=profile) self.git_name=self._git_name() self.env = { "OPENSHIFT_CLIENT_LOCATION": self.latest_release["openshift_client_location"], "SNAPPY_DATA_SERVER_URL": self.SNAPPY_DATA_SERVER_URL, "SNAPPY_DATA_SERVER_USERNAME": self.SNAPPY_DATA_SERVER_USERNAME, "SNAPPY_DATA_SERVER_PASSWORD": self.SNAPPY_DATA_SERVER_PASSWORD, "SNAPPY_USER_FOLDER": self.git_name }
def to_v1_kubernetes_pod(self): """ Convert to support k8s V1Pod :return: k8s.V1Pod """ import kubernetes.client.models as k8s meta = k8s.V1ObjectMeta( labels=self.labels, name=self.name, namespace=self.namespace, ) spec = k8s.V1PodSpec( init_containers=self.init_containers, containers=[ k8s.V1Container( image=self.image, command=self.cmds, name="base", env=[k8s.V1EnvVar(name=key, value=val) for key, val in self.envs.items()], args=self.args, image_pull_policy=self.image_pull_policy, ) ], image_pull_secrets=self.image_pull_secrets, service_account_name=self.service_account_name, dns_policy=self.dnspolicy, host_network=self.hostnetwork, tolerations=self.tolerations, affinity=self.affinity, security_context=self.security_context, ) pod = k8s.V1Pod( spec=spec, metadata=meta, ) for port in _extract_ports(self.ports): pod = port.attach_to_pod(pod) volumes = _extract_volumes(self.volumes) for volume in volumes: pod = volume.attach_to_pod(pod) for volume_mount in _extract_volume_mounts(self.volume_mounts): pod = volume_mount.attach_to_pod(pod) for secret in self.secrets: pod = secret.attach_to_pod(pod) for runtime_info in self.pod_runtime_info_envs: pod = runtime_info.attach_to_pod(pod) pod = _extract_resources(self.resources).attach_to_pod(pod) return pod
def test_construct_pod(self, mock_uuid): path = sys.path[ 0] + '/tests/kubernetes/pod_generator_base_with_secrets.yaml' worker_config = PodGenerator.deserialize_model_file(path) mock_uuid.return_value = self.static_uuid executor_config = k8s.V1Pod(spec=k8s.V1PodSpec(containers=[ k8s.V1Container(name='', resources=k8s.V1ResourceRequirements(limits={ 'cpu': '1m', 'memory': '1G' })) ])) result = PodGenerator.construct_pod( dag_id=self.dag_id, task_id=self.task_id, pod_id='pod_id', kube_image='airflow_image', try_number=self.try_number, date=self.execution_date, args=['command'], pod_override_object=executor_config, base_worker_pod=worker_config, namespace='test_namespace', scheduler_job_id='uuid', ) expected = self.expected expected.metadata.labels = self.labels expected.metadata.labels['app'] = 'myapp' expected.metadata.annotations = self.annotations expected.metadata.name = 'pod_id.' + self.static_uuid.hex expected.metadata.namespace = 'test_namespace' expected.spec.containers[0].args = ['command'] expected.spec.containers[0].image = 'airflow_image' expected.spec.containers[0].resources = { 'limits': { 'cpu': '1m', 'memory': '1G' } } expected.spec.containers[0].env.append( k8s.V1EnvVar( name="AIRFLOW_IS_K8S_EXECUTOR_POD", value='True', )) result_dict = self.k8s_client.sanitize_for_serialization(result) expected_dict = self.k8s_client.sanitize_for_serialization( self.expected) assert expected_dict == result_dict
def test_full_pod_spec(self): pod_spec = k8s.V1Pod( metadata=k8s.V1ObjectMeta(labels={ "foo": "bar", "fizz": "buzz" }, namespace="default", name="test-pod"), spec=k8s.V1PodSpec( containers=[ k8s.V1Container( name="base", image="perl", command=["/bin/bash"], args=[ "-c", 'echo {\\"hello\\" : \\"world\\"} | cat > /airflow/xcom/return.json' ], env=[k8s.V1EnvVar(name="env_name", value="value")], ) ], restart_policy="Never", ), ) k = KubernetesPodOperator( task_id="task" + self.get_current_task_name(), in_cluster=False, full_pod_spec=pod_spec, do_xcom_push=True, is_delete_operator_pod=False, ) context = create_context(k) result = k.execute(context) assert result is not None assert k.pod.metadata.labels == { 'fizz': 'buzz', 'foo': 'bar', 'airflow_version': mock.ANY, 'dag_id': 'dag', 'run_id': 'manual__2016-01-01T0100000100-da4d1ce7b', 'kubernetes_pod_operator': 'True', 'task_id': mock.ANY, 'try_number': '1', 'already_checked': 'True', } assert k.pod.spec.containers[0].env == [ k8s.V1EnvVar(name="env_name", value="value") ] assert result == {"hello": "world"}
def test_full_pod_spec(self): pod_spec = k8s.V1Pod( metadata=k8s.V1ObjectMeta(name="hello", labels={"foo": "bar"}, namespace="mynamespace"), spec=k8s.V1PodSpec(containers=[ k8s.V1Container( name="base", image="ubuntu:16.04", command=["something"], ) ]), ) k = KubernetesPodOperator( task_id="task", in_cluster=False, do_xcom_push=False, cluster_context="default", full_pod_spec=pod_spec, ) pod = k.create_pod_request_obj() assert pod.metadata.name == pod_spec.metadata.name assert pod.metadata.labels == pod_spec.metadata.labels assert pod.metadata.namespace == pod_spec.metadata.namespace assert pod.spec.containers[0].image == pod_spec.spec.containers[ 0].image assert pod.spec.containers[0].command == pod_spec.spec.containers[ 0].command # kwargs take precedence, however image = "some.custom.image:andtag" name_base = "world" k = KubernetesPodOperator( task_id="task", in_cluster=False, do_xcom_push=False, cluster_context="default", full_pod_spec=pod_spec, name=name_base, image=image, ) pod = k.create_pod_request_obj() # make sure the kwargs takes precedence (and that name is randomized) assert pod.metadata.name.startswith(name_base) assert pod.metadata.name != name_base assert pod.spec.containers[0].image == image
def test_not_adopt_unassigned_task(self, mock_kube_client): """ We should not adopt any tasks that were not assigned by the scheduler. This ensures that there is no contention over pod management. """ executor = self.kubernetes_executor executor.scheduler_job_id = "modified" pod_ids = {"foobar": {}} pod = k8s.V1Pod( metadata=k8s.V1ObjectMeta( name="foo", labels={"airflow-worker": "bar", "dag_id": "dag", "task_id": "task"} ) ) executor.adopt_launched_task(mock_kube_client, pod=pod, pod_ids=pod_ids) assert not mock_kube_client.patch_namespaced_pod.called assert pod_ids == {"foobar": {}}
def test_invalid_executor_config(self, mock_get_kube_client, mock_kubernetes_job_watcher): executor = self.kubernetes_executor executor.start() assert executor.event_buffer == {} executor.execute_async( key=('dag', 'task', datetime.utcnow(), 1), queue=None, command=['airflow', 'tasks', 'run', 'true', 'some_parameter'], executor_config=k8s.V1Pod( spec=k8s.V1PodSpec( containers=[k8s.V1Container(name="base", image="myimage", image_pull_policy="Always")] ) ), ) assert list(executor.event_buffer.values())[0][1] == "Invalid executor_config passed"
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 get_default_executor_config(dag_config: DagConfig, executor_image='airflow-ansible'): return { "pod_override": k8s.V1Pod( spec=k8s.V1PodSpec( containers=[ k8s.V1Container( name="base", image=f"{dag_config.executor_image['repository']}/{executor_image}:{dag_config.executor_image['tag']}", image_pull_policy="Always", volume_mounts=[ get_empty_dir_volume_mount()] ) ], volumes=[get_empty_dir_volume_mount()] ) ) }
def test_add_custom_label(self): from kubernetes.client import models as k8s pod = PodGenerator.construct_pod( namespace="test", worker_uuid="test", pod_id="test", dag_id="test", task_id="test", try_number=1, date="23-07-2020", command="test", kube_executor_config=None, worker_config=k8s.V1Pod(metadata=k8s.V1ObjectMeta( labels={"airflow-test": "airflow-task-pod"}, annotations={"my.annotation": "foo"}))) self.assertIn("airflow-test", pod.metadata.labels) self.assertIn("my.annotation", pod.metadata.annotations)
def test_try_adopt_task_instances(self, mock_adopt_completed_pods, mock_adopt_launched_task): executor = self.kubernetes_executor executor.scheduler_job_id = "10" mock_ti = mock.MagicMock(queued_by_job_id="1", external_executor_id="1", dag_id="dag", task_id="task") pod = k8s.V1Pod(metadata=k8s.V1ObjectMeta(name="foo", labels={ "dag_id": "dag", "task_id": "task" })) pod_id = create_pod_id(dag_id="dag", task_id="task") mock_kube_client = mock.MagicMock() mock_kube_client.list_namespaced_pod.return_value.items = [pod] executor.kube_client = mock_kube_client # First adoption executor.try_adopt_task_instances([mock_ti]) mock_kube_client.list_namespaced_pod.assert_called_once_with( namespace='default', label_selector='airflow-worker=1') mock_adopt_launched_task.assert_called_once_with( mock_kube_client, pod, {pod_id: mock_ti}) mock_adopt_completed_pods.assert_called_once() # We aren't checking the return value of `try_adopt_task_instances` because it relies on # `adopt_launched_task` mutating its arg. This should be refactored, but not right now. # Second adoption (queued_by_job_id and external_executor_id no longer match) mock_kube_client.reset_mock() mock_adopt_launched_task.reset_mock() mock_adopt_completed_pods.reset_mock() mock_ti.queued_by_job_id = "10" # scheduler_job would have updated this after the first adoption executor.scheduler_job_id = "20" executor.try_adopt_task_instances([mock_ti]) mock_kube_client.list_namespaced_pod.assert_called_once_with( namespace='default', label_selector='airflow-worker=10') mock_adopt_launched_task.assert_called_once_with( mock_kube_client, pod, {pod_id: mock_ti}) mock_adopt_completed_pods.assert_called_once()
def mock_kubernetes_read_namespaced_pod(*_args, **_kwargs): """ Represents the mocked output of kubernetes.client.read_namespaced_pod """ return models.V1Pod( metadata=models.V1ObjectMeta( namespace="default", name="gordo-test-pod-name-1234", labels={"app": "gordo-model-builder"}, ), status=models.V1PodStatus(phase="Running"), spec=models.V1PodSpec(containers=[ models.V1Container( name="some-generated-test-container-name", env=[ models.V1EnvVar(name="MACHINE_NAME", value="test-machine-name") ], ) ]), )
def get_executor_config_with_cluster_access(dag_config: DagConfig, release: OpenshiftRelease, executor_image="airflow-ansible"): return { "pod_override": k8s.V1Pod( spec=k8s.V1PodSpec( containers=[ k8s.V1Container( name="base", image=f"{dag_config.executor_image['repository']}/{executor_image}:{dag_config.executor_image['tag']}", image_pull_policy="Always", env=[ get_kubeadmin_password(release) ], volume_mounts=[ get_kubeconfig_volume_mount()] ) ], volumes=[get_kubeconfig_volume(release)] ) ) }
def test_adopt_launched_task(self, mock_kube_client): executor = self.kubernetes_executor executor.scheduler_job_id = "modified" pod_ids = {"dagtask": {}} pod = k8s.V1Pod( metadata=k8s.V1ObjectMeta( name="foo", labels={"airflow-worker": "bar", "dag_id": "dag", "task_id": "task"} ) ) executor.adopt_launched_task(mock_kube_client, pod=pod, pod_ids=pod_ids) assert mock_kube_client.patch_namespaced_pod.call_args[1] == { 'body': { 'metadata': { 'labels': {'airflow-worker': 'modified', 'dag_id': 'dag', 'task_id': 'task'}, 'name': 'foo', } }, 'name': 'foo', 'namespace': None, } assert pod_ids == {}
def _new_pod(self, immortalcontainer): """Returns the pod definition to create the pod for an ImmortalContainer""" labels = dict(controller=immortalcontainer['metadata']['name']) return models.V1Pod( metadata=models.V1ObjectMeta( name=immortalcontainer['metadata']['name'] + "-immortalpod", labels=labels, namespace=immortalcontainer['metadata']['namespace'], owner_references=[ models.V1OwnerReference( api_version=self.custom_group + "/" + self.custom_version, controller=True, kind=self.custom_kind, name=immortalcontainer['metadata']['name'], uid=immortalcontainer['metadata']['uid']) ]), spec=models.V1PodSpec(containers=[ models.V1Container(name="acontainer", image=immortalcontainer['spec']['image']) ]))
def __init__(self, dag, version, release_stream, platform, profile): self.exec_config = { "pod_override": k8s.V1Pod(spec=k8s.V1PodSpec( containers=[ k8s.V1Container( name="base", image="quay.io/keithwhitley4/airflow-ansible:2.0.0", image_pull_policy="Always", volume_mounts=[ kubeconfig.get_empty_dir_volume_mount() ]) ], volumes=[kubeconfig.get_empty_dir_volume_mount()])) } # General DAG Configuration self.dag = dag self.platform = platform # e.g. aws self.version = version # e.g. 4.6/4.7, major.minor only self.release_stream = release_stream # true release stream to follow. Nightlies, CI, etc. self.profile = profile # e.g. default/ovn # Specific Task Configuration self.vars = var_loader.build_task_vars(task="install", version=version, platform=platform, profile=profile) # Airflow Variables self.ansible_orchestrator = Variable.get("ansible_orchestrator", deserialize_json=True) self.release_stream_base_url = Variable.get("release_stream_base_url") self.install_secrets = Variable.get(f"openshift_install_config", deserialize_json=True) self.aws_creds = Variable.get("aws_creds", deserialize_json=True) self.gcp_creds = Variable.get("gcp_creds", deserialize_json=True) self.azure_creds = Variable.get("azure_creds", deserialize_json=True)
def __init__(self, dag, version, release_stream, platform, profile): self.exec_config = { "pod_override": k8s.V1Pod(spec=k8s.V1PodSpec(containers=[ k8s.V1Container( name="base", image="quay.io/keithwhitley4/airflow-ansible:2.0.0", image_pull_policy="Always", volume_mounts=[kubeconfig.get_kubeconfig_volume_mount()]) ], volumes=[ kubeconfig.get_kubeconfig_volume( version, platform, profile) ])) } # General DAG Configuration self.dag = dag self.platform = platform # e.g. aws self.version = version # e.g. 4.6/4.7, major.minor only self.release_stream = release_stream # true release stream to follow. Nightlies, CI, etc. self.profile = profile # e.g. default/ovn # Specific Task Configuration self.vars = var_loader.build_task_vars(task="index", version=version, platform=platform, profile=profile) self.release_stream_base_url = Variable.get("release_stream_base_url") latest_release = var_loader.get_latest_release_from_stream( self.release_stream_base_url, self.release_stream) self.env = { "OPENSHIFT_CLIENT_LOCATION": latest_release["openshift_client_location"], "RELEASE_STREAM": self.release_stream }
def __init__(self, dag, version, release_stream, platform, profile, default_args): self.exec_config = { "pod_override": k8s.V1Pod(spec=k8s.V1PodSpec(containers=[ k8s.V1Container( name="base", image="quay.io/keithwhitley4/airflow-ansible:2.0.0", image_pull_policy="Always", volume_mounts=[kubeconfig.get_kubeconfig_volume_mount()]) ], volumes=[ kubeconfig.get_kubeconfig_volume( version, platform, profile) ])) } # General DAG Configuration self.dag = dag self.platform = platform # e.g. aws self.version = version # e.g. stable/.next/.future self.release_stream = release_stream self.profile = profile # e.g. default/ovn self.default_args = default_args # Specific Task Configuration self.vars = var_loader.build_task_vars(task="benchmarks", version=version, platform=platform, profile=profile) self.release_stream_base_url = Variable.get("release_stream_base_url") latest_release = var_loader.get_latest_release_from_stream( self.release_stream_base_url, self.release_stream) self.env = { "OPENSHIFT_CLIENT_LOCATION": latest_release["openshift_client_location"] }
def test_pod_template_file_with_full_pod_spec(self): fixture = sys.path[0] + '/tests/kubernetes/basic_pod.yaml' pod_spec = k8s.V1Pod( metadata=k8s.V1ObjectMeta(labels={ "foo": "bar", "fizz": "buzz" }, ), spec=k8s.V1PodSpec(containers=[ k8s.V1Container( name="base", env=[k8s.V1EnvVar(name="env_name", value="value")], ) ]), ) k = KubernetesPodOperator( task_id="task" + self.get_current_task_name(), in_cluster=False, pod_template_file=fixture, full_pod_spec=pod_spec, do_xcom_push=True, ) context = create_context(k) result = k.execute(context) assert result is not None assert k.pod.metadata.labels == { 'fizz': 'buzz', 'foo': 'bar', 'airflow_version': mock.ANY, 'dag_id': 'dag', 'execution_date': mock.ANY, 'kubernetes_pod_operator': 'True', 'task_id': mock.ANY, 'try_number': '1', } assert k.pod.spec.containers[0].env == [ k8s.V1EnvVar(name="env_name", value="value") ] assert result == {"hello": "world"}
def test_full_pod_spec(self): pod_spec = k8s.V1Pod( metadata=k8s.V1ObjectMeta(labels={ "foo": "bar", "fizz": "buzz" }, namespace="default", name="test-pod"), spec=k8s.V1PodSpec( containers=[ k8s.V1Container( name="base", image="perl", command=["/bin/bash"], args=[ "-c", 'echo {\\"hello\\" : \\"world\\"} | cat > /airflow/xcom/return.json' ], env=[k8s.V1EnvVar(name="env_name", value="value")], ) ], restart_policy="Never", ), ) k = KubernetesPodOperator( task_id="task" + self.get_current_task_name(), in_cluster=False, full_pod_spec=pod_spec, do_xcom_push=True, ) context = create_context(k) result = k.execute(context) assert result is not None assert k.pod.metadata.labels == {'fizz': 'buzz', 'foo': 'bar'} assert k.pod.spec.containers[0].env == [ k8s.V1EnvVar(name="env_name", value="value") ] assert result == {"hello": "world"}
def test_not_adopt_unassigned_task(self, mock_kube_client): """ We should not adopt any tasks that were not assigned by the scheduler. This ensures that there is no contention over pod management. """ executor = self.kubernetes_executor executor.scheduler_job_id = "modified" pod_ids = {"foobar": {}} pod = k8s.V1Pod(metadata=k8s.V1ObjectMeta( name="foo", labels={"airflow-worker": "bar"}, annotations={ 'dag_id': 'dag', 'execution_date': datetime.utcnow().isoformat(), 'task_id': 'task', 'try_number': '1', }, )) executor.adopt_launched_task(mock_kube_client, pod=pod, pod_ids=pod_ids) assert not mock_kube_client.patch_namespaced_pod.called assert pod_ids == {"foobar": {}}
def setUp(self): self.static_uuid = uuid.UUID('cf4a56d2-8101-4217-b027-2af6216feb48') self.deserialize_result = { 'apiVersion': 'v1', 'kind': 'Pod', 'metadata': {'name': 'memory-demo', 'namespace': 'mem-example'}, 'spec': { 'containers': [ { 'args': ['--vm', '1', '--vm-bytes', '150M', '--vm-hang', '1'], 'command': ['stress'], 'image': 'apache/airflow:stress-2020.07.10-1.0.4', 'name': 'memory-demo-ctr', 'resources': {'limits': {'memory': '200Mi'}, 'requests': {'memory': '100Mi'}}, } ] }, } self.envs = {'ENVIRONMENT': 'prod', 'LOG_LEVEL': 'warning'} self.secrets = [ # This should be a secretRef Secret('env', None, 'secret_a'), # This should be a single secret mounted in volumeMounts Secret('volume', '/etc/foo', 'secret_b'), # This should produce a single secret mounted in env Secret('env', 'TARGET', 'secret_b', 'source_b'), ] self.execution_date = parser.parse('2020-08-24 00:00:00.000000') self.execution_date_label = datetime_to_label_safe_datestring(self.execution_date) self.dag_id = 'dag_id' self.task_id = 'task_id' self.try_number = 3 self.labels = { 'airflow-worker': 'uuid', 'dag_id': self.dag_id, 'execution_date': self.execution_date_label, 'task_id': self.task_id, 'try_number': str(self.try_number), 'airflow_version': __version__.replace('+', '-'), 'kubernetes_executor': 'True', } self.annotations = { 'dag_id': self.dag_id, 'task_id': self.task_id, 'execution_date': self.execution_date.isoformat(), 'try_number': str(self.try_number), } self.metadata = { 'labels': self.labels, 'name': 'pod_id-' + self.static_uuid.hex, 'namespace': 'namespace', 'annotations': self.annotations, } self.resources = k8s.V1ResourceRequirements( requests={ "cpu": 1, "memory": "1Gi", "ephemeral-storage": "2Gi", }, limits={"cpu": 2, "memory": "2Gi", "ephemeral-storage": "4Gi", 'nvidia.com/gpu': 1}, ) self.k8s_client = ApiClient() self.expected = k8s.V1Pod( api_version="v1", kind="Pod", metadata=k8s.V1ObjectMeta( namespace="default", name='myapp-pod-' + self.static_uuid.hex, labels={'app': 'myapp'}, ), spec=k8s.V1PodSpec( containers=[ k8s.V1Container( name='base', image='busybox', command=['sh', '-c', 'echo Hello Kubernetes!'], env=[ k8s.V1EnvVar(name='ENVIRONMENT', value='prod'), k8s.V1EnvVar( name="LOG_LEVEL", value='warning', ), k8s.V1EnvVar( name='TARGET', value_from=k8s.V1EnvVarSource( secret_key_ref=k8s.V1SecretKeySelector(name='secret_b', key='source_b') ), ), ], env_from=[ k8s.V1EnvFromSource(config_map_ref=k8s.V1ConfigMapEnvSource(name='configmap_a')), k8s.V1EnvFromSource(config_map_ref=k8s.V1ConfigMapEnvSource(name='configmap_b')), k8s.V1EnvFromSource(secret_ref=k8s.V1SecretEnvSource(name='secret_a')), ], ports=[k8s.V1ContainerPort(name="foo", container_port=1234)], resources=k8s.V1ResourceRequirements( requests={'memory': '100Mi'}, limits={ 'memory': '200Mi', }, ), ) ], security_context=k8s.V1PodSecurityContext( fs_group=2000, run_as_user=1000, ), host_network=True, image_pull_secrets=[ k8s.V1LocalObjectReference(name="pull_secret_a"), k8s.V1LocalObjectReference(name="pull_secret_b"), ], ), )
def test_from_obj(self): result = PodGenerator.from_obj( { "pod_override": k8s.V1Pod( api_version="v1", kind="Pod", metadata=k8s.V1ObjectMeta(name="foo", annotations={"test": "annotation"}), spec=k8s.V1PodSpec( containers=[ k8s.V1Container( name="base", volume_mounts=[ k8s.V1VolumeMount( mount_path="/foo/", name="example-kubernetes-test-volume" ) ], ) ], volumes=[ k8s.V1Volume( name="example-kubernetes-test-volume", host_path=k8s.V1HostPathVolumeSource(path="/tmp/"), ) ], ), ) } ) result = self.k8s_client.sanitize_for_serialization(result) assert { 'apiVersion': 'v1', 'kind': 'Pod', 'metadata': { 'name': 'foo', 'annotations': {'test': 'annotation'}, }, 'spec': { 'containers': [ { 'name': 'base', 'volumeMounts': [{'mountPath': '/foo/', 'name': 'example-kubernetes-test-volume'}], } ], 'volumes': [{'hostPath': {'path': '/tmp/'}, 'name': 'example-kubernetes-test-volume'}], }, } == result result = PodGenerator.from_obj( { "KubernetesExecutor": { "annotations": {"test": "annotation"}, "volumes": [ { "name": "example-kubernetes-test-volume", "hostPath": {"path": "/tmp/"}, }, ], "volume_mounts": [ { "mountPath": "/foo/", "name": "example-kubernetes-test-volume", }, ], } } ) result_from_pod = PodGenerator.from_obj( { "pod_override": k8s.V1Pod( metadata=k8s.V1ObjectMeta(annotations={"test": "annotation"}), spec=k8s.V1PodSpec( containers=[ k8s.V1Container( name="base", volume_mounts=[ k8s.V1VolumeMount( name="example-kubernetes-test-volume", mount_path="/foo/" ) ], ) ], volumes=[k8s.V1Volume(name="example-kubernetes-test-volume", host_path="/tmp/")], ), ) } ) result = self.k8s_client.sanitize_for_serialization(result) result_from_pod = self.k8s_client.sanitize_for_serialization(result_from_pod) expected_from_pod = { 'metadata': {'annotations': {'test': 'annotation'}}, 'spec': { 'containers': [ { 'name': 'base', 'volumeMounts': [{'mountPath': '/foo/', 'name': 'example-kubernetes-test-volume'}], } ], 'volumes': [{'hostPath': '/tmp/', 'name': 'example-kubernetes-test-volume'}], }, } assert ( result_from_pod == expected_from_pod ), "There was a discrepency between KubernetesExecutor and pod_override" assert { 'apiVersion': 'v1', 'kind': 'Pod', 'metadata': { 'annotations': {'test': 'annotation'}, }, 'spec': { 'containers': [ { 'args': [], 'command': [], 'env': [], 'envFrom': [], 'name': 'base', 'ports': [], 'volumeMounts': [{'mountPath': '/foo/', 'name': 'example-kubernetes-test-volume'}], } ], 'hostNetwork': False, 'imagePullSecrets': [], 'volumes': [{'hostPath': {'path': '/tmp/'}, 'name': 'example-kubernetes-test-volume'}], }, } == result
log = logging.getLogger(__name__) try: from kubernetes.client import models as k8s with DAG( dag_id='example_kubernetes_executor_config', schedule_interval=None, start_date=datetime(2021, 1, 1), catchup=False, tags=['example3'], ) as dag: # You can use annotations on your kubernetes pods! start_task_executor_config = { "pod_override": k8s.V1Pod(metadata=k8s.V1ObjectMeta( annotations={"test": "annotation"})) } @task(executor_config=start_task_executor_config) def start_task(): print_stuff() start_task = start_task() # [START task_with_volume] executor_config_volume_mount = { "pod_override": k8s.V1Pod(spec=k8s.V1PodSpec( containers=[ k8s.V1Container( name="base",