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 = {