schedulers/kubernetes_scheduler: add support for resource instance-type node selectors

[d4l3k]

This allows specifying specific instance types when scheduling kubernetes jobs. It uses node_selectors and the node.kubernetes.io/instance-type label on nodes to limit pods to specific instances.

To avoid instance type cpu and memory hitting issues with the node reserved cpu/mem this will subtract a small amount of CPU and memory from the requested resources. Limits remains the same.

Also adds g4dn.xlarge resource type.

• https://kubernetes.io/docs/tasks/configure-pod-container/assign-pods-nodes/
• https://kubernetes.io/docs/tasks/administer-cluster/reserve-compute-resources/

Test plan:

Unit tests, updated kube dist integration test to specify instance type

spec:
  containers:
  - command:
    - python
    - -m
    - torch.distributed.run
    - --rdzv_backend
    - etcd
    - --rdzv_endpoint
    - etcd-server:2379
    - --rdzv_id
    - cv-trainer-smvh1095z11h5
    - --nnodes
    - "2"
    - --nproc_per_node
    - "1"
    - -m
    - torchx.examples.apps.lightning_classy_vision.train
    - --load_path
    - ""
    - --log_path
    - /tmp/logs
    - --epochs
    - "1"
    - --output_path
    - s3://torchx-test/integration-tests/runner_wh5wn4cz4b7fcd/output
    env:
    - name: TORCHX_RANK0_HOST
      value: localhost
    - name: VC_WORKER_0_HOSTS
      valueFrom:
        configMapKeyRef:
          key: VC_WORKER_0_HOSTS
          name: cv-trainer-smvh1095z11h5-svc
    - name: VC_WORKER_0_NUM
      valueFrom:
        configMapKeyRef:
          key: VC_WORKER_0_NUM
          name: cv-trainer-smvh1095z11h5-svc
    - name: VC_WORKER_1_HOSTS
      valueFrom:
        configMapKeyRef:
          key: VC_WORKER_1_HOSTS
          name: cv-trainer-smvh1095z11h5-svc
    - name: VC_WORKER_1_NUM
      valueFrom:
        configMapKeyRef:
          key: VC_WORKER_1_NUM
          name: cv-trainer-smvh1095z11h5-svc
    - name: VK_TASK_INDEX
      value: "0"
    - name: VC_TASK_INDEX
      value: "0"
    image: 495572122715.dkr.ecr.us-west-2.amazonaws.com/torchx/integration-tests:canary_runner_wh5wn4cz4b7fcd_torchx
    imagePullPolicy: IfNotPresent
    name: worker-0
    resources:
      limits:
        cpu: "2"
        memory: 8192M
        nvidia.com/gpu: "1"
      requests:
        cpu: 1900m
        memory: 7168M
        nvidia.com/gpu: "1"
    terminationMessagePath: /dev/termination-log
    terminationMessagePolicy: File
    volumeMounts:
    - mountPath: /dev/shm
      name: dshm
    - mountPath: /etc/volcano
      name: cv-trainer-smvh1095z11h5-svc
    - mountPath: /var/run/secrets/kubernetes.io/serviceaccount
      name: default-token-z8vb9
      readOnly: true
  dnsPolicy: ClusterFirst
  enableServiceLinks: true
  hostname: cv-trainer-smvh1095z11h5-worker-0-0
  nodeName: ip-192-168-16-165.us-west-2.compute.internal
  nodeSelector:
    node.kubernetes.io/instance-type: p3.2xlarge

https://github.com/pytorch/torchx/runs/5698855673?check_suite_focus=true

[codecov]

Codecov Report

Merging #433 (984a8c0) into main (434c013) will increase coverage by 0.02%.
The diff coverage is 100.00%.

@@            Coverage Diff             @@
##             main     #433      +/-   ##
==========================================
+ Coverage   94.52%   94.54%   +0.02%     
==========================================
  Files          64       64              
  Lines        3782     3798      +16     
==========================================
+ Hits         3575     3591      +16     
  Misses        207      207              

Impacted Files	Coverage Δ
torchx/schedulers/kubernetes_scheduler.py	93.08% <100.00%> (+0.40%)	⬆️
torchx/specs/__init__.py	96.15% <100.00%> (ø)
torchx/specs/named_resources_aws.py	100.00% <100.00%> (ø)

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 434c013...984a8c0. Read the comment docs.

[facebook-github-bot]

@d4l3k has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.