def from_obj(obj) -> k8s.V1Pod:
"""Converts to pod from obj"""
if obj is None:
return k8s.V1Pod()
if isinstance(obj, PodGenerator):
return obj.gen_pod()
if not isinstance(obj, dict):
raise TypeError(
'Cannot convert a non-dictionary or non-PodGenerator '
'object into a KubernetesExecutorConfig')
namespaced = obj.get(Executors.KubernetesExecutor, {})
resources = namespaced.get('resources')
if resources is None:
requests = {
'cpu': namespaced.get('request_cpu'),
'memory': namespaced.get('request_memory')
}
limits = {
'cpu': namespaced.get('limit_cpu'),
'memory': namespaced.get('limit_memory')
}
all_resources = list(requests.values()) + list(limits.values())
if all(r is None for r in all_resources):
resources = None
else:
resources = k8s.V1ResourceRequirements(requests=requests,
limits=limits)
annotations = namespaced.get('annotations', {})
gcp_service_account_key = namespaced.get('gcp_service_account_key',
None)
if annotations is not None and gcp_service_account_key is not None:
annotations.update({
'iam.cloud.google.com/service-account':
gcp_service_account_key
})
pod_spec_generator = PodGenerator(
image=namespaced.get('image'),
envs=namespaced.get('env'),
cmds=namespaced.get('cmds'),
args=namespaced.get('args'),
labels=namespaced.get('labels'),
node_selectors=namespaced.get('node_selectors'),
name=namespaced.get('name'),
ports=namespaced.get('ports'),
volumes=namespaced.get('volumes'),
volume_mounts=namespaced.get('volume_mounts'),
namespace=namespaced.get('namespace'),
image_pull_policy=namespaced.get('image_pull_policy'),
restart_policy=namespaced.get('restart_policy'),
image_pull_secrets=namespaced.get('image_pull_secrets'),
init_containers=namespaced.get('init_containers'),
service_account_name=namespaced.get('service_account_name'),
resources=resources,
annotations=namespaced.get('annotations'),
affinity=namespaced.get('affinity'),
hostnetwork=namespaced.get('hostnetwork'),
tolerations=namespaced.get('tolerations'),
security_context=namespaced.get('security_context'),
configmaps=namespaced.get('configmaps'),
dnspolicy=namespaced.get('dnspolicy'),
pod=namespaced.get('pod'),
extract_xcom=namespaced.get('extract_xcom'),
)
return pod_spec_generator.gen_pod()
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': '2.0.0.dev0',
'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")
]),
)
match_expressions=[
k8s.V1LabelSelectorRequirement(key='app', operator='In', values=['airflow'])
]
),
topology_key='kubernetes.io/hostname',
)
]
)
)
# Use k8s_client.V1Toleration to define node tolerations
k8s_tolerations = [k8s.V1Toleration(key='dedicated', operator='Equal', value='airflow')]
# Use k8s_client.V1ResourceRequirements to define resource limits
k8s_resource_requirements = k8s.V1ResourceRequirements(
requests={'memory': '512Mi'}, limits={'memory': '512Mi'}
)
kube_exec_config_resource_limits = {
"pod_override": k8s.V1Pod(
spec=k8s.V1PodSpec(
containers=[
k8s.V1Container(
name="base",
resources=k8s_resource_requirements,
)
],
affinity=k8s_affinity,
tolerations=k8s_tolerations,
)
)
def test_construct_pod(self, mock_uuid):
mock_uuid.return_value = self.static_uuid
worker_config = k8s.V1Pod(
metadata=k8s.V1ObjectMeta(name='gets-overridden-by-dynamic-args',
annotations={'should': 'stay'}),
spec=k8s.V1PodSpec(containers=[
k8s.V1Container(name='doesnt-override',
resources=k8s.V1ResourceRequirements(
limits={
'cpu': '1m',
'memory': '1G'
}),
security_context=k8s.V1SecurityContext(
run_as_user=1))
]))
executor_config = k8s.V1Pod(spec=k8s.V1PodSpec(containers=[
k8s.V1Container(name='doesnt-override-either',
resources=k8s.V1ResourceRequirements(limits={
'cpu': '2m',
'memory': '2G'
}))
]))
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)
expected_metadata = dict(self.metadata)
expected_metadata['annotations'].update({'should': 'stay'})
self.assertEqual(
{
'apiVersion': 'v1',
'kind': 'Pod',
'metadata': expected_metadata,
'spec': {
'containers': [{
'args': [],
'command': ['command'],
'env': [],
'envFrom': [],
'name': 'base',
'image': 'kube_image',
'ports': [],
'resources': {
'limits': {
'cpu': '2m',
'memory': '2G'
}
},
'volumeMounts': [],
'securityContext': {
'runAsUser': 1
}
}],
'hostNetwork':
False,
'imagePullSecrets': [],
'volumes': []
}
}, sanitized_result)