in gpudirect-tcpxo/topology-scheduler/schedule-daemon.py [0:0]
def schedule_pod_on_node(v1, pod_name, pod_namespace, node, gate_name):
"""Schedules a pod on a given node."""
try:
pod = v1.read_namespaced_pod(pod_name, pod_namespace)
if any(gate.name == gate_name for gate in pod.spec.scheduling_gates):
new_gates = [
gate for gate in pod.spec.scheduling_gates if gate.name != gate_name
]
pod.spec.affinity = {
'nodeAffinity': {
'requiredDuringSchedulingIgnoredDuringExecution': {
'nodeSelectorTerms': [{
'matchExpressions': [{
'key': 'kubernetes.io/hostname',
'operator': 'In',
'values': [node['name']],
}]
}]
}
}
}
pod.spec.scheduling_gates = new_gates
v1.replace_namespaced_pod(pod_name, pod_namespace, pod)
print(
'Pod %s/%s scheduled on %s with topology %s', pod_namespace, pod_name, node['name'], node_topology_key(node)
)
except ApiException as e:
print(f'Exception when removing scheduling gate: {e}')