def test_create_with_affinity(self):
name_base = "test"
affinity = {
"nodeAffinity": {
"preferredDuringSchedulingIgnoredDuringExecution": [
{
"weight": 1,
"preference": {
"matchExpressions": [{"key": "disktype", "operator": "In", "values": ["ssd"]}]
},
}
]
}
}
k = KubernetesPodOperator(
namespace="default",
image="ubuntu:16.04",
cmds=["bash", "-cx"],
arguments=["echo 10"],
labels={"foo": "bar"},
name=name_base,
task_id="task",
in_cluster=False,
do_xcom_push=False,
cluster_context="default",
affinity=affinity,
)
pod = k.create_pod_request_obj()
sanitized_pod = self.sanitize_for_serialization(pod)
assert isinstance(pod.spec.affinity, k8s.V1Affinity)
assert sanitized_pod["spec"]["affinity"] == affinity
k8s_api_affinity = k8s.V1Affinity(
node_affinity=k8s.V1NodeAffinity(
preferred_during_scheduling_ignored_during_execution=[
k8s.V1PreferredSchedulingTerm(
weight=1,
preference=k8s.V1NodeSelectorTerm(
match_expressions=[
k8s.V1NodeSelectorRequirement(key="disktype", operator="In", values=["ssd"])
]
),
)
]
),
)
k = KubernetesPodOperator(
namespace="default",
image="ubuntu:16.04",
cmds=["bash", "-cx"],
arguments=["echo 10"],
labels={"foo": "bar"},
name=name_base,
task_id="task",
in_cluster=False,
do_xcom_push=False,
cluster_context="default",
affinity=k8s_api_affinity,
)
pod = k.create_pod_request_obj()
sanitized_pod = self.sanitize_for_serialization(pod)
assert isinstance(pod.spec.affinity, k8s.V1Affinity)
assert sanitized_pod["spec"]["affinity"] == affinity
def __init__( # pylint: disable=too-many-arguments,too-many-locals
# fmt: on
self,
*,
namespace: Optional[str] = None,
image: Optional[str] = None,
name: Optional[str] = None,
cmds: Optional[List[str]] = None,
arguments: Optional[List[str]] = None,
ports: Optional[List[k8s.V1ContainerPort]] = None,
volume_mounts: Optional[List[k8s.V1VolumeMount]] = None,
volumes: Optional[List[k8s.V1Volume]] = None,
env_vars: Optional[List[k8s.V1EnvVar]] = None,
env_from: Optional[List[k8s.V1EnvFromSource]] = None,
secrets: Optional[List[Secret]] = None,
in_cluster: Optional[bool] = None,
cluster_context: Optional[str] = None,
labels: Optional[Dict] = None,
reattach_on_restart: bool = True,
startup_timeout_seconds: int = 120,
get_logs: bool = True,
image_pull_policy: str = 'IfNotPresent',
annotations: Optional[Dict] = None,
resources: Optional[k8s.V1ResourceRequirements] = None,
affinity: Optional[k8s.V1Affinity] = None,
config_file: Optional[str] = None,
node_selectors: Optional[dict] = None,
node_selector: Optional[dict] = None,
image_pull_secrets: Optional[List[k8s.V1LocalObjectReference]] = None,
service_account_name: str = 'default',
is_delete_operator_pod: bool = False,
hostnetwork: bool = False,
tolerations: Optional[List[k8s.V1Toleration]] = None,
security_context: Optional[Dict] = None,
dnspolicy: Optional[str] = None,
schedulername: Optional[str] = None,
full_pod_spec: Optional[k8s.V1Pod] = None,
init_containers: Optional[List[k8s.V1Container]] = None,
log_events_on_failure: bool = False,
do_xcom_push: bool = False,
pod_template_file: Optional[str] = None,
priority_class_name: Optional[str] = None,
pod_runtime_info_envs: List[PodRuntimeInfoEnv] = None,
termination_grace_period: Optional[int] = None,
configmaps: Optional[str] = None,
**kwargs,
) -> None:
if kwargs.get('xcom_push') is not None:
raise AirflowException("'xcom_push' was deprecated, use 'do_xcom_push' instead")
super().__init__(resources=None, **kwargs)
self.do_xcom_push = do_xcom_push
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.env_vars = convert_env_vars(env_vars) if env_vars else []
if pod_runtime_info_envs:
self.env_vars.extend([convert_pod_runtime_info_env(p) for p in pod_runtime_info_envs])
env_vars = self.add_template_fields_to_env_vars(env_vars)
self.env_vars = env_vars
self.env_from = env_from or []
if configmaps:
self.env_from.extend([convert_configmap(c) for c in configmaps])
self.ports = [convert_port(p) for p in ports] if ports else []
self.volume_mounts = [convert_volume_mount(v) for v in volume_mounts] if volume_mounts else []
self.volumes = [convert_volume(volume) for volume in volumes] if volumes else []
self.secrets = secrets or []
self.in_cluster = in_cluster
self.cluster_context = cluster_context
self.reattach_on_restart = reattach_on_restart
self.get_logs = get_logs
self.image_pull_policy = image_pull_policy
if node_selectors:
# Node selectors is incorrect based on k8s API
warnings.warn("node_selectors is deprecated. Please use node_selector instead.")
self.node_selector = node_selectors or {}
elif node_selector:
self.node_selector = node_selector or {}
else:
self.node_selector = None
self.annotations = annotations or {}
self.affinity = convert_affinity(affinity) if affinity else k8s.V1Affinity()
self.k8s_resources = convert_resources(resources) if resources else {}
self.config_file = config_file
self.image_pull_secrets = convert_image_pull_secrets(image_pull_secrets) if image_pull_secrets else []
self.service_account_name = service_account_name
self.is_delete_operator_pod = is_delete_operator_pod
self.hostnetwork = hostnetwork
self.tolerations = [convert_toleration(toleration) for toleration in tolerations] \
if tolerations else []
self.security_context = security_context or {}
self.dnspolicy = dnspolicy
self.schedulername = schedulername
self.full_pod_spec = full_pod_spec
self.init_containers = init_containers or []
self.log_events_on_failure = log_events_on_failure
self.priority_class_name = priority_class_name
self.pod_template_file = pod_template_file
self.name = self._set_name(name)
self.termination_grace_period = termination_grace_period
self.client: CoreV1Api = None
self.pod: k8s.V1Pod = None
def test_create_with_affinity(self):
name_base = 'test'
affinity = {
'nodeAffinity': {
'preferredDuringSchedulingIgnoredDuringExecution': [{
"weight": 1,
"preference": {
"matchExpressions": [{
"key": "disktype",
"operator": "In",
"values": ["ssd"]
}]
},
}]
}
}
k = KubernetesPodOperator(
namespace='default',
image="ubuntu:16.04",
cmds=["bash", "-cx"],
arguments=["echo 10"],
labels={"foo": "bar"},
name=name_base,
task_id="task",
in_cluster=False,
do_xcom_push=False,
cluster_context='default',
affinity=affinity,
)
result = k.create_pod_request_obj()
client = ApiClient()
self.assertEqual(type(result.spec.affinity), k8s.V1Affinity)
self.assertEqual(
client.sanitize_for_serialization(result)['spec']['affinity'],
affinity)
k8s_api_affinity = k8s.V1Affinity(node_affinity=k8s.V1NodeAffinity(
preferred_during_scheduling_ignored_during_execution=[
k8s.V1PreferredSchedulingTerm(
weight=1,
preference=k8s.V1NodeSelectorTerm(match_expressions=[
k8s.V1NodeSelectorRequirement(
key="disktype", operator="In", values=["ssd"])
]),
)
]), )
k = KubernetesPodOperator(
namespace='default',
image="ubuntu:16.04",
cmds=["bash", "-cx"],
arguments=["echo 10"],
labels={"foo": "bar"},
name=name_base,
task_id="task",
in_cluster=False,
do_xcom_push=False,
cluster_context='default',
affinity=k8s_api_affinity,
)
result = k.create_pod_request_obj()
self.assertEqual(type(result.spec.affinity), k8s.V1Affinity)
self.assertEqual(
client.sanitize_for_serialization(result)['spec']['affinity'],
affinity)
args=["echo 10"],
)
affinity = k8s.V1Affinity(
node_affinity=k8s.V1NodeAffinity(
preferred_during_scheduling_ignored_during_execution=[
k8s.V1PreferredSchedulingTerm(
weight=1,
preference=k8s.V1NodeSelectorTerm(match_expressions=[
k8s.V1NodeSelectorRequirement(
key="disktype", operator="in", values=["ssd"])
]),
)
]),
pod_affinity=k8s.V1PodAffinity(
required_during_scheduling_ignored_during_execution=[
k8s.V1WeightedPodAffinityTerm(
weight=1,
pod_affinity_term=k8s.V1PodAffinityTerm(
label_selector=k8s.V1LabelSelector(match_expressions=[
k8s.V1LabelSelectorRequirement(
key="security", operator="In", values="S1")
]),
topology_key="failure-domain.beta.kubernetes.io/zone",
),
)
]),
)
tolerations = [k8s.V1Toleration(key="key", operator="Equal", value="value")]
image=f"{worker_container_repository}:{worker_container_tag}"),
]))
}
@task(executor_config=kube_exec_config_special)
def base_image_override_task():
print_stuff()
base_image_task = base_image_override_task()
# Use k8s_client.V1Affinity to define node affinity
k8s_affinity = k8s.V1Affinity(pod_anti_affinity=k8s.V1PodAntiAffinity(
required_during_scheduling_ignored_during_execution=[
k8s.V1PodAffinityTerm(
label_selector=k8s.V1LabelSelector(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 = {
