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_merge_objects(self):
base_annotations = {'foo1': 'bar1'}
base_labels = {'foo1': 'bar1'}
client_annotations = {'foo2': 'bar2'}
base_obj = k8s.V1ObjectMeta(annotations=base_annotations, labels=base_labels)
client_obj = k8s.V1ObjectMeta(annotations=client_annotations)
res = merge_objects(base_obj, client_obj)
client_obj.labels = base_labels
assert client_obj == res
def handleSecret(self, key: str, value: T.Dict[str, str],
full_path: str) -> None:
if not self.jobSecret:
assert self.object
name = self.dnsify("jobsecret." + str(self.object.get("id", "")))
metadata = K.V1ObjectMeta(name=name)
assert self.manifest_list
clusterMeta = self.manifest_list.clusterMeta()
if clusterMeta:
metadata.annotations = clusterMeta.annotations
self.jobSecret = K.V1Secret(metadata=metadata, data={})
self.jobSecret.api_version = 'v1'
self.jobSecret.kind = 'Secret'
self.manifest.append(self.jobSecret)
sourceSecretName = value.get("source", "")
if not sourceSecretName:
# logging.critical("{}.source missing", full_path)
return
source_key = utils.dnsify(sourceSecretName)
assert self.manifest_list
sourceSecret = self.manifest_list.manifest(pluginName='secret',
manifestName=source_key)
if not sourceSecret:
# logging.warning("secret {} missing", source_key)
self.jobSecret.data[key] = "{} not found".format(source_key)
return
[v] = sourceSecret[0].data.values()
self.jobSecret.data[key] = v
def build_pod_spec(name, bucket, data_dir):
metadata = k8s.V1ObjectMeta(name=make_unique_pod_name(name), )
container = k8s.V1Container(
name=name,
lifecycle=k8s.V1Lifecycle(
post_start=k8s.V1Handler(_exec=k8s.V1ExecAction(command=[
"gcsfuse",
"--log-file",
"/var/log/gcs_fuse.log",
"--temp-dir",
"/tmp",
"--debug_gcs",
bucket,
data_dir,
])),
pre_stop=k8s.V1Handler(_exec=k8s.V1ExecAction(
command=["fusermount", "-u", data_dir])),
),
security_context=k8s.V1SecurityContext(
privileged=True,
capabilities=k8s.V1Capabilities(add=["SYS_ADMIN"])),
)
pod = k8s.V1Pod(metadata=metadata,
spec=k8s.V1PodSpec(containers=[container]))
return pod
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)
self.assertIsNotNone(result)
self.assertEqual(k.pod.metadata.labels, {'fizz': 'buzz', 'foo': 'bar'})
self.assertEqual(k.pod.spec.containers[0].env, [k8s.V1EnvVar(name="env_name", value="value")])
self.assertDictEqual(result, {"hello": "world"})
def test_adopt_launched_task(self, mock_kube_client):
executor = self.kubernetes_executor
executor.scheduler_job_id = "modified"
annotations = {
'dag_id': 'dag',
'execution_date': datetime.utcnow().isoformat(),
'task_id': 'task',
'try_number': '1',
}
ti_key = annotations_to_key(annotations)
pod = k8s.V1Pod(
metadata=k8s.V1ObjectMeta(name="foo",
labels={"airflow-worker": "bar"},
annotations=annotations))
pod_ids = {ti_key: {}}
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'
},
'annotations': annotations,
'name': 'foo',
}
},
'name': 'foo',
'namespace': None,
}
assert pod_ids == {}
assert executor.running == {ti_key}
def initCronJob(self, key: str, value: str, full_path: str) -> None:
name = self.dnsify(value)
metadata = K.V1ObjectMeta(name=name)
assert self.manifest_list
clusterMeta = self.manifest_list.clusterMeta()
if clusterMeta:
metadata.annotations = clusterMeta.annotations
# intentionally written this way so one can easily scan down paths
container1 = K.V1Container(
name="job",
resources=K.V1ResourceRequirements(limits={}, requests={}),
)
self.manifest = Manifest(
data=[
K.V1beta1CronJob(
metadata=metadata,
kind="CronJob",
api_version="batch/v1beta1",
spec=K.V1beta1CronJobSpec(
schedule="* * * * *",
suspend=True,
job_template=K.V1beta1JobTemplateSpec(spec=K.V1JobSpec(
template=K.V1PodTemplateSpec(spec=K.V1PodSpec(
containers=[container1])))),
),
)
],
pluginName="metronome",
manifestName=name,
)
def setUp(self):
self.watcher = KubernetesJobWatcher(
namespace="airflow",
multi_namespace_mode=False,
watcher_queue=mock.MagicMock(),
resource_version="0",
scheduler_job_id="123",
kube_config=mock.MagicMock(),
)
self.kube_client = mock.MagicMock()
self.core_annotations = {
"dag_id": "dag",
"task_id": "task",
"execution_date": "dt",
"try_number": "1",
}
self.pod = k8s.V1Pod(
metadata=k8s.V1ObjectMeta(
name="foo",
annotations={
"airflow-worker": "bar",
**self.core_annotations
},
namespace="airflow",
resource_version="456",
),
status=k8s.V1PodStatus(phase="Pending"),
)
self.events = []
def test_encode_k8s_v1pod(self):
from kubernetes.client import models as k8s
pod = k8s.V1Pod(
metadata=k8s.V1ObjectMeta(
name="foo",
namespace="bar",
),
spec=k8s.V1PodSpec(
containers=[k8s.V1Container(
name="foo",
image="bar",
)]),
)
self.assertEqual(
json.loads(json.dumps(pod, cls=utils_json.AirflowJsonEncoder)),
{
"metadata": {
"name": "foo",
"namespace": "bar"
},
"spec": {
"containers": [{
"image": "bar",
"name": "foo"
}]
},
},
)
def translate_pool(self, key: str, value: str, full_path: str) -> None:
assert self.object is not None
objects = migrate(self.object)
cluster_annotations = {}
assert self.manifest_list is not None
cluster_metadata = self.manifest_list.clusterMeta()
if cluster_metadata is not None and cluster_metadata.annotations:
cluster_annotations = cluster_metadata.annotations
if not any(objects):
return
self.manifest = system.Manifest(
pluginName="ingress",
manifestName=self.dnsify(value),
)
assert self.manifest is not None
for obj in objects:
if not obj:
continue
if obj.metadata:
obj.metadata.annotations.update(cluster_annotations)
else:
obj.metadata = models.V1ObjectMeta(
annotations=cluster_annotations, )
self.manifest.append(obj)
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 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)
self.assertIsNotNone(result)
self.assertEqual(k.pod.metadata.labels, {'fizz': 'buzz', 'foo': 'bar'})
self.assertEqual(k.pod.spec.containers[0].env,
[k8s.V1EnvVar(name="env_name", value="value")])
self.assertDictEqual(result, {"hello": "world"})
def test_pending_pod_timeout_multi_namespace_mode(
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="foo90",
labels={"airflow-worker": "123"},
creation_timestamp=now - timedelta(seconds=500),
namespace="anothernamespace",
)),
]
mock_kube_client.list_pod_for_all_namespaces.return_value.items = pending_pods
config = {
('kubernetes', 'namespace'): 'mynamespace',
('kubernetes', 'multi_namespace_mode'): 'true',
('kubernetes', 'kube_client_request_args'): '{"sentinel": "foo"}',
}
with conf_vars(config):
executor = KubernetesExecutor()
executor.job_id = "123"
executor.start()
executor._check_worker_pods_pending_timeout()
mock_kube_client.list_pod_for_all_namespaces.assert_called_once_with(
field_selector='status.phase=Pending',
label_selector='airflow-worker=123',
limit=100,
sentinel='foo',
)
mock_delete_pod.assert_called_once_with('foo90', 'anothernamespace')
def test_full_pod_spec(self, mock_client, monitor_mock, start_mock):
from airflow.utils.state import State
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,
)
monitor_mock.return_value = (State.SUCCESS, None)
context = self.create_context(k)
k.execute(context=context)
assert start_mock.call_args[0][
0].metadata.name == pod_spec.metadata.name
assert start_mock.call_args[0][
0].metadata.labels == pod_spec.metadata.labels
assert start_mock.call_args[0][
0].metadata.namespace == pod_spec.metadata.namespace
assert start_mock.call_args[0][0].spec.containers[
0].image == pod_spec.spec.containers[0].image
assert start_mock.call_args[0][0].spec.containers[
0].command == pod_spec.spec.containers[0].command
# kwargs take precedence, however
start_mock.reset_mock()
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,
)
context = self.create_context(k)
k.execute(context=context)
# make sure the kwargs takes precedence (and that name is randomized)
assert start_mock.call_args[0][0].metadata.name.startswith(name_base)
assert start_mock.call_args[0][0].metadata.name != name_base
assert start_mock.call_args[0][0].spec.containers[0].image == image
def mock_iter_resp_lines(resp):
for name in ['test1', 'test2', 'test3']:
yield json.dumps({
'type':
'ADDED',
'object':
models.V1DeploymentConfig(metadata=k8s_models.V1ObjectMeta(
name=name, resource_version=1)).to_dict()
})
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_merge_objects_empty(self):
annotations = {'foo1': 'bar1'}
base_obj = k8s.V1ObjectMeta(annotations=annotations)
client_obj = None
res = merge_objects(base_obj, client_obj)
assert base_obj == res
client_obj = k8s.V1ObjectMeta()
res = merge_objects(base_obj, client_obj)
assert base_obj == res
client_obj = k8s.V1ObjectMeta(annotations=annotations)
base_obj = None
res = merge_objects(base_obj, client_obj)
assert client_obj == res
base_obj = k8s.V1ObjectMeta()
res = merge_objects(base_obj, client_obj)
assert client_obj == res
def test_merge_objects_empty(self):
annotations = {'foo1': 'bar1'}
base_obj = k8s.V1ObjectMeta(annotations=annotations)
client_obj = None
res = merge_objects(base_obj, client_obj)
self.assertEqual(base_obj, res)
client_obj = k8s.V1ObjectMeta()
res = merge_objects(base_obj, client_obj)
self.assertEqual(base_obj, res)
client_obj = k8s.V1ObjectMeta(annotations=annotations)
base_obj = None
res = merge_objects(base_obj, client_obj)
self.assertEqual(client_obj, res)
base_obj = k8s.V1ObjectMeta()
res = merge_objects(base_obj, client_obj)
self.assertEqual(client_obj, res)
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 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 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_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 ensure_kubernetes_namespace(self):
kubernetes_client = KubernetesClient()
try:
kubernetes_client.core.read_namespace(
self.job_resources["namespace"])
except kubernetes.client.exceptions.ApiException as e:
if e.status != 404:
raise e
kubernetes_client.core.create_namespace(
body=kubernetes_models.V1Namespace(
api_version="v1",
kind="Namespace",
metadata=kubernetes_models.V1ObjectMeta(
name=self.job_resources["namespace"]),
))
def create_pvc(self):
"""
Creates k8s pvc definition object
:return:
"""
pvc = k8s.V1PersistentVolumeClaim(
metadata=k8s.V1ObjectMeta(name=self.pvc_name,
namespace=self._namespace),
spec=k8s.V1PersistentVolumeClaimSpec(
access_modes=self._access_modes,
storage_class_name=self._storage_class_name,
resources=k8s.V1ResourceRequirements(
requests={"storage": self._volumes_size}),
),
)
return pvc
def migrate_lb(pool: Dict[str, Any]) -> List[models.V1Service]:
output = []
backends = pool.get("backends", {})
frontends = pool.get("frontends", {})
for frontend in frontends.values():
if "HTTP" in frontend.get("protocol").upper():
continue
backend = backends.get(frontend.get("default_backend"))
if not backend:
warnings.warn(
TCP_BACKEND_WARNING.format(frontend.get("name", "UNKNOWN")))
continue
# TODO(jkoelker) figure out targetPort
port = models.V1ServicePort(
port=frontend["port"],
target_port=0,
protocol=frontend.get("protocol").upper(),
)
spec = models.V1ServiceSpec(
type="LoadBalancer",
ports=[port],
selector={
"app": backend["service"]["name"],
},
)
metadata = models.V1ObjectMeta(
annotations={},
name=pool["name"],
namespace=pool.get("namespace"),
)
lb = models.V1Service(
api_version="v1",
kind="Service",
metadata=metadata,
spec=spec,
)
output.append(lb)
return output
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_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")
],
)
]),
)
