def schedule_pod_on_node()

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}')