def test_image_pull_secrets_correctly_set(self, mock_client, monitor_mock,
start_mock):
from airflow.utils.state import State
fake_pull_secrets = "fakeSecret"
k = KubernetesPodOperator(
namespace='default',
image="ubuntu:16.04",
cmds=["bash", "-cx"],
arguments=["echo 10"],
labels={"foo": "bar"},
name="test",
task_id="task",
in_cluster=False,
do_xcom_push=False,
image_pull_secrets=[k8s.V1LocalObjectReference(fake_pull_secrets)],
cluster_context='default',
)
monitor_mock.return_value = (State.SUCCESS, None)
context = self.create_context(k)
k.execute(context=context)
self.assertEqual(
start_mock.call_args[0][0].spec.image_pull_secrets,
[k8s.V1LocalObjectReference(name=fake_pull_secrets)],
)
def test_image_pull_secrets_correctly_set(self, mock_client,
await_pod_completion_mock,
create_mock):
fake_pull_secrets = "fakeSecret"
k = KubernetesPodOperator(
namespace='default',
image="ubuntu:16.04",
cmds=["bash", "-cx"],
arguments=["echo 10"],
labels={"foo": "bar"},
name="test",
task_id="task",
in_cluster=False,
do_xcom_push=False,
image_pull_secrets=fake_pull_secrets,
cluster_context='default',
)
mock_pod = MagicMock()
mock_pod.status.phase = 'Succeeded'
await_pod_completion_mock.return_value = mock_pod
context = create_context(k)
k.execute(context=context)
assert create_mock.call_args[1]['pod'].spec.image_pull_secrets == [
k8s.V1LocalObjectReference(name=fake_pull_secrets)
]
def _get_image_pull_secrets(self):
"""Extracts any image pull secrets for fetching container(s)"""
if not self.kube_config.image_pull_secrets:
return []
pull_secrets = self.kube_config.image_pull_secrets.split(',')
return list(
map(lambda name: k8s.V1LocalObjectReference(name), pull_secrets))
def test_image_pull_secrets_correctly_set(self):
fake_pull_secrets = "fakeSecret"
k = KubernetesPodOperator(
namespace="default",
image="ubuntu:16.04",
cmds=["bash", "-cx"],
arguments=["echo 10"],
labels={"foo": "bar"},
name="test",
task_id="task",
in_cluster=False,
do_xcom_push=False,
image_pull_secrets=[k8s.V1LocalObjectReference(fake_pull_secrets)],
cluster_context="default",
)
pod = k.create_pod_request_obj()
assert pod.spec.image_pull_secrets == [k8s.V1LocalObjectReference(name=fake_pull_secrets)]
def convert_image_pull_secrets(image_pull_secrets) -> List[k8s.V1LocalObjectReference]:
"""
Converts a PodRuntimeInfoEnv into an k8s.V1EnvVar
:param image_pull_secrets:
:return:
"""
if isinstance(image_pull_secrets, str):
secrets = image_pull_secrets.split(",")
return [k8s.V1LocalObjectReference(name=secret) for secret in secrets]
else:
return image_pull_secrets
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 __init__(
self,
image: Optional[str] = None,
name: Optional[str] = None,
namespace: Optional[str] = None,
volume_mounts: Optional[List[Union[k8s.V1VolumeMount, dict]]] = None,
envs: Optional[Dict[str, str]] = None,
cmds: Optional[List[str]] = None,
args: Optional[List[str]] = None,
labels: Optional[Dict[str, str]] = None,
node_selectors: Optional[Dict[str, str]] = None,
ports: Optional[List[Union[k8s.V1ContainerPort, dict]]] = None,
volumes: Optional[List[Union[k8s.V1Volume, dict]]] = None,
image_pull_policy: Optional[str] = None,
restart_policy: Optional[str] = None,
image_pull_secrets: Optional[str] = None,
init_containers: Optional[List[k8s.V1Container]] = None,
service_account_name: Optional[str] = None,
resources: Optional[Union[k8s.V1ResourceRequirements, dict]] = None,
annotations: Optional[Dict[str, str]] = None,
affinity: Optional[dict] = None,
hostnetwork: bool = False,
tolerations: Optional[list] = None,
security_context: Optional[Union[k8s.V1PodSecurityContext,
dict]] = None,
configmaps: Optional[List[str]] = None,
dnspolicy: Optional[str] = None,
schedulername: Optional[str] = None,
extract_xcom: bool = False,
priority_class_name: Optional[str] = None,
):
self.pod = k8s.V1Pod()
self.pod.api_version = 'v1'
self.pod.kind = 'Pod'
# Pod Metadata
self.metadata = k8s.V1ObjectMeta()
self.metadata.labels = labels
self.metadata.name = name
self.metadata.namespace = namespace
self.metadata.annotations = annotations
# Pod Container
self.container = k8s.V1Container(name='base')
self.container.image = image
self.container.env = []
if envs:
if isinstance(envs, dict):
for key, val in envs.items():
self.container.env.append(k8s.V1EnvVar(name=key,
value=val))
elif isinstance(envs, list):
self.container.env.extend(envs)
configmaps = configmaps or []
self.container.env_from = []
for configmap in configmaps:
self.container.env_from.append(
k8s.V1EnvFromSource(config_map_ref=k8s.V1ConfigMapEnvSource(
name=configmap)))
self.container.command = cmds or []
self.container.args = args or []
if image_pull_policy:
self.container.image_pull_policy = image_pull_policy
self.container.ports = ports or []
self.container.resources = resources
self.container.volume_mounts = volume_mounts or []
# Pod Spec
self.spec = k8s.V1PodSpec(containers=[])
self.spec.security_context = security_context
self.spec.tolerations = tolerations
if dnspolicy:
self.spec.dns_policy = dnspolicy
self.spec.scheduler_name = schedulername
self.spec.host_network = hostnetwork
self.spec.affinity = affinity
self.spec.service_account_name = service_account_name
self.spec.init_containers = init_containers
self.spec.volumes = volumes or []
self.spec.node_selector = node_selectors
if restart_policy:
self.spec.restart_policy = restart_policy
self.spec.priority_class_name = priority_class_name
self.spec.image_pull_secrets = []
if image_pull_secrets:
for image_pull_secret in image_pull_secrets.split(','):
self.spec.image_pull_secrets.append(
k8s.V1LocalObjectReference(name=image_pull_secret))
# Attach sidecar
self.extract_xcom = extract_xcom
def test_pod_template_file_override_in_executor_config(
self, mock_get_kube_client, mock_run_pod_async):
current_folder = pathlib.Path(__file__).parent.absolute()
template_file = str(
(current_folder / "kubernetes_executor_template_files" /
"basic_template.yaml").absolute())
mock_kube_client = mock.patch('kubernetes.client.CoreV1Api',
autospec=True)
mock_get_kube_client.return_value = mock_kube_client
with conf_vars({('kubernetes', 'pod_template_file'): ''}):
executor = self.kubernetes_executor
executor.start()
assert executor.event_buffer == {}
assert executor.task_queue.empty()
execution_date = datetime.utcnow()
executor.execute_async(
key=('dag', 'task', execution_date, 1),
queue=None,
command=['airflow', 'tasks', 'run', 'true', 'some_parameter'],
executor_config={
"pod_template_file":
template_file,
"pod_override":
k8s.V1Pod(
metadata=k8s.V1ObjectMeta(
labels={"release": "stable"}),
spec=k8s.V1PodSpec(containers=[
k8s.V1Container(name="base", image="airflow:3.6")
], ),
),
},
)
assert not executor.task_queue.empty()
task = executor.task_queue.get_nowait()
_, _, expected_executor_config, expected_pod_template_file = task
# Test that the correct values have been put to queue
assert expected_executor_config.metadata.labels == {
'release': 'stable'
}
assert expected_pod_template_file == template_file
self.kubernetes_executor.kube_scheduler.run_next(task)
mock_run_pod_async.assert_called_once_with(
k8s.V1Pod(
api_version="v1",
kind="Pod",
metadata=k8s.V1ObjectMeta(
name=mock.ANY,
namespace="default",
annotations={
'dag_id': 'dag',
'execution_date': execution_date.isoformat(),
'task_id': 'task',
'try_number': '1',
},
labels={
'airflow-worker':
'5',
'airflow_version':
mock.ANY,
'dag_id':
'dag',
'execution_date':
datetime_to_label_safe_datestring(execution_date),
'kubernetes_executor':
'True',
'mylabel':
'foo',
'release':
'stable',
'task_id':
'task',
'try_number':
'1',
},
),
spec=k8s.V1PodSpec(
containers=[
k8s.V1Container(
name="base",
image="airflow:3.6",
args=[
'airflow', 'tasks', 'run', 'true',
'some_parameter'
],
env=[
k8s.V1EnvVar(
name='AIRFLOW_IS_K8S_EXECUTOR_POD',
value='True')
],
)
],
image_pull_secrets=[
k8s.V1LocalObjectReference(name='airflow-registry')
],
scheduler_name='default-scheduler',
security_context=k8s.V1PodSecurityContext(
fs_group=50000, run_as_user=50000),
),
))
def __init__( # pylint: disable=too-many-arguments,too-many-locals
self,
image,
name=None,
namespace=None,
volume_mounts=None,
envs=None,
cmds=None,
args=None,
labels=None,
node_selectors=None,
ports=None,
volumes=None,
image_pull_policy='IfNotPresent',
restart_policy='Never',
image_pull_secrets=None,
init_containers=None,
service_account_name=None,
resources=None,
annotations=None,
affinity=None,
hostnetwork=False,
tolerations=None,
security_context=None,
configmaps=None,
dnspolicy=None,
pod=None,
extract_xcom=False,
):
self.ud_pod = pod
self.pod = k8s.V1Pod()
self.pod.api_version = 'v1'
self.pod.kind = 'Pod'
# Pod Metadata
self.metadata = k8s.V1ObjectMeta()
self.metadata.labels = labels
self.metadata.name = name + "-" + str(uuid.uuid4())[:8] if name else None
self.metadata.namespace = namespace
self.metadata.annotations = annotations
# Pod Container
self.container = k8s.V1Container(name='base')
self.container.image = image
self.container.env = []
if envs:
if isinstance(envs, dict):
for key, val in envs.items():
self.container.env.append(k8s.V1EnvVar(
name=key,
value=val
))
elif isinstance(envs, list):
self.container.env.extend(envs)
configmaps = configmaps or []
self.container.env_from = []
for configmap in configmaps:
self.container.env_from.append(k8s.V1EnvFromSource(
config_map_ref=k8s.V1ConfigMapEnvSource(
name=configmap
)
))
self.container.command = cmds or []
self.container.args = args or []
self.container.image_pull_policy = image_pull_policy
self.container.ports = ports or []
self.container.resources = resources
self.container.volume_mounts = volume_mounts or []
# Pod Spec
self.spec = k8s.V1PodSpec(containers=[])
self.spec.security_context = security_context
self.spec.tolerations = tolerations
self.spec.dns_policy = dnspolicy
self.spec.host_network = hostnetwork
self.spec.affinity = affinity
self.spec.service_account_name = service_account_name
self.spec.init_containers = init_containers
self.spec.volumes = volumes or []
self.spec.node_selector = node_selectors
self.spec.restart_policy = restart_policy
self.spec.image_pull_secrets = []
if image_pull_secrets:
for image_pull_secret in image_pull_secrets.split(','):
self.spec.image_pull_secrets.append(k8s.V1LocalObjectReference(
name=image_pull_secret
))
# Attach sidecar
self.extract_xcom = extract_xcom
env_from=configmaps,
name="airflow-test-pod",
task_id="task",
affinity=affinity,
is_delete_operator_pod=True,
hostnetwork=False,
tolerations=tolerations,
init_containers=[init_container],
priority_class_name="medium",
)
# [START howto_operator_k8s_private_image]
quay_k8s = KubernetesPodOperator(
namespace="default",
image="quay.io/apache/bash",
image_pull_secrets=[k8s.V1LocalObjectReference("testquay")],
cmds=["bash", "-cx"],
arguments=["echo", "10", "echo pwd"],
labels={"foo": "bar"},
name="airflow-private-image-pod",
is_delete_operator_pod=True,
in_cluster=True,
task_id="task-two",
get_logs=True,
)
# [END howto_operator_k8s_private_image]
# [START howto_operator_k8s_write_xcom]
write_xcom = KubernetesPodOperator(
namespace="default",
image="alpine",
default_args = {
'owner': 'airflow',
'depends_on_past': False,
'start_date': datetime.utcnow()
}
with DAG('dataload-01',
default_args=default_args,
schedule_interval=None,
max_active_tasks=10,
max_active_runs=10) as dag:
load_resources = V1ResourceRequirements(requests={"memory": "18Gi"},
limits={"memory": "18Gi"})
node_selector = {"loader-node": "true"}
image_pull_secrets = [k8s.V1LocalObjectReference('falkonry-pull-secret')]
envs = list()
load_file = "{{ dag_run.conf['load_file'] }}"
compact_file = "{{ dag_run.conf['compact_file'] }}"
task_id = "{{ dag_run.conf['task_id'] }}"
envs.append(
k8s.V1EnvVar(name="AWS_SECRET_ACCESS_KEY",
value_from=k8s.V1EnvVarSource(
secret_key_ref=k8s.V1SecretKeySelector(
key="secret-id", name="aws-key-prod"))))
envs.append(
k8s.V1EnvVar(name="AWS_DEFAULT_REGION",
value_from=k8s.V1EnvVarSource(
secret_key_ref=k8s.V1SecretKeySelector(
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,
annotations=self.annotations,
)
if self.image_pull_secrets:
image_pull_secrets = [
k8s.V1LocalObjectReference(i)
for i in self.image_pull_secrets.split(",")
]
else:
image_pull_secrets = []
spec = k8s.V1PodSpec(
init_containers=self.init_containers,
containers=[
k8s.V1Container(
image=self.image,
command=self.cmds,
env_from=[],
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=image_pull_secrets,
service_account_name=self.service_account_name,
node_selector=self.node_selectors,
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 configmap_name in self.configmaps:
env_var = k8s.V1EnvFromSource(
config_map_ref=k8s.V1ConfigMapEnvSource(name=configmap_name, ))
pod.spec.containers[0].env_from.append(env_var)
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 __init__(
self,
image=None,
name=None,
namespace=None,
volume_mounts=None,
envs=None,
cmds=None,
args=None,
labels=None,
node_selectors=None,
ports=None,
volumes=None,
image_pull_policy=None,
restart_policy=None,
image_pull_secrets=None,
init_containers=None,
service_account_name=None,
resources=None,
annotations=None,
affinity=None,
hostnetwork=False,
tolerations=None,
security_context=None,
configmaps=None,
dnspolicy=None,
schedulername=None,
priority_class_name=None,
pod=None,
pod_template_file=None,
extract_xcom=False,
):
if pod_template_file:
self.ud_pod = self.deserialize_model_file(pod_template_file)
else:
self.ud_pod = pod
self.pod = k8s.V1Pod()
self.pod.api_version = 'v1'
self.pod.kind = 'Pod'
# Pod Metadata
self.metadata = k8s.V1ObjectMeta()
self.metadata.labels = labels
self.metadata.name = name
self.metadata.namespace = namespace
self.metadata.annotations = annotations
# Pod Container
self.container = k8s.V1Container(name='base')
self.container.image = image
self.container.env = []
if envs:
if isinstance(envs, dict):
for key, val in envs.items():
self.container.env.append(k8s.V1EnvVar(name=key,
value=val))
elif isinstance(envs, list):
self.container.env.extend(envs)
configmaps = configmaps or []
self.container.env_from = []
for configmap in configmaps:
self.container.env_from.append(
k8s.V1EnvFromSource(config_map_ref=k8s.V1ConfigMapEnvSource(
name=configmap)))
self.container.command = cmds or []
self.container.args = args or []
self.container.image_pull_policy = image_pull_policy
self.container.ports = ports or []
self.container.resources = resources
self.container.volume_mounts = [
v.to_k8s_client_obj()
for v in _extract_volume_mounts(volume_mounts)
]
# Pod Spec
self.spec = k8s.V1PodSpec(containers=[])
self.spec.security_context = security_context
self.spec.tolerations = tolerations
self.spec.dns_policy = dnspolicy
self.spec.scheduler_name = schedulername
self.spec.host_network = hostnetwork
self.spec.affinity = affinity
self.spec.service_account_name = service_account_name
self.spec.init_containers = init_containers
self.spec.volumes = volumes or []
self.spec.node_selector = node_selectors
self.spec.restart_policy = restart_policy
self.spec.priority_class_name = priority_class_name
self.spec.image_pull_secrets = []
if image_pull_secrets:
for image_pull_secret in image_pull_secrets.split(','):
self.spec.image_pull_secrets.append(
k8s.V1LocalObjectReference(name=image_pull_secret))
# Attach sidecar
self.extract_xcom = extract_xcom
def _to_real_pod(
self,
cmds: List[str],
args: List[str],
namespace: str,
name: str,
image: str,
envs: Dict[str, str],
labels: Dict[str, str],
annotations: Dict[str, str],
resources: "DbndExtendedResources",
secrets: List["Secret"],
) -> k8s.V1Pod:
# TODO add yaml template as basis
BASE_CONTAINER_NAME = "base"
kc: KubernetesEngineConfig = self
meta = k8s.V1ObjectMeta(labels=labels,
name=name,
namespace=namespace,
annotations=annotations)
if kc.image_pull_secrets:
image_pull_secrets = [
k8s.V1LocalObjectReference(i)
for i in kc.image_pull_secrets.split(",")
]
else:
image_pull_secrets = []
spec = k8s.V1PodSpec(
# init_containers=kc.init_containers,
containers=[
k8s.V1Container(
image=image,
command=cmds,
env_from=[],
name=BASE_CONTAINER_NAME,
env=[
k8s.V1EnvVar(name=key, value=val)
for key, val in envs.items()
],
args=args,
image_pull_policy=kc.image_pull_policy,
)
],
image_pull_secrets=image_pull_secrets,
service_account_name=kc.service_account_name,
node_selector=kc.node_selectors,
# dns_policy=kc.dnspolicy,
host_network=kc.hostnetwork,
tolerations=kc.tolerations,
affinity=kc.affinity,
security_context=kc.security_context,
)
k8_pod = k8s.V1Pod(spec=spec, metadata=meta)
for configmap_name in kc.configmaps:
env_var = k8s.V1EnvFromSource(
config_map_ref=k8s.V1ConfigMapEnvSource(name=configmap_name))
k8_pod.spec.containers[0].env_from.append(env_var)
volumes = kc.volumes or []
for volume in volumes:
k8_pod = volume_shims.attach_to_pod(k8_pod, volume)
mounts = kc.volume_mounts or []
for volume_mount in mounts:
k8_pod = attach_volume_mount(k8_pod, volume_mount)
secret: Secret
for secret in secrets:
if AIRFLOW_ABOVE_10:
k8_pod = secret.attach_to_pod(k8_pod)
else:
k8_pod = attach_to_pod(secret, k8_pod)
k8_pod = resources.attach_to_pod(k8_pod)
return k8_pod