package main
import (
"encoding/json"
"log"
"net/http"
"github.com/julienschmidt/httprouter"
extenderv1 "k8s.io/kube-scheduler/extender/v1"
)
func main() {
router := httprouter.New()
router.POST("/filter", filterHandler)
router.POST("/prioritize", prioritizeHandler)
log.Fatal(http.ListenAndServe(":8888", router))
}
// 필터 핸들러
func filterHandler(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
var extenderArgs extenderv1.ExtenderArgs
var extenderFilterResult extenderv1.ExtenderFilterResult
if err := json.NewDecoder(r.Body).Decode(&extenderArgs); err != nil {
http.Error(w, err.Error(), http.StatusBadRequest)
return
}
// 모든 노드를 허용하는 간단한 예제
extenderFilterResult.Nodes = extenderArgs.Nodes
extenderFilterResult.FailedNodes = make(map[string]string)
// 특정 조건에 따라 노드 필터링
// 예: GPU 요구 사항이 있는 포드에 대해 GPU가 있는 노드만 허용
if requiresGPU(&extenderArgs.Pod) {
filteredNodes := &extenderv1.NodeList{
Items: make([]extenderv1.Node, 0),
}
for _, node := range extenderArgs.Nodes.Items {
if hasGPU(&node) {
filteredNodes.Items = append(filteredNodes.Items, node)
} else {
extenderFilterResult.FailedNodes[node.Name] = "Node does not have GPU"
}
}
extenderFilterResult.Nodes = filteredNodes
}
if err := json.NewEncoder(w).Encode(extenderFilterResult); err != nil {
http.Error(w, err.Error(), http.StatusInternalServerError)
return
}
}
// 우선순위 핸들러
func prioritizeHandler(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
var extenderArgs extenderv1.ExtenderArgs
var hostPriorityList extenderv1.HostPriorityList
if err := json.NewDecoder(r.Body).Decode(&extenderArgs); err != nil {
http.Error(w, err.Error(), http.StatusBadRequest)
return
}
// 각 노드에 점수 할당
hostPriorityList = make(extenderv1.HostPriorityList, len(extenderArgs.Nodes.Items))
for i, node := range extenderArgs.Nodes.Items {
// 간단한 예제: 모든 노드에 동일한 점수 할당
hostPriorityList[i] = extenderv1.HostPriority{
Host: node.Name,
Score: 1,
}
// 특정 조건에 따라 점수 조정
// 예: GPU 메모리가 많은 노드에 더 높은 점수 할당
if requiresGPU(&extenderArgs.Pod) && hasGPU(&node) {
gpuMemory := getGPUMemory(&node)
hostPriorityList[i].Score = int64(gpuMemory / 1024) // GB 단위로 변환
}
}
if err := json.NewEncoder(w).Encode(hostPriorityList); err != nil {
http.Error(w, err.Error(), http.StatusInternalServerError)
return
}
}
// GPU 요구 사항 확인 함수
func requiresGPU(pod *extenderv1.Pod) bool {
// 포드의 리소스 요청에서 GPU 요구 사항 확인
for _, container := range pod.Spec.Containers {
if _, ok := container.Resources.Requests["nvidia.com/gpu"]; ok {
return true
}
}
return false
}
// 노드에 GPU가 있는지 확인하는 함수
func hasGPU(node *extenderv1.Node) bool {
// 노드의 용량에서 GPU 확인
if _, ok := node.Status.Capacity["nvidia.com/gpu"]; ok {
return true
}
return false
}
// 노드의 GPU 메모리 확인 함수
func getGPUMemory(node *extenderv1.Node) int {
// 노드 레이블에서 GPU 메모리 확인
if memoryStr, ok := node.Labels["gpu.nvidia.com/memory"]; ok {
var memory int
fmt.Sscanf(memoryStr, "%d", &memory)
return memory
}
return 0
}
FROM golang:1.17 as builder
WORKDIR /go/src/example.com/gpu-scheduler
COPY . .
RUN CGO_ENABLED=0 GOOS=linux go build -a -installsuffix cgo -o kube-scheduler .
FROM alpine:3.14
RUN apk --no-cache add ca-certificates
WORKDIR /
COPY --from=builder /go/src/example.com/gpu-scheduler/kube-scheduler .
ENTRYPOINT ["/kube-scheduler"]
