diff --git a/charts/hami/README.md b/charts/hami/README.md index 9df0a684f5..2917e930c5 100644 --- a/charts/hami/README.md +++ b/charts/hami/README.md @@ -77,6 +77,8 @@ This document provides detailed descriptions of all configurable values paramete | `scheduler.livenessProbe` | Whether to enable liveness probe | `false` | | `scheduler.leaderElect` | Whether to enable leader election | `true` | | `scheduler.replicas` | Number of replicas | `1` | +| `scheduler.podDisruptionBudget.minAvailable` | Minimum number of available scheduler pods during voluntary disruptions (only rendered when `scheduler.leaderElect` is `true`) | `1` | +| `scheduler.podDisruptionBudget.maxUnavailable` | Maximum number of unavailable scheduler pods during voluntary disruptions; set `minAvailable` to `null` when using this | `nil` | ### Kube Scheduler Configuration diff --git a/charts/hami/templates/scheduler/pdb.yaml b/charts/hami/templates/scheduler/pdb.yaml new file mode 100644 index 0000000000..84ed916ec7 --- /dev/null +++ b/charts/hami/templates/scheduler/pdb.yaml @@ -0,0 +1,27 @@ +{{- if .Values.scheduler.leaderElect }} +apiVersion: policy/v1 +kind: PodDisruptionBudget +metadata: + name: {{ include "hami-vgpu.scheduler" . }} + namespace: {{ include "hami-vgpu.namespace" . }} + labels: + app.kubernetes.io/component: hami-scheduler + {{- include "hami-vgpu.labels" . | nindent 4 }} + {{- with .Values.global.labels }} + {{- toYaml . | nindent 4 }} + {{- end }} + {{- if .Values.global.annotations }} + annotations: {{ toYaml .Values.global.annotations | nindent 4}} + {{- end }} +spec: + {{- with .Values.scheduler.podDisruptionBudget.minAvailable }} + minAvailable: {{ . }} + {{- end }} + {{- with .Values.scheduler.podDisruptionBudget.maxUnavailable }} + maxUnavailable: {{ . }} + {{- end }} + selector: + matchLabels: + app.kubernetes.io/component: hami-scheduler + {{- include "hami-vgpu.selectorLabels" . | nindent 6 }} +{{- end }} diff --git a/charts/hami/values.yaml b/charts/hami/values.yaml index a7e13d1b11..2e662688e9 100644 --- a/charts/hami/values.yaml +++ b/charts/hami/values.yaml @@ -83,6 +83,11 @@ scheduler: leaderElect: true # when leaderElect is true, replicas is available, otherwise replicas is 1. replicas: 1 + # PodDisruptionBudget for the hami-scheduler Deployment, only rendered when leaderElect is true. + # Set minAvailable to null when overriding maxUnavailable, since only one of the two may be set. + podDisruptionBudget: + minAvailable: 1 + maxUnavailable: kubeScheduler: # @param enabled indicate whether to run kube-scheduler container in the scheduler pod, it's true by default. enabled: true diff --git a/cmd/hami-cli/README.md b/cmd/hami-cli/README.md new file mode 100644 index 0000000000..f5604d189b --- /dev/null +++ b/cmd/hami-cli/README.md @@ -0,0 +1,78 @@ +# hami-cli + +A small, read-only inspection tool for HAMi vGPU device allocations. + +HAMi's scheduler writes device allocation decisions into pod and node +annotations under the `hami.io/` prefix (see +[docs/develop/protocol.md](../../docs/develop/protocol.md)). `hami-cli` +decodes those existing annotations and prints them as a table, so you don't +have to hand-decode `kubectl get pod -o yaml` output to answer "which pod +holds which GPU slice on which node". + +It is a pure reader of cluster state: it makes no scheduler-side changes and +requires no GPU hardware to build, test, or run against a cluster that has +HAMi-managed pods. + +## Installation + +Build from source with the rest of HAMi's binaries: + +```bash +make build +# binary is written to bin/hami-cli +``` + +Or build just this binary directly: + +```bash +go build -o bin/hami-cli ./cmd/hami-cli +``` + +## Usage + +`hami-cli` uses the same kubeconfig resolution as `kubectl`: it reads +`$KUBECONFIG` if set, falls back to `~/.kube/config`, and falls back to +in-cluster configuration when running inside a pod. + +List every HAMi device allocation in the cluster: + +```bash +hami-cli get allocations +``` + +```text +NODE NAMESPACE POD CONTAINER DEVICE TYPE DEVICE UUID MEMORY CORE +node67-4v100 default train-job trainer NVIDIA GPU-0fc3eda5-e98b-a25b-5b0d-cf5c855d1448 3000 0 +``` + +Filter by node or namespace: + +```bash +hami-cli get allocations --node node67-4v100 +hami-cli get allocations --namespace ml-team +``` + +## Vendor coverage + +`hami-cli` does not hardcode a list of supported vendors. Every HAMi device +backend registers its per-pod annotation key under the pattern +`hami.io/-devices-to-allocate` (documented in +[docs/develop/protocol.md](../../docs/develop/protocol.md)); `hami-cli` +matches any annotation following that pattern and decodes it with the same +`device.DecodePodDevices` function the scheduler itself uses. This means +allocations from every currently-supported vendor (NVIDIA, Cambricon, +Ascend, AMD, Hygon, Iluvatar, Kunlun, Metax, Mthreads, Biren, Enflame, +AWS Neuron, Vast.ai) are decoded uniformly, including vendors whose +annotation key is only known at runtime via chart configuration (Ascend, +Iluvatar). + +If a pod's HAMi annotations are malformed, `hami-cli` prints a warning to +stderr and skips that pod rather than failing the whole command. + +## Limitations + +- Shows what the scheduler *requested/recorded* for each container, not the + device's live runtime utilization (see `vGPUmonitor`'s Prometheus metrics + for that). +- Node-side device capacity/registration (`hami.io/node-*-register`) is not + yet cross-referenced against pod allocations in the table output. diff --git a/cmd/hami-cli/allocations.go b/cmd/hami-cli/allocations.go new file mode 100644 index 0000000000..ee6b5d18b0 --- /dev/null +++ b/cmd/hami-cli/allocations.go @@ -0,0 +1,190 @@ +/* +Copyright 2024 The HAMi Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package main + +import ( + "context" + "fmt" + "io" + "os" + "regexp" + "sort" + "text/tabwriter" + + "github.com/spf13/cobra" + corev1 "k8s.io/api/core/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/client-go/kubernetes" + + "github.com/Project-HAMi/HAMi/pkg/device" + "github.com/Project-HAMi/HAMi/pkg/util" + "github.com/Project-HAMi/HAMi/pkg/util/client" +) + +// podDeviceAllocateAnnotation matches the per-vendor "hami.io/-devices-to-allocate" +// annotation keys documented in docs/develop/protocol.md. Every device backend registers +// its own key under this suffix (see device.InRequestDevices assignments), so matching by +// pattern lets hami-cli decode any vendor's allocations without importing vendor packages. +var podDeviceAllocateAnnotation = regexp.MustCompile(`^hami\.io/(.+)-devices-to-allocate$`) + +type allocationRow struct { + Node string + Namespace string + Pod string + Container string + DeviceType string + DeviceUUID string + RequestedMem int32 + RequestedCore int32 +} + +var allocationsCmd = &cobra.Command{ + Use: "allocations", + Short: "List HAMi device allocations decoded from hami.io/* pod annotations", + RunE: func(cmd *cobra.Command, args []string) error { + c, err := client.NewClient() + if err != nil { + return fmt.Errorf("failed to build kubernetes client: %w", err) + } + return runAllocations(cmd.Context(), c, cmd.OutOrStdout()) + }, +} + +var ( + nodeFilter string + namespaceFilter string +) + +func init() { + allocationsCmd.Flags().StringVar(&nodeFilter, "node", "", "only show allocations on this node") + allocationsCmd.Flags().StringVar(&namespaceFilter, "namespace", "", "only show pods in this namespace (default: all namespaces)") +} + +func runAllocations(ctx context.Context, c kubernetes.Interface, out io.Writer) error { + pods, err := c.CoreV1().Pods(namespaceFilter).List(ctx, metav1.ListOptions{}) + if err != nil { + return fmt.Errorf("failed to list pods: %w", err) + } + + rows := collectAllocationRows(pods.Items) + if nodeFilter != "" { + filtered := rows[:0] + for _, r := range rows { + if r.Node == nodeFilter { + filtered = append(filtered, r) + } + } + rows = filtered + } + + printAllocationTable(out, rows) + return nil +} + +// collectAllocationRows decodes hami.io/*-devices-to-allocate annotations on every pod +// into a flat, sorted list of allocation rows. Pods with no HAMi annotations are skipped +// silently; pods with malformed HAMi annotations are skipped with a warning so that one +// bad pod cannot hide the rest of the cluster's allocation state. +func collectAllocationRows(pods []corev1.Pod) []allocationRow { + var rows []allocationRow + for _, pod := range pods { + node := pod.Annotations[util.AssignedNodeAnnotations] + if node == "" { + continue + } + + checklist := map[string]string{} + for key := range pod.Annotations { + if podDeviceAllocateAnnotation.MatchString(key) { + checklist[key] = key + } + } + if len(checklist) == 0 { + continue + } + + podDevices, err := device.DecodePodDevices(checklist, pod.Annotations) + if err != nil { + fmt.Fprintf(os.Stderr, "warning: skipping pod %s/%s: %v\n", pod.Namespace, pod.Name, err) + continue + } + + containerNames := podContainerNames(&pod) + for _, containers := range podDevices { + for ctrIdx, ctrDevices := range containers { + ctrName := fmt.Sprintf("container[%d]", ctrIdx) + if ctrIdx < len(containerNames) { + ctrName = containerNames[ctrIdx] + } + for _, d := range ctrDevices { + rows = append(rows, allocationRow{ + Node: node, + Namespace: pod.Namespace, + Pod: pod.Name, + Container: ctrName, + DeviceType: d.Type, + DeviceUUID: d.UUID, + RequestedMem: d.Usedmem, + RequestedCore: d.Usedcores, + }) + } + } + } + } + + sort.Slice(rows, func(i, j int) bool { + a, b := rows[i], rows[j] + if a.Node != b.Node { + return a.Node < b.Node + } + if a.Namespace != b.Namespace { + return a.Namespace < b.Namespace + } + if a.Pod != b.Pod { + return a.Pod < b.Pod + } + if a.Container != b.Container { + return a.Container < b.Container + } + return a.DeviceUUID < b.DeviceUUID + }) + return rows +} + +// podContainerNames returns container names in the same order used to build the +// hami.io/*-devices-to-allocate annotation: init containers first, then regular +// containers (see device.Resourcereqs for the matching encode-side order). +func podContainerNames(pod *corev1.Pod) []string { + names := make([]string, 0, len(pod.Spec.InitContainers)+len(pod.Spec.Containers)) + for _, c := range pod.Spec.InitContainers { + names = append(names, c.Name) + } + for _, c := range pod.Spec.Containers { + names = append(names, c.Name) + } + return names +} + +func printAllocationTable(out io.Writer, rows []allocationRow) { + w := tabwriter.NewWriter(out, 0, 4, 2, ' ', 0) + fmt.Fprintln(w, "NODE\tNAMESPACE\tPOD\tCONTAINER\tDEVICE TYPE\tDEVICE UUID\tMEMORY\tCORE") + for _, r := range rows { + fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\t%s\t%d\t%d\n", + r.Node, r.Namespace, r.Pod, r.Container, r.DeviceType, r.DeviceUUID, r.RequestedMem, r.RequestedCore) + } + w.Flush() +} diff --git a/cmd/hami-cli/allocations_test.go b/cmd/hami-cli/allocations_test.go new file mode 100644 index 0000000000..8e5599495c --- /dev/null +++ b/cmd/hami-cli/allocations_test.go @@ -0,0 +1,181 @@ +/* +Copyright 2024 The HAMi Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package main + +import ( + "bytes" + "context" + "strings" + "testing" + + corev1 "k8s.io/api/core/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/client-go/kubernetes/fake" +) + +func nvidiaPod(namespace, name, node string) corev1.Pod { + return corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Namespace: namespace, + Name: name, + Annotations: map[string]string{ + "hami.io/vgpu-node": node, + "hami.io/vgpu-devices-to-allocate": "GPU-0fc3eda5-e98b-a25b-5b0d-cf5c855d1448,NVIDIA,3000,0:;", + }, + }, + Spec: corev1.PodSpec{ + Containers: []corev1.Container{{Name: "trainer"}}, + }, + } +} + +func cambriconPod(namespace, name, node string) corev1.Pod { + return corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Namespace: namespace, + Name: name, + Annotations: map[string]string{ + "hami.io/vgpu-node": node, + "hami.io/cambricon-mlu-devices-to-allocate": "MLU-45013011-2257-0000-0000-000000000000,MLU,23308,0:;", + }, + }, + Spec: corev1.PodSpec{ + Containers: []corev1.Container{{Name: "infer"}}, + }, + } +} + +func TestCollectAllocationRows_MultiVendor(t *testing.T) { + pods := []corev1.Pod{ + nvidiaPod("default", "nvidia-job", "node-a"), + cambriconPod("default", "mlu-job", "node-b"), + } + + rows := collectAllocationRows(pods) + if len(rows) != 2 { + t.Fatalf("expected 2 rows, got %d: %+v", len(rows), rows) + } + + // rows are sorted by node, so node-a (nvidia) comes before node-b (cambricon). + nvidiaRow := rows[0] + if nvidiaRow.Node != "node-a" || nvidiaRow.Pod != "nvidia-job" || nvidiaRow.Container != "trainer" { + t.Errorf("unexpected nvidia row: %+v", nvidiaRow) + } + if nvidiaRow.DeviceUUID != "GPU-0fc3eda5-e98b-a25b-5b0d-cf5c855d1448" || nvidiaRow.DeviceType != "NVIDIA" { + t.Errorf("unexpected nvidia device fields: %+v", nvidiaRow) + } + if nvidiaRow.RequestedMem != 3000 || nvidiaRow.RequestedCore != 0 { + t.Errorf("unexpected nvidia usage fields: %+v", nvidiaRow) + } + + mluRow := rows[1] + if mluRow.Node != "node-b" || mluRow.Pod != "mlu-job" || mluRow.Container != "infer" { + t.Errorf("unexpected cambricon row: %+v", mluRow) + } + if mluRow.DeviceUUID != "MLU-45013011-2257-0000-0000-000000000000" || mluRow.DeviceType != "MLU" { + t.Errorf("unexpected cambricon device fields: %+v", mluRow) + } + if mluRow.RequestedMem != 23308 { + t.Errorf("unexpected cambricon memory: %+v", mluRow) + } +} + +func TestCollectAllocationRows_SkipsPodsWithoutHamiAnnotations(t *testing.T) { + pods := []corev1.Pod{ + {ObjectMeta: metav1.ObjectMeta{Namespace: "default", Name: "plain-pod"}}, + } + + rows := collectAllocationRows(pods) + if len(rows) != 0 { + t.Fatalf("expected 0 rows for a pod with no HAMi annotations, got %d", len(rows)) + } +} + +func TestCollectAllocationRows_SkipsMalformedAnnotationWithoutFailingOthers(t *testing.T) { + malformed := corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Namespace: "default", + Name: "broken-job", + Annotations: map[string]string{ + "hami.io/vgpu-node": "node-a", + "hami.io/vgpu-devices-to-allocate": "not-a-valid-device-record", + }, + }, + } + good := nvidiaPod("default", "good-job", "node-a") + + rows := collectAllocationRows([]corev1.Pod{malformed, good}) + if len(rows) != 1 { + t.Fatalf("expected the malformed pod to be skipped and the good pod kept, got %d rows: %+v", len(rows), rows) + } + if rows[0].Pod != "good-job" { + t.Errorf("expected surviving row to belong to good-job, got %q", rows[0].Pod) + } +} + +func TestCollectAllocationRows_InitContainerNameMapping(t *testing.T) { + pod := corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Namespace: "default", + Name: "multi-ctr", + Annotations: map[string]string{ + "hami.io/vgpu-node": "node-a", + "hami.io/vgpu-devices-to-allocate": "GPU-init,NVIDIA,1000,0:;" + + "GPU-main,NVIDIA,2000,0:;", + }, + }, + Spec: corev1.PodSpec{ + InitContainers: []corev1.Container{{Name: "setup"}}, + Containers: []corev1.Container{{Name: "trainer"}}, + }, + } + + rows := collectAllocationRows([]corev1.Pod{pod}) + if len(rows) != 2 { + t.Fatalf("expected 2 rows, got %d: %+v", len(rows), rows) + } + if rows[0].Container != "setup" { + t.Errorf("expected first row's container to resolve to init container %q, got %q", "setup", rows[0].Container) + } + if rows[1].Container != "trainer" { + t.Errorf("expected second row's container to resolve to %q, got %q", "trainer", rows[1].Container) + } +} + +func TestRunAllocations_NamespaceAndNodeFilters(t *testing.T) { + client := fake.NewSimpleClientset( + new(nvidiaPod("team-a", "job-1", "node-a")), + new(cambriconPod("team-b", "job-2", "node-b")), + ) + + namespaceFilter = "team-a" + nodeFilter = "" + defer func() { namespaceFilter = ""; nodeFilter = "" }() + + var out bytes.Buffer + if err := runAllocations(context.Background(), client, &out); err != nil { + t.Fatalf("runAllocations returned error: %v", err) + } + + got := out.String() + if !strings.Contains(got, "job-1") { + t.Errorf("expected output to contain job-1, got:\n%s", got) + } + if strings.Contains(got, "job-2") { + t.Errorf("expected namespace filter to exclude job-2, got:\n%s", got) + } +} diff --git a/cmd/hami-cli/main.go b/cmd/hami-cli/main.go new file mode 100644 index 0000000000..2141c1011d --- /dev/null +++ b/cmd/hami-cli/main.go @@ -0,0 +1,48 @@ +/* +Copyright 2024 The HAMi Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package main + +import ( + "github.com/spf13/cobra" + klog "k8s.io/klog/v2" + + "github.com/Project-HAMi/HAMi/pkg/util" + "github.com/Project-HAMi/HAMi/pkg/version" +) + +var rootCmd = &cobra.Command{ + Use: "hami-cli", + Short: "Read-only inspection tool for HAMi vGPU device allocations", +} + +var getCmd = &cobra.Command{ + Use: "get", + Short: "Display one or many HAMi resources", +} + +func init() { + rootCmd.AddCommand(getCmd) + getCmd.AddCommand(allocationsCmd) + rootCmd.AddCommand(version.VersionCmd) + rootCmd.PersistentFlags().AddGoFlagSet(util.InitKlogFlags()) +} + +func main() { + if err := rootCmd.Execute(); err != nil { + klog.Fatal(err) + } +} diff --git a/version.mk b/version.mk index 00a7ead79f..6b213a6e07 100644 --- a/version.mk +++ b/version.mk @@ -1,6 +1,6 @@ GO=go GO111MODULE=on -CMDS=scheduler vGPUmonitor +CMDS=scheduler vGPUmonitor hami-cli DEVICES=nvidia OUTPUT_DIR=bin TARGET_PLATFORMS=linux/amd64