def get_node_affinity(
affinities: PVector["NodeAffinity"]) -> Optional[V1NodeAffinity]:
# convert NodeAffinity into V1NodeSelectorRequirement
match_expressions = []
for aff in affinities:
op = aff["operator"]
val = aff["value"]
# operator and value are assumed to be validated
if op in {NodeAffinityOperator.IN, NodeAffinityOperator.NOT_IN}:
val = [str(v) for v in val]
elif op in {NodeAffinityOperator.GT, NodeAffinityOperator.LT}:
val = [str(val)]
elif op in {
NodeAffinityOperator.EXISTS,
NodeAffinityOperator.DOES_NOT_EXIST
}:
val = []
else:
continue
match_expressions.append(
V1NodeSelectorRequirement(key=str(aff["key"]),
operator=op,
values=val))
# package into V1NodeAffinity
if not match_expressions:
return None
return V1NodeAffinity(
# this means that the selectors are only used during scheduling.
# changing it while the pod is running will not cause an eviction.
required_during_scheduling_ignored_during_execution=V1NodeSelector(
node_selector_terms=[
V1NodeSelectorTerm(match_expressions=match_expressions)
], ), )
def my_pipeline():
affinity = V1Affinity(
node_affinity=V1NodeAffinity(
required_during_scheduling_ignored_during_execution=V1NodeSelector(
node_selector_terms=[V1NodeSelectorTerm(
match_expressions=[V1NodeSelectorRequirement(
key='beta.kubernetes.io/instance-type', operator='In', values=['p2.xlarge'])])])))
some_op().add_affinity(affinity)
def affinity_pipeline(
):
"""A pipeline with affinity"""
affinity = V1Affinity(
node_affinity=V1NodeAffinity(
required_during_scheduling_ignored_during_execution=V1NodeSelector(
node_selector_terms=[V1NodeSelectorTerm(
match_expressions=[V1NodeSelectorRequirement(
key='kubernetes.io/os',
operator='In',
values=['linux'])])])))
echo_op().add_affinity(affinity)
def test_get_node_affinity_ok():
affinities = pvector([
NodeAffinity(key="label0", operator="In", value=[1, 2, 3]),
NodeAffinity(key="label1", operator="NotIn", value=[3, 2, 1]),
NodeAffinity(key="label2", operator="Gt", value=1),
NodeAffinity(key="label3", operator="Lt", value=2),
NodeAffinity(key="label4", operator="Exists", value="hi"),
NodeAffinity(key="label5", operator="DoesNotExist", value="bye"),
])
assert get_node_affinity(affinities) == V1NodeAffinity(
required_during_scheduling_ignored_during_execution=V1NodeSelector(
node_selector_terms=[
V1NodeSelectorTerm(match_expressions=[
V1NodeSelectorRequirement(
key="label0",
operator="In",
values=["1", "2", "3"],
),
V1NodeSelectorRequirement(
key="label1",
operator="NotIn",
values=["3", "2", "1"],
),
V1NodeSelectorRequirement(
key="label2",
operator="Gt",
values=["1"],
),
V1NodeSelectorRequirement(
key="label3",
operator="Lt",
values=["2"],
),
V1NodeSelectorRequirement(
key="label4",
operator="Exists",
values=[],
),
V1NodeSelectorRequirement(
key="label5",
operator="DoesNotExist",
values=[],
),
])
]))
def _set_preemptible(task):
task.add_toleration(toleration)
node_selector_term = V1NodeSelectorTerm(match_expressions=[
V1NodeSelectorRequirement(key='cloud.google.com/gke-preemptible',
operator='In',
values=['true'])
])
if hard_constraint:
node_affinity = V1NodeAffinity(
required_during_scheduling_ignored_during_execution=
V1NodeSelector(node_selector_terms=[node_selector_term]))
else:
node_affinity = V1NodeAffinity(
preferred_during_scheduling_ignored_during_execution=
V1PreferredSchedulingTerm(preference=node_selector_term,
weight=50))
affinity = V1Affinity(node_affinity=node_affinity)
task.add_affinity(affinity=affinity)
return task
