def __init__(self,
namespace,
image,
name,
cmds=None,
arguments=None,
env_vars=None,
secrets=None,
in_cluster=False,
labels=None,
startup_timeout_seconds=120,
get_logs=True,
image_pull_policy='IfNotPresent',
annotations=None,
resources=None,
*args,
**kwargs):
super(KubernetesPodOperator, self).__init__(*args, **kwargs)
self.image = image
self.namespace = namespace
self.cmds = cmds or []
self.arguments = arguments or []
self.labels = labels or {}
self.startup_timeout_seconds = startup_timeout_seconds
self.name = name
self.env_vars = env_vars or {}
self.secrets = secrets or []
self.in_cluster = in_cluster
self.get_logs = get_logs
self.image_pull_policy = image_pull_policy
self.annotations = annotations or {}
self.resources = resources or Resources()
def __init__(self,
namespace,
image,
name,
cmds=None,
arguments=None,
volume_mounts=None,
volumes=None,
env_vars=None,
secrets=None,
in_cluster=False,
cluster_context=None,
labels=None,
startup_timeout_seconds=120,
get_logs=True,
image_pull_policy='IfNotPresent',
annotations=None,
resources=None,
affinity=None,
config_file=None,
xcom_push=False,
node_selectors=None,
image_pull_secrets=None,
service_account_name="default",
is_delete_operator_pod=False,
hostnetwork=False,
tolerations=None,
configmaps=None,
security_context=None,
*args,
**kwargs):
super(KubernetesPodOperator, self).__init__(*args, **kwargs)
self.image = image
self.namespace = namespace
self.cmds = cmds or []
self.arguments = arguments or []
self.labels = labels or {}
self.startup_timeout_seconds = startup_timeout_seconds
self.name = name
self.env_vars = env_vars or {}
self.volume_mounts = volume_mounts or []
self.volumes = volumes or []
self.secrets = secrets or []
self.in_cluster = in_cluster
self.cluster_context = cluster_context
self.get_logs = get_logs
self.image_pull_policy = image_pull_policy
self.node_selectors = node_selectors or {}
self.annotations = annotations or {}
self.affinity = affinity or {}
self.xcom_push = xcom_push
self.resources = resources or Resources()
self.config_file = config_file
self.image_pull_secrets = image_pull_secrets
self.service_account_name = service_account_name
self.is_delete_operator_pod = is_delete_operator_pod
self.hostnetwork = hostnetwork
self.tolerations = tolerations or []
self.configmaps = configmaps or []
self.security_context = security_context or {}
def _set_resources(self, resources):
# Legacy
inputResource = Resources()
if resources:
for item in resources.keys():
setattr(inputResource, item, resources[item])
return inputResource
def convert_resources(container: V1Container) -> Resources:
resources: V1ResourceRequirements = container.resources
if not resources:
return Resources()
limits = resources.limits or {}
requests = resources.requests or {}
gpu_limit = limits.get(
"nvidia.com/gpu", limits.get("amd.com/gpu")
)
return Resources(
request_memory=requests.get("memory"),
request_cpu=requests.get("cpu"),
limit_memory=limits.get("memory"),
limit_cpu=limits.get("cpu"),
limit_gpu=gpu_limit,
)
def make_pod(self, namespace, worker_uuid, pod_id, dag_id, task_id,
execution_date, try_number, airflow_command,
kube_executor_config):
volumes_dict, volume_mounts_dict = self._get_volumes_and_mounts()
worker_init_container_spec = self._get_init_containers()
resources = Resources(
request_memory=kube_executor_config.request_memory,
request_cpu=kube_executor_config.request_cpu,
limit_memory=kube_executor_config.limit_memory,
limit_cpu=kube_executor_config.limit_cpu,
limit_gpu=kube_executor_config.limit_gpu)
gcp_sa_key = kube_executor_config.gcp_service_account_key
annotations = dict(kube_executor_config.annotations
) or self.kube_config.kube_annotations
if gcp_sa_key:
annotations['iam.cloud.google.com/service-account'] = gcp_sa_key
volumes = [value for value in volumes_dict.values()
] + kube_executor_config.volumes
volume_mounts = [value for value in volume_mounts_dict.values()
] + kube_executor_config.volume_mounts
affinity = kube_executor_config.affinity or self.kube_config.kube_affinity
tolerations = kube_executor_config.tolerations or self.kube_config.kube_tolerations
return Pod(
namespace=namespace,
name=pod_id,
image=kube_executor_config.image or self.kube_config.kube_image,
image_pull_policy=(kube_executor_config.image_pull_policy
or self.kube_config.kube_image_pull_policy),
cmds=airflow_command,
labels=self._get_labels(
kube_executor_config.labels, {
'airflow-worker': worker_uuid,
'dag_id': dag_id,
'task_id': task_id,
'execution_date': execution_date,
'try_number': str(try_number),
'airflow_version': airflow_version.replace('+', '-'),
'kubernetes_executor': 'True',
}),
envs=self._get_environment(),
secrets=self._get_secrets(),
service_account_name=self.kube_config.worker_service_account_name,
image_pull_secrets=self.kube_config.image_pull_secrets,
init_containers=worker_init_container_spec,
volumes=volumes,
volume_mounts=volume_mounts,
resources=resources,
annotations=annotations,
node_selectors=(kube_executor_config.node_selectors
or self.kube_config.kube_node_selectors),
affinity=affinity,
tolerations=tolerations,
security_context=self._get_security_context(),
configmaps=self._get_configmaps())
def test_extract_limits(self):
# Test when resources is not empty
resources = Resources(limit_memory='1Gi', limit_cpu=1)
pod = Pod('v3.14', {}, [], resources=resources)
self.expected['spec']['containers'][0]['resources'] = {
'limits': {
'memory': '1Gi',
'cpu': 1
}
}
KubernetesRequestFactory.extract_resources(pod, self.input_req)
self.assertEqual(self.expected, self.input_req)
def make_pod(self, namespace, worker_uuid, pod_id, dag_id, task_id,
execution_date, airflow_command, kube_executor_config):
volumes_dict, volume_mounts_dict = self.init_volumes_and_mounts()
worker_init_container_spec = self._get_init_containers(
copy.deepcopy(volume_mounts_dict))
resources = Resources(
request_memory=kube_executor_config.request_memory,
request_cpu=kube_executor_config.request_cpu,
limit_memory=kube_executor_config.limit_memory,
limit_cpu=kube_executor_config.limit_cpu)
gcp_sa_key = kube_executor_config.gcp_service_account_key
annotations = dict(kube_executor_config.annotations)
if gcp_sa_key:
annotations['iam.cloud.google.com/service-account'] = gcp_sa_key
volumes = [value for value in volumes_dict.values()
] + kube_executor_config.volumes
volume_mounts = [value for value in volume_mounts_dict.values()
] + kube_executor_config.volume_mounts
affinity = kube_executor_config.affinity or self.kube_config.kube_affinity
tolerations = kube_executor_config.tolerations or self.kube_config.kube_tolerations
return Pod(
namespace=namespace,
name=pod_id,
image=kube_executor_config.image or self.kube_config.kube_image,
image_pull_policy=(kube_executor_config.image_pull_policy
or self.kube_config.kube_image_pull_policy),
cmds=airflow_command,
labels={
'airflow-worker': worker_uuid,
'dag_id': dag_id,
'task_id': task_id,
'execution_date': execution_date
},
envs=self._get_environment(),
secrets=self._get_secrets(),
service_account_name=self.kube_config.worker_service_account_name,
image_pull_secrets=self.kube_config.image_pull_secrets,
init_containers=worker_init_container_spec,
volumes=volumes,
volume_mounts=volume_mounts,
resources=resources,
annotations=annotations,
node_selectors=(kube_executor_config.node_selectors
or self.kube_config.kube_node_selectors),
affinity=affinity,
tolerations=tolerations)
def make_pod(self, namespace, worker_uuid, pod_id, dag_id, task_id, execution_date,
airflow_command, kube_executor_config):
volumes, volume_mounts = self.init_volumes_and_mounts()
worker_init_container_spec = self._get_init_containers(
copy.deepcopy(volume_mounts))
resources = Resources(
request_memory=kube_executor_config.request_memory,
request_cpu=kube_executor_config.request_cpu,
limit_memory=kube_executor_config.limit_memory,
limit_cpu=kube_executor_config.limit_cpu
)
gcp_sa_key = kube_executor_config.gcp_service_account_key
annotations = {
'iam.cloud.google.com/service-account': gcp_sa_key
} if gcp_sa_key else {}
return Pod(
namespace=namespace,
name=pod_id,
image=kube_executor_config.image or self.kube_config.kube_image,
image_pull_policy=(kube_executor_config.image_pull_policy or
self.kube_config.kube_image_pull_policy),
cmds=['bash', '-cx', '--'],
args=[airflow_command],
labels={
'airflow-worker': worker_uuid,
'dag_id': dag_id,
'task_id': task_id,
'execution_date': execution_date
},
envs=self._get_environment(),
secrets=self._get_secrets(),
service_account_name=self.kube_config.worker_service_account_name,
image_pull_secrets=self.kube_config.image_pull_secrets,
init_containers=worker_init_container_spec,
volumes=volumes,
volume_mounts=volume_mounts,
resources=resources,
annotations=annotations,
node_selectors=(kube_executor_config.node_selectors or
self.kube_config.kube_node_selectors),
affinity=kube_executor_config.affinity
)
def make_pod(self, namespace, worker_uuid, pod_id, dag_id, task_id,
execution_date, airflow_command, kube_executor_config):
volumes, volume_mounts = self.init_volumes_and_mounts()
worker_init_container_spec = self._get_init_containers(
copy.deepcopy(volume_mounts))
resources = Resources(
request_memory=kube_executor_config.request_memory,
request_cpu=kube_executor_config.request_cpu,
limit_memory=kube_executor_config.limit_memory,
limit_cpu=kube_executor_config.limit_cpu)
gcp_sa_key = kube_executor_config.gcp_service_account_key
annotations = {
'iam.cloud.google.com/service-account': gcp_sa_key
} if gcp_sa_key else {}
airflow_command = airflow_command.replace("-sd", "-i -sd")
airflow_path = airflow_command.split('-sd')[-1]
airflow_path = self.worker_airflow_home + airflow_path.split('/')[-1]
airflow_command = airflow_command.split(
'-sd')[0] + '-sd ' + airflow_path
return Pod(
namespace=namespace,
name=pod_id,
image=kube_executor_config.image or self.kube_config.kube_image,
cmds=['bash', '-cx', '--'],
args=[airflow_command],
labels={
'airflow-worker': worker_uuid,
'dag_id': dag_id,
'task_id': task_id,
'execution_date': execution_date
},
envs=self._get_environment(),
secrets=self._get_secrets(),
service_account_name=self.kube_config.worker_service_account_name,
image_pull_secrets=self.kube_config.image_pull_secrets,
init_containers=worker_init_container_spec,
volumes=volumes,
volume_mounts=volume_mounts,
resources=resources,
annotations=annotations)
def test_extract_all_resources(self):
# Test when resources is not empty
resources = Resources(request_memory='1Gi',
request_cpu=1,
limit_memory='2Gi',
limit_cpu=2,
limit_gpu=3)
pod = Pod('v3.14', {}, [], resources=resources)
self.expected['spec']['containers'][0]['resources'] = {
'requests': {
'memory': '1Gi',
'cpu': 1
},
'limits': {
'memory': '2Gi',
'cpu': 2,
'nvidia.com/gpu': 3
}
}
KubernetesRequestFactory.extract_resources(pod, self.input_req)
self.assertEqual(self.input_req, self.expected)
def make_pod(self, namespace, worker_uuid, pod_id, dag_id, task_id,
execution_date, airflow_command, kube_executor_config):
volumes, volume_mounts = self._get_volumes_and_mounts()
worker_init_container_spec = self._get_init_containers(
copy.deepcopy(volume_mounts))
resources = Resources(
request_memory=kube_executor_config.request_memory,
request_cpu=kube_executor_config.request_cpu,
limit_memory=kube_executor_config.limit_memory,
limit_cpu=kube_executor_config.limit_cpu)
gcp_sa_key = kube_executor_config.gcp_service_account_key
annotations = {
"iam.cloud.google.com/service-account": gcp_sa_key
} if gcp_sa_key else {}
return Pod(
namespace=namespace,
name=pod_id,
image=kube_executor_config.image or self.kube_config.kube_image,
cmds=["bash", "-cx", "--"],
args=[airflow_command],
labels={
"airflow-worker": worker_uuid,
"dag_id": dag_id,
"task_id": task_id,
"execution_date": execution_date
},
envs=self._get_environment(),
secrets=self._get_secrets(),
service_account_name=self.kube_config.worker_service_account_name,
image_pull_secrets=self.kube_config.image_pull_secrets,
init_containers=worker_init_container_spec,
volumes=volumes,
volume_mounts=volume_mounts,
resources=resources,
annotations=annotations)
concurrency=100)
start = KubernetesPodOperator(
namespace='dev',
image="python:3.6",
cmds=["python", "-c"],
arguments=[
f"from time import sleep; sleep(5); print('slept for 5 seconds')"
],
labels={"foo": "bar"},
name=f"sleeper-agent-start",
task_id=f"sleeper-agent-start-task",
get_logs=True,
dag=dag,
affinity=solver_affinity,
resources=Resources(request_cpu='100m'),
in_cluster=True)
end = KubernetesPodOperator(
namespace='dev',
image="python:3.6",
cmds=["python", "-c"],
arguments=[
f"from time import sleep; sleep(5); print('slept for 5 seconds')"
],
labels={"foo": "bar"},
name=f"sleeper-agent-end",
task_id=f"sleeper-agent-end-task",
get_logs=True,
dag=dag,
affinity=solver_affinity,
resources=Resources(request_cpu='100m'),
# USEFUL VARIABLES
gcloud_image = Variable.get("gcloud_image")
image = Variable.get("pod_image")
cluster_name = Variable.get("cluster_name")
service_account_email = Variable.get("service_account_email")
_MAIN_DAG_ID = "test-dag-spawn-pools"
pool_name = "test-dag-spawn-pool"
machine_type = "n1-standard-4"
num_nodes = 1
node_taint = "only-test-pod"
pod_resources = Resources(request_memory = "10000Mi",
limit_memory = "10000Mi",
limit_cpu = "2",
request_cpu = "2")
default_args = {
"owner": "airflow",
"start_date": datetime(2018, 12, 1),
"email": ["*****@*****.**"],
# Define some variables to cancel email sending
"email_on_failure": False,
"email_on_retry": False,
"retries": 3,
"retry_delay": timedelta(seconds=5),
'provide_context': True,
'catchup': True
}
gke_location = "us-central1-a"
gke_cluster_name = "bq-load-gke-1"
# Built from repo https://github.com/mozilla/probe-scraper
probe_scraper_image = 'gcr.io/moz-fx-data-airflow-prod-88e0/probe-scraper:latest'
probe_scraper_args = [
'python3', '-m', 'probe_scraper.runner', '--out-dir',
'/app/probe_data', '--cache-dir', '/app/probe_cache',
'--output-bucket', 'net-mozaws-prod-us-west-2-data-pitmo',
'--cache-bucket', 'telemetry-airflow-cache', '--env', 'prod'
]
# Cluster autoscaling works on pod resource requests, instead of usage
resources = Resources(request_memory='13312Mi',
request_cpu=None,
limit_memory='20480Mi',
limit_cpu=None)
probe_scraper = GKEPodOperator(
task_id="probe_scraper",
gcp_conn_id=gcp_conn_id,
project_id=connection.project_id,
location=gke_location,
cluster_name=gke_cluster_name,
name='probe-scraper',
namespace='default',
# Needed to scale the highmem pool from 0 -> 1
resources=resources,
# This python job requires 13 GB of memory, thus the highmem node pool
node_selectors={"nodepool": "highmem"},
# Due to the nature of the container run, we set get_logs to False,
def test_display_resources(self):
resources_string = str(Resources('1Gi', 1))
self.assertEqual(
resources_string,
"Request: [cpu: 1, memory: 1Gi], Limit: [cpu: None, memory: None, gpu: None]"
)
def _set_resources(self, resources):
return Resources(**resources) if resources else Resources()
# The environment variable the secret is specified under
deploy_target='SQL_CONN',
# Name of secret in Kubernetes
secret='airflow-secrets',
# Key of the secret within Kubernetes
key='sql_alchemy_conn')
resources_obj = Resources(
# Amount of memory requested, can use E, P, T, G, M, K, Ei, Pi,
# Ti, Gi, Mi, Ki as suffixes
request_memory='200Mi',
# Tells container to attempt to use specified number of cpus. One 'cpu' is
# equivalent to 1 AWS vCPU, 1 GCP Core, 1 Azure vCore or 1 Hyperthread on
# a bare-metal intel processor with Hyperthreading
# If CPU request exceeds all of your node's capacities it will fail to ever
# get scheduled.
request_cpu='2',
# If memory limit is exceeded the Pod goes up for termination, if no
# limit is specified there is no upper bound on the amount of memory it can
# use. You can also specify a default memory limit on a per-namespace basis
limit_memory='100Mi',
# If cpu request exceeds your node's capacity, it will fail to ever get
# scheduled. The m suffix stands for milli-cpus, therefore .5 cpu and 500m
# cpu are equivalent.
limit_cpu='500m')
# Creates a volume of type emptyDir without any configs
volumes = [Volume(name='empty-vol', configs={'emptyDir', {}})]
# Used to mount pod level volumes to a running container
volume_mounts = [
VolumeMount(name='test-vol-mount',
split_date_parts((from_date + timedelta(days=i)), partition)
for i in range(delta.days + 1)
]
seen = set()
parts = []
# loops through every day and pulls out unique set of date parts
for p in all_parts:
if p["part"] not in seen:
seen.add(p["part"])
parts.append({k: v for k, v in p.items()})
return parts
# Set the resources for the task pods
pod_resources = Resources(request_memory="500Mi", request_cpu="500m")
# Default settings for all DAGs
pod_defaults = dict(
get_logs=True,
image_pull_policy="Always",
in_cluster=True,
is_delete_operator_pod=True,
namespace=os.environ["NAMESPACE"],
# resources=pod_resources,
cmds=["/bin/bash", "-c"],
)
# Default environment variables for worker pods
env = os.environ.copy()
pod_env_vars = {
def make_pod(self, namespace, worker_uuid, pod_id, dag_id, task_id,
execution_date, try_number, airflow_command,
kube_executor_config):
volumes_dict, volume_mounts_dict = self._get_volumes_and_mounts()
worker_init_container_spec = self._get_init_containers()
#TODO add resourses merge
#
resources = Resources(
request_memory=kube_executor_config.request_memory,
request_cpu=kube_executor_config.request_cpu,
limit_memory=kube_executor_config.limit_memory,
limit_cpu=kube_executor_config.limit_cpu,
limit_gpu=kube_executor_config.limit_gpu)
#hack for creating default resources
#TODO make it safe
if resources.is_empty_resource_request(
) and self.kube_config.default_limits:
resources = Resources(
request_memory=deep_get(self.kube_config.default_limits,
'request.memory'),
request_cpu=deep_get(self.kube_config.default_limits,
'request.cpu'),
limit_memory=deep_get(self.kube_config.default_limits,
'limit.memory'),
limit_cpu=deep_get(self.kube_config.default_limits,
'limit.cpu'))
gcp_sa_key = kube_executor_config.gcp_service_account_key
annotations = dict(kube_executor_config.annotations
) or self.kube_config.kube_annotations
if gcp_sa_key:
annotations['iam.cloud.google.com/service-account'] = gcp_sa_key
volumes = [value for value in volumes_dict.values()
] + kube_executor_config.volumes
volume_mounts = [value for value in volume_mounts_dict.values()
] + kube_executor_config.volume_mounts
affinity = kube_executor_config.affinity or self.kube_config.kube_affinity
tolerations = kube_executor_config.tolerations or self.kube_config.kube_tolerations
return Pod(
namespace=namespace,
name=pod_id,
image=kube_executor_config.image or self.kube_config.kube_image,
image_pull_policy=(kube_executor_config.image_pull_policy
or self.kube_config.kube_image_pull_policy),
cmds=airflow_command,
labels=self._get_labels(
kube_executor_config.labels, {
'airflow-worker': worker_uuid,
'dag_id': dag_id,
'task_id': task_id,
'execution_date': execution_date,
'try_number': str(try_number),
}),
envs=self._get_environment(),
secrets=self._get_secrets(),
service_account_name=self.kube_config.worker_service_account_name,
image_pull_secrets=self.kube_config.image_pull_secrets,
init_containers=worker_init_container_spec,
volumes=volumes,
volume_mounts=volume_mounts,
resources=resources,
annotations=annotations,
node_selectors=(kube_executor_config.node_selectors
or self.kube_config.kube_node_selectors),
affinity=affinity,
tolerations=tolerations,
security_context=self._get_security_context(),
configmaps=self._get_configmaps(),
scheduler_name=self.kube_config.scheduler_name)
def setUp(self):
self.simple_pod_request_factory = SimplePodRequestFactory()
self.xcom_pod_request_factory = ExtractXcomPodRequestFactory()
self.pod = Pod(
image='busybox',
envs={
'ENVIRONMENT': 'prod',
'LOG_LEVEL': 'warning'
},
name='myapp-pod',
cmds=['sh', '-c', 'echo Hello Kubernetes!'],
labels={'app': 'myapp'},
image_pull_secrets='pull_secret_a,pull_secret_b',
configmaps=['configmap_a', 'configmap_b'],
ports=[{'name': 'foo', 'containerPort': 1234}],
resources=Resources('1Gi', 1, '2Gi', 2, 1),
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'),
],
security_context={
'runAsUser': 1000,
'fsGroup': 2000,
}
)
self.maxDiff = None
self.expected = {
'apiVersion': 'v1',
'kind': 'Pod',
'metadata': {
'name': 'myapp-pod',
'labels': {'app': 'myapp'},
'annotations': {}},
'spec': {
'containers': [{
'name': 'base',
'image': 'busybox',
'command': [
'sh', '-c', 'echo Hello Kubernetes!'
],
'imagePullPolicy': 'IfNotPresent',
'args': [],
'env': [{
'name': 'ENVIRONMENT',
'value': 'prod'
}, {
'name': 'LOG_LEVEL',
'value': 'warning'
}, {
'name': 'TARGET',
'valueFrom': {
'secretKeyRef': {
'name': 'secret_b',
'key': 'source_b'
}
}
}],
'envFrom': [{
'secretRef': {
'name': 'secret_a'
}
}, {
'configMapRef': {
'name': 'configmap_a'
}
}, {
'configMapRef': {
'name': 'configmap_b'
}
}],
'resources': {
'requests': {
'memory': '1Gi',
'cpu': 1
},
'limits': {
'memory': '2Gi',
'cpu': 2,
'nvidia.com/gpu': 1
},
},
'ports': [{'name': 'foo', 'containerPort': 1234}],
'volumeMounts': [{
'mountPath': '/etc/foo',
'name': 'secretvol0',
'readOnly': True
}]
}],
'restartPolicy': 'Never',
'nodeSelector': {},
'volumes': [{
'name': 'secretvol0',
'secret': {
'secretName': 'secret_b'
}
}],
'imagePullSecrets': [
{'name': 'pull_secret_a'},
{'name': 'pull_secret_b'}
],
'affinity': {},
'securityContext': {
'runAsUser': 1000,
'fsGroup': 2000,
},
}
}