Add support for Sakkara

Co-authored-by: Lixiang "Eric" Luo <lgl@@users.noreply.github.com>

Author: dgrove-oss

tools/pytorchjob-generator/chart/templates/_helpers.tpl

@@ -35,7 +35,7 @@ metadata:
 terminationGracePeriodSeconds: {{ .Values.terminationGracePeriodSeconds }}
 {{- end }}
 {{- if .Values.bypassCoscheduler }}
-schedulerName: default-scheduler
+schedulerName: {{ .Values.schedulerName }}
 {{- end }}
 priorityClassName: {{ .Values.priority }}
 affinity:
@@ -81,8 +81,14 @@ envFrom:
   - configMapRef:
       name: {{ .Values.ncclGdrEnvConfigMap }}
 {{- end }}
-{{- if or .Values.environmentVariables .Values.sshGitCloneConfig .Values.mountNVMe .Values.topologyFileConfigMap }}
+{{- if or .Values.environmentVariables .Values.sshGitCloneConfig .Values.mountNVMe .Values.topologyFileConfigMap ( eq .Values.schedulerName "sakkara" ) }}
 env:
+    {{- if eq .Values.schedulerName "sakkara" }}
+    - name: SAKKARA_RANK
+      valueFrom:
+        fieldRef:
+          fieldPath: metadata.labels['sakkara.member.rank']
+    {{- end }}
     {{- if .Values.topologyFileConfigMap }}
     - name: NCCL_TOPO_FILE
       value: /var/run/nvidia-topologyd/virtualTopology.xml
@@ -146,6 +152,10 @@ command:
       #
       # User commands
       #
+      {{- if eq .Values.schedulerName "sakkara" }}
+      echo "Sakkara is enabled: using Sakkara-assigned rank instead of the default PyTorchJob rank"
+      export RANK=$SAKKARA_RANK
+      {{- end }}
       {{- range $command := .Values.setupCommands }}
       {{ $command }}
       {{- end }}

tools/pytorchjob-generator/chart/values.schema.json

@@ -115,6 +115,10 @@
             { "type": "array" }
         ]},
         "bypassCoscheduler": { "type": "boolean" },
+        "schedulerName": { "oneOf": [
+            { "type": "null" },
+            { "type": "string", "enum": ["sakkara", "default-scheduler" ] }
+        ]},
         "serviceAccountName":  { "oneOf" : [
             { "type": "null" },
             { "$ref": "#/$defs/rfc1123Label" }

tools/pytorchjob-generator/examples/helloworld-sakkara.settings.yaml

@@ -0,0 +1,28 @@
+namespace: my-namespace       # namespace to deploy to (required)
+jobName: my-job               # name of the generated AppWrapper and PyTorchJob objects (required)
+queueName: default-queue      # local queue to submit to (default: default-queue)
+
+bypassCoscheduler: true
+schedulerName: sakkara
+# If additional constraints are used, specify the configmap here:
+#customLabels:
+#  - key: sakkara.group.name
+#    value: my-topogrp-0
+
+numPods: 4                    # total pod count including master and worker pods (default: 1)
+numCpusPerPod: 500m           # requested number of cpus per pod (default: 1)
+numGpusPerPod: 8              # requested number of gpus per pod (default: 0)
+totalMemoryPerPod: 1Gi        # requested amount of memory per pod (default: 1Gi)
+
+priority: default-priority    # default-priority (default), low-priority, or high-priority
+
+# container image for the pods (required)
+containerImage: ghcr.io/foundation-model-stack/base:pytorch-latest-nightly-20230126
+
+# setup commands to run in each pod (optional)
+setupCommands:
+- git clone https://github.com/dbarnett/python-helloworld
+- cd python-helloworld
+
+# main program to invoke via torchrun (optional)
+mainProgram: helloworld.py