[Feat] ✨ Add Mock GPU for OnlyCPU Node #7

Author: Harry-zklcdc

.github/workflows/build.yml

@@ -27,15 +27,17 @@ jobs:
           COMMIT=$(git rev-parse --short HEAD)
           GO_VERSION=$(go env GOVERSION)
           BUILD_TIME=$(date +%FT%T%z)
-          GOOS=linux GOARCH=amd64 go build -ldflags="-s -w -X main.BRANCH=${BRANCH} -X main.VERSION=${VERSION} -X main.COMMIT=${COMMIT} -X main.GoVersion=${GO_VERSION} -X main.BuildTime=${BUILD_TIME}" -trimpath -o gpu-docker-api cmd/gpu-docker-api/main.go && tar -zcvf gpu-docker-api-linux-amd64.tar.gz gpu-docker-api
+          GOOS=linux GOARCH=amd64 go build -ldflags="-s -w -X main.BRANCH=${BRANCH} -X main.VERSION=${VERSION} -X main.COMMIT=${COMMIT} -X main.GoVersion=${GO_VERSION} -X main.BuildTime=${BUILD_TIME}" -tag "nvidia" -trimpath -o gpu-docker-api cmd/gpu-docker-api/main.go && tar -zcvf gpu-docker-api-nvidia-linux-amd64.tar.gz gpu-docker-api
+          GOOS=linux GOARCH=amd64 go build -ldflags="-s -w -X main.BRANCH=${BRANCH} -X main.VERSION=${VERSION} -X main.COMMIT=${COMMIT} -X main.GoVersion=${GO_VERSION} -X main.BuildTime=${BUILD_TIME}" -tag "mock" -trimpath -o gpu-docker-api cmd/gpu-docker-api/main.go && tar -zcvf gpu-docker-api-mock-linux-amd64.tar.gz gpu-docker-api
       - name: Create Release and Upload Release Asset
         uses: softprops/action-gh-release@v1
         with:
           files: |
-            gpu-docker-api-linux-amd64.tar.gz
+            gpu-docker-api-nvidia-linux-amd64.tar.gz
+            gpu-docker-api-mock-linux-amd64.tar.gz
 
-  docker-build:
-    name: docker-build
+  docker-build-nvidia:
+    name: docker-build-nvidia
     runs-on: ubuntu-latest
     permissions:
       packages: write
@@ -69,4 +71,45 @@ jobs:
           push: true
           tags: |
             docker.io/xshengtech/gpu-docker-api:${{ github.ref_name }}
+            docker.io/xshengtech/gpu-docker-api:${{ github.ref_name }}-nvidia
             docker.io/xshengtech/gpu-docker-api:latest
+            docker.io/xshengtech/gpu-docker-api:latest-nvidia
+            docker.io/xshengtech/gpu-docker-api:nvidia
+
+  docker-build-mock:
+    name: docker-build-mock
+    runs-on: ubuntu-latest
+    permissions:
+      packages: write
+      contents: read
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Setup timezone
+        uses: zcong1993/setup-timezone@master
+        with:
+          timezone: Asia/Shanghai
+
+      - name: Login to DockerHub
+        uses: docker/login-action@v2
+        with:
+          username: ${{ secrets.DOCKERHUB_USERNAME }}
+          password: ${{ secrets.DOCKERHUB_TOKEN }}
+
+      - name: Set up QEMU
+        uses: docker/setup-qemu-action@v2
+      
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v2
+
+      - name: Build and push
+        uses: docker/build-push-action@v4
+        with:
+          platforms: linux/amd64
+          context: .
+          file: ./Dockerfile  
+          push: true
+          tags: |
+            docker.io/xshengtech/gpu-docker-api:${{ github.ref_name }}-mock
+            docker.io/xshengtech/gpu-docker-api:latest-mock
+            docker.io/xshengtech/gpu-docker-api:mock

Dockerfile

@@ -15,14 +15,14 @@ COPY go.sum .
 RUN go mod download
 
 COPY . .
-RUN make linux
+RUN make nvidia_linux
 
 FROM ubuntu:22.04
 
 VOLUME /data
 WORKDIR /data
 
-COPY --from=builder /build/bin/gpu-docker-api-linux-amd64 /data/gpu-docker-api
+COPY --from=builder /build/bin/gpu-docker-api-nvidia-linux-amd64 /data/gpu-docker-api
 
 EXPOSE 2378
 

Dockerfile.mock

@@ -0,0 +1,29 @@
+FROM golang:1.23.6-alpine AS builder
+LABEL stage=gobuilder \
+      mainatiner=https://github.com/XShengTech/gpu-docker-api
+
+# RUN sed -i 's/dl-cdn.alpinelinux.org/mirrors.tuna.tsinghua.edu.cn/g' /etc/apk/repositories
+RUN apk add gcc g++ make libffi-dev openssl-dev libtool git
+
+ENV CGO_ENABLED=0
+# ENV GOPROXY=https://goproxy.cn,direct
+
+WORKDIR /build
+
+COPY go.mod .
+COPY go.sum .
+RUN go mod download
+
+COPY . .
+RUN make mock_linux
+
+FROM ubuntu:22.04
+
+VOLUME /data
+WORKDIR /data
+
+COPY --from=builder /build/bin/gpu-docker-api-mock-linux-amd64 /data/gpu-docker-api
+
+EXPOSE 2378
+
+ENTRYPOINT ["./gpu-docker-api"]

