Feat/scheduler pdb

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
 

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

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

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/<slug>-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.

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/<slug>-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()
+}