Support newly added option for GPUs in Dirigent.

Signed-off-by: Luka Simić <git@kocka.tech>

Author: KockaAdmiralac

pkg/config/parser.go

@@ -60,6 +60,7 @@ type LoaderConfiguration struct {
 	RpsImage                    string  `json:"RpsImage"`
 	RpsRuntimeMs                int     `json:"RpsRuntimeMs"`
 	RpsMemoryMB                 int     `json:"RpsMemoryMB"`
+	RpsRequestedGpu             int     `json:"RpsRequestedGpu"`
 	RpsIterationMultiplier      int     `json:"RpsIterationMultiplier"`
 	RpsDataSizeMB               float64 `json:"RpsDataSizeMB"`
 	RpsFile                     string  `json:"RpsFile"`

pkg/driver/deployment/dirigent.go

@@ -2,9 +2,6 @@ package deployment
 
 import (
 	"fmt"
-	log "github.com/sirupsen/logrus"
-	"github.com/vhive-serverless/loader/pkg/common"
-	"github.com/vhive-serverless/loader/pkg/config"
 	"io"
 	"math/rand"
 	"net"
@@ -14,6 +11,10 @@ import (
 	"strings"
 	"sync"
 	"time"
+
+	log "github.com/sirupsen/logrus"
+	"github.com/vhive-serverless/loader/pkg/common"
+	"github.com/vhive-serverless/loader/pkg/config"
 )
 
 type dirigentDeployer struct{}
@@ -47,6 +48,7 @@ func (*dirigentDeployer) Deploy(cfg *config.Configuration) {
 				dirigentConfig.RegistrationServer,
 				cfg.LoaderConfiguration.BusyLoopOnSandboxStartup,
 				cfg.LoaderConfiguration.PrepullMode,
+				cfg.LoaderConfiguration.RpsRequestedGpu,
 			)
 		}(i)
 	}
@@ -80,7 +82,7 @@ var checkClient = &http.Client{
 	},
 }
 
-func deployDirigent(function *common.Function, controlPlaneAddress string, busyLoopOnColdStart bool, prepullMode string) {
+func deployDirigent(function *common.Function, controlPlaneAddress string, busyLoopOnColdStart bool, prepullMode string, requestedGpu int) {
 	metadata := function.DirigentMetadata
 
 	if metadata == nil {
@@ -95,6 +97,7 @@ func deployDirigent(function *common.Function, controlPlaneAddress string, busyL
 		"scaling_lower_bound": {strconv.Itoa(metadata.ScalingLowerBound)},
 		"requested_cpu":       {strconv.Itoa(function.CPURequestsMilli)},
 		"requested_memory":    {strconv.Itoa(function.MemoryRequestsMiB)},
+		"requested_gpu":       {strconv.Itoa(requestedGpu)},
 		"env_vars":            metadata.EnvVars,     // FORMAT: arg1=value1 arg2=value2 ...
 		"program_args":        metadata.ProgramArgs, // FORMAT: arg1 arg2 ...
 		"prepull_mode":        {prepullMode},