Makefile

@@ -13,30 +13,38 @@ CURRENT_DIR = $(shell pwd)
 BUILD_DIR = ${CURRENT_DIR}/cmd/${BINARY}
 BIN_DIR= ${CURRENT_DIR}/bin
 
+NVIDIA = nvidia
+MOCK = mock
+
 LDFLAGS = -ldflags "-s -w -X main.BRANCH=${BRANCH} -X main.VERSION=${VERSION} -X main.COMMIT=${COMMIT} -X main.GoVersion=${GO_VERSION} -X main.BuildTime=${BUILD_TIME}"
 
 all: fmt imports clean linux darwin windows
 
 build: clean linux darwin windows
 
-linux:
+nvidia_linux:
+	cd ${BUILD_DIR}; \
+	GOOS=linux GOARCH=${GOARCH} go build ${LDFLAGS} -tags "${NVIDIA}" -o ${BIN_DIR}/${BINARY}-${NVIDIA}-linux-${GOARCH} . ; \
+	cd - >/dev/null
+
+nvidia_linux_no_ldflags:
 	cd ${BUILD_DIR}; \
-	GOOS=linux GOARCH=${GOARCH} go build ${LDFLAGS} -o ${BIN_DIR}/${BINARY}-linux-${GOARCH} . ; \
+	GOOS=linux GOARCH=${GOARCH} go build -tags "${NVIDIA}" -o ${BIN_DIR}/${BINARY}-${NVIDIA}-linux-${GOARCH} . ; \
 	cd - >/dev/null
 
-linux_no_ldflags:
+nvidia_darwin:
 	cd ${BUILD_DIR}; \
-	GOOS=linux GOARCH=${GOARCH} go build -o ${BIN_DIR}/${BINARY}-linux-${GOARCH} . ; \
+	GOOS=darwin GOARCH=${GOARCH} go build ${LDFLAGS} -tags "${NVIDIA}" -o ${BIN_DIR}/${BINARY}-${NVIDIA}-darwin-${GOARCH} . ; \
 	cd - >/dev/null
 
-darwin:
+nvidia_windows:
 	cd ${BUILD_DIR}; \
-	GOOS=darwin GOARCH=${GOARCH} go build ${LDFLAGS} -o ${BIN_DIR}/${BINARY}-darwin-${GOARCH} . ; \
+	GOOS=windows GOARCH=${GOARCH} go build ${LDFLAGS} -tags "${NVIDIA}" -o ${BIN_DIR}/${BINARY}-${NVIDIA}-windows-${GOARCH}.exe . ; \
 	cd - >/dev/null
 
-windows:
+mock_linux:
 	cd ${BUILD_DIR}; \
-	GOOS=windows GOARCH=${GOARCH} go build ${LDFLAGS} -o ${BIN_DIR}/${BINARY}-windows-${GOARCH}.exe . ; \
+	GOOS=linux GOARCH=${GOARCH} go build ${LDFLAGS} -tags "${MOCK}" -o ${BIN_DIR}/${BINARY}-${MOCK}-linux-${GOARCH} . ; \
 	cd - >/dev/null
 
 docker_build:

internal/schedulers/gpuscheduler.go

