Add shared GPUs for deployments using DRA and MPS/timeslicing

Dockerfile

@@ -1,47 +1,56 @@
-############################
-# STEP 1 build executable binary
-############################
-FROM --platform=$BUILDPLATFORM golang:alpine AS builder
+FROM --platform=$BUILDPLATFORM golang:alpine AS build-0
+
 # Install git.
-RUN apk update && apk add --no-cache git=~2
+RUN apk update && \
+    apk add --no-cache git=~2
 
 # Set up working directory
 WORKDIR /app
+
 # Copy go.mod and go.sum separately so we only invalidate the downloading layers if we need to
 COPY go.mod go.sum ./
 
 # Fetch dependencies and build the binary
 ENV GO111MODULE=on
-RUN go mod download
+ENV CGO_ENABLED=0
+
+RUN go mod download -x
 
 # Copy the rest of the project to ensure code changes doesnt trigger re-download of all deps
 COPY . .
 
-RUN CGO_ENABLED=0 GOOS=$TARGETOS GOARCH=$TARGETARCH go build -a -installsuffix cgo -o main .
+FROM --platform=$BUILDPLATFORM build-0 AS build-1
 
+ARG TARGETOS
+ARG TARGETARCH
 
-############################
-# STEP 2 build a small image
-############################
+RUN CGO_ENABLED=0 GOOS=${TARGETOS} GOARCH=${TARGETARCH} \
+  go build -a -installsuffix cgo -o go-deploy_${TARGETOS}_${TARGETARCH} .
+
+
+# Runner 
 FROM alpine:3
 
 # Set up the working directory
 WORKDIR /go
 
+ARG TARGETOS
+ARG TARGETARCH
+
 # Copy the binary from the builder stage
-COPY --from=builder /app/main .
+COPY --from=build-1 --chmod=777 /app/go-deploy_${TARGETOS}_${TARGETARCH} /usr/bin/go-deploy
 
 # Copy the "index" folder
-COPY --from=builder /app/index index
+COPY --from=build-1 /app/index index
 
 # Copy the "docs" folder
-COPY --from=builder /app/docs docs
+COPY --from=build-1 /app/docs docs
 
 # Set environment variables and expose necessary port
 ENV PORT=8080
 ENV GIN_MODE=release
 EXPOSE 8080
 
 # Run the Go Gin binary
-ENTRYPOINT ["./main"]
+ENTRYPOINT ["/usr/bin/go-deploy"]
 

Makefile

@@ -13,7 +13,7 @@ all: build
 build:
 	@echo "Building the application..."
 	@mkdir -p $(BUILD_DIR)
-	@CGO_ENABLED=0 go build -o $(BUILD_DIR)/$(BINARY_NAME)$(EXT) .
+	@CGO_ENABLED=1 go build -ldflags="-X github.com/NVIDIA/k8s-dra-driver-gpu/internal/info.version=v1.34.2" -o $(BUILD_DIR)/$(BINARY_NAME)$(EXT) .
 	@echo "Build complete."
 
 run: build

cmd/server.go

@@ -5,6 +5,7 @@ import (
 	"errors"
 	argFlag "flag"
 	"fmt"
+	"net"
 	"net/http"
 	"os"
 	"time"
@@ -19,9 +20,11 @@ import (
 	"github.com/kthcloud/go-deploy/pkg/log"
 	"github.com/kthcloud/go-deploy/pkg/metrics"
 	"github.com/kthcloud/go-deploy/routers"
+	"github.com/kthcloud/go-deploy/service/utils"
 )
 
 type Options struct {
+	Ctx   context.Context
 	Flags FlagDefinitionList
 	Mode  string
 }
@@ -74,8 +77,7 @@ func Create(opts *Options) *App {
 	}
 	log.Printf("%sInitialization completed%s", log.Orange, log.Reset)
 
-	ctx, cancel := context.WithCancel(context.Background())
-
+	ctx, cancel := context.WithCancel(utils.FirstNonZero(opts.Ctx, context.Background()))
 	for _, flag := range opts.Flags {
 		// Handle api worker separately
 		if flag.Name == "api" {
@@ -100,6 +102,9 @@ func Create(opts *Options) *App {
 		httpServer = &http.Server{
 			Addr:    fmt.Sprintf("0.0.0.0:%d", config.Config.Port),
 			Handler: routers.NewRouter(),
+			BaseContext: func(_ net.Listener) context.Context {
+				return ctx
+			},
 		}
 
 		go func() {
@@ -126,11 +131,14 @@ func (app *App) Stop() {
 		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 		defer cancel()
 		if err := app.httpServer.Shutdown(ctx); err != nil {
-			log.Fatalln(fmt.Errorf("failed to shutdown server. details: %w", err))
+			log.Errorln(fmt.Errorf("failed to gracefully shutdown server. details: %w", err))
+			if closeErr := app.httpServer.Close(); closeErr != nil {
+				log.Fatalln("Force close failed:", closeErr)
+			}
 		}
 
 		<-ctx.Done()
-		log.Println("Saiting for http server to shutdown...")
+		log.Println("Waiting for http server to shutdown...")
 	}
 
 	shutdown()