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",