@@ -1,3 +1,5 @@
+//go:build !mock
+
 package schedulers
 
 import (

internal/schedulers/gpuscheduler_mock.go

@@ -0,0 +1,208 @@
+//go:build mock
+
+package schedulers
+
+import (
+	"encoding/json"
+	"strconv"
+	"strings"
+	"sync"
+
+	"github.com/pkg/errors"
+
+	"github.com/mayooot/gpu-docker-api/internal/etcd"
+	"github.com/mayooot/gpu-docker-api/internal/workQueue"
+	"github.com/mayooot/gpu-docker-api/internal/xerrors"
+)
+
+const (
+	gpuStatusMapKey = "gpuStatusMapKey"
+)
+
+var GpuScheduler *gpuScheduler
+
+type gpu struct {
+	Index int     `json:"index"`
+	UUID  *string `json:"uuid"`
+}
+
+type gpuScheduler struct {
+	sync.RWMutex
+
+	AvailableGpuNums int             `json:"availableGpuNums"`
+	GpuStatusMap     map[string]byte `json:"gpuStatusMap"`
+}
+
+func InitGPuScheduler() error {
+	var err error
+	GpuScheduler, err = initGpuFormEtcd()
+	if err != nil {
+		return errors.Wrap(err, "initFormEtcd failed")
+	}
+
+	if GpuScheduler.AvailableGpuNums == 0 || len(GpuScheduler.GpuStatusMap) == 0 {
+		// if it has not been initialized
+		gpus, err := getAllGpuUUID()
+		if err != nil {
+			return errors.Wrap(err, "getAllGpuUUID failed")
+		}
+
+		GpuScheduler.AvailableGpuNums = len(gpus)
+		for i := 0; i < len(gpus); i++ {
+			GpuScheduler.GpuStatusMap[*gpus[i].UUID] = 0
+		}
+	}
+	return nil
+}
+
+func CloseGpuScheduler() error {
+	return etcd.Put(etcd.Gpus, gpuStatusMapKey, GpuScheduler.serialize())
+}
+
+func initGpuFormEtcd() (s *gpuScheduler, err error) {
+	bytes, err := etcd.GetValue(etcd.Gpus, gpuStatusMapKey)
+	if err != nil {
+		if xerrors.IsNotExistInEtcdError(err) {
+			err = nil
+		} else {
+			return s, err
+		}
+	}
+
+	s = &gpuScheduler{
+		GpuStatusMap: make(map[string]byte),
+	}
+	if len(bytes) != 0 {
+		err = json.Unmarshal(bytes, &s)
+	}
+	return s, err
+}
+
+// Apply for a specified number of gpus
+func (gs *gpuScheduler) Apply(num int) ([]string, error) {
+	if num <= 0 || num > gs.AvailableGpuNums {
+		return nil, errors.New("num must be greater than 0 and less than " + strconv.Itoa(gs.AvailableGpuNums))
+	}
+
+	gs.Lock()
+	defer gs.Unlock()
+
+	var availableGpus []string
+	for k, v := range gs.GpuStatusMap {
+		if v == 0 {
+			gs.GpuStatusMap[k] = 1
+			availableGpus = append(availableGpus, k)
+			if len(availableGpus) == num {
+				break
+			}
+		}
+	}
+
+	if len(availableGpus) < num {
+		gs.restore(availableGpus)
+		return nil, xerrors.NewGpuNotEnoughError()
+	}
+
+	go gs.putToEtcd()
+
+	return availableGpus, nil
+}
+
+// Restore a specified number of gpu
+func (gs *gpuScheduler) Restore(gpus []string) {
+	if len(gpus) <= 0 || len(gpus) > gs.AvailableGpuNums {
+		return
+	}
+
+	gs.Lock()
+	defer gs.Unlock()
+
+	gs.restore(gpus)
+
+	go gs.putToEtcd()
+}
+
+func (gs *gpuScheduler) restore(gpus []string) {
+	if len(gpus) <= 0 || len(gpus) > gs.AvailableGpuNums {
+		return
+	}
+
+	for _, gpu := range gpus {
+		gs.GpuStatusMap[gpu] = 0
+	}
+}
+
+func (gs *gpuScheduler) serialize() *string {
+	gs.RLock()
+	defer gs.RUnlock()
+
+	bytes, _ := json.Marshal(gs)
+	tmp := string(bytes)
+	return &tmp
+}
+
+func (gs *gpuScheduler) GetGpuStatus() map[string]byte {
+	gs.RLock()
+	defer gs.RUnlock()
+
+	copyMap := make(map[string]byte, len(gs.GpuStatusMap))
+	for k, v := range gs.GpuStatusMap {
+		copyMap[k] = v
+	}
+
+	return copyMap
+}
+
+func (gs *gpuScheduler) putToEtcd() {
+	workQueue.Queue <- etcd.PutKeyValue{
+		Resource: etcd.Gpus,
+		Key:      gpuStatusMapKey,
+		Value:    GpuScheduler.serialize(),
+	}
+}
+
+func getAllGpuUUID() ([]*gpu, error) {
+	uuids := []string{
+		"MockGPU-0",
+		"MockGPU-1",
+		"MockGPU-2",
+		"MockGPU-3",
+		"MockGPU-4",
+		"MockGPU-5",
+		"MockGPU-6",
+		"MockGPU-7",
+	}
+	gpuList := []*gpu{}
+	for i, uuid := range uuids {
+		gpuList = append(gpuList, &gpu{
+			Index: i,
+			UUID:  &uuid,
+		})
+	}
+
+	return gpuList, nil
+}
+
+func parseOutput(output string) (gpuList []*gpu, err error) {
+	lines := strings.Split(output, "\n")
+	gpuList = make([]*gpu, 0, len(lines))
+	for _, line := range lines {
+		if line == "" {
+			continue
+		}
+
+		fields := strings.Split(line, ", ")
+		if len(fields) == 2 {
+			index, err := strconv.Atoi(fields[0])
+			if err != nil {
+				return gpuList, errors.Errorf("invaild index: %s, ", fields[0])
+			}
+			uuid := "nvidia.com/gpu=" + fields[1]
+			gpuList = append(gpuList, &gpu{
+				Index: index,
+				UUID:  &uuid,
+			})
+		}
+	}
+	return
+}