[Feat/Item-40057463] Support Country Code, Logging, Networking, etc. …

…fields (#12)

go.mod

@@ -3,7 +3,7 @@ module github.com/SaladTechnologies/virtual-kubelet-saladcloud
 go 1.20
 
 require (
-	github.com/lucklypriyansh-2/salad-client v0.0.0-20230902193233-7c5a02b4c6d7
+	github.com/lucklypriyansh-2/salad-client v0.0.0-20231005114315-47f7f0403e30
 	github.com/prometheus/client_model v0.4.0
 	github.com/sirupsen/logrus v1.9.3
 	github.com/virtual-kubelet/virtual-kubelet v1.10.0
@@ -37,7 +37,7 @@ require (
 	github.com/google/gnostic-models v0.6.8 // indirect
 	github.com/google/go-cmp v0.5.9 // indirect
 	github.com/google/gofuzz v1.2.0 // indirect
-	github.com/google/uuid v1.3.0 // indirect
+	github.com/google/uuid v1.3.0
 	github.com/gorilla/mux v1.8.0 // indirect
 	github.com/grpc-ecosystem/go-grpc-prometheus v1.2.0 // indirect
 	github.com/grpc-ecosystem/grpc-gateway/v2 v2.16.0 // indirect
@@ -52,7 +52,7 @@ require (
 	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
 	github.com/modern-go/reflect2 v1.0.2 // indirect
 	github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
-	github.com/pkg/errors v0.9.1 // indirect
+	github.com/pkg/errors v0.9.1
 	github.com/prometheus/client_golang v1.16.0 // indirect
 	github.com/prometheus/common v0.44.0 // indirect
 	github.com/prometheus/procfs v0.10.1 // indirect
@@ -97,7 +97,7 @@ require (
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 	k8s.io/apimachinery v0.28.2
 	k8s.io/apiserver v0.28.1 // indirect
-	k8s.io/client-go v0.28.1 // indirect
+	k8s.io/client-go v0.28.1
 	k8s.io/component-base v0.28.1 // indirect
 	k8s.io/klog/v2 v2.100.1 // indirect
 	k8s.io/kms v0.28.1 // indirect

go.sum

@@ -133,6 +133,10 @@ github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
 github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
 github.com/lucklypriyansh-2/salad-client v0.0.0-20230902193233-7c5a02b4c6d7 h1:N7D71xaSmphJuSZdtJnhq/B5sF2/qEkId00LhrqWusg=
 github.com/lucklypriyansh-2/salad-client v0.0.0-20230902193233-7c5a02b4c6d7/go.mod h1:VjkYRyCaHqCXvxGIDnUCRg8QEcvOuOkhNnbHsrAPETc=
+github.com/lucklypriyansh-2/salad-client v0.0.0-20231005112118-e2e44b2b9380 h1:LQo11CfzUSuxI3VNrm0/ZNxBnIpZ3slvfXWWun5Buvo=
+github.com/lucklypriyansh-2/salad-client v0.0.0-20231005112118-e2e44b2b9380/go.mod h1:aFjkAhEyFj9mGwJ+w4rE0tWMTjwE9OBNhl5/Q9dFNDE=
+github.com/lucklypriyansh-2/salad-client v0.0.0-20231005114315-47f7f0403e30 h1:HrlYiLV6NAoalMkbfDnEjvO3FNJ9qcPTG+DZ47kktiA=
+github.com/lucklypriyansh-2/salad-client v0.0.0-20231005114315-47f7f0403e30/go.mod h1:aFjkAhEyFj9mGwJ+w4rE0tWMTjwE9OBNhl5/Q9dFNDE=
 github.com/mailru/easyjson v0.7.7 h1:UGYAvKxe3sBsEDzO8ZeWOSlIQfWFlxbzLZe7hwFURr0=
 github.com/mailru/easyjson v0.7.7/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJJLY9Nlc=
 github.com/matttproud/golang_protobuf_extensions v1.0.4 h1:mmDVorXM7PCGKw94cs5zkfA9PSy5pEvNWRP0ET0TIVo=
@@ -184,6 +188,7 @@ github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/
 github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
 github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
 github.com/stretchr/testify v1.8.3 h1:RP3t2pwF7cMEbC1dqtB6poj3niw/9gnV4Cjg5oW5gtY=
+github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk=
 github.com/tmc/grpc-websocket-proxy v0.0.0-20220101234140-673ab2c3ae75 h1:6fotK7otjonDflCTK0BCfls4SPy3NcCVb5dqqmbRknE=
 github.com/virtual-kubelet/virtual-kubelet v1.10.0 h1:eV/mFFqThOJLz7Gjn1Ev8LchanGKGA2qZlsW6wipb4g=
 github.com/virtual-kubelet/virtual-kubelet v1.10.0/go.mod h1:7Pvdei1p82C9uWS1VzLrnXbHTwQcGBoqShahChpacgI=

internal/provider/provider.go

@@ -5,7 +5,9 @@ import (
 	"encoding/json"
 	"github.com/SaladTechnologies/virtual-kubelet-saladcloud/internal/models"
 	"github.com/SaladTechnologies/virtual-kubelet-saladcloud/internal/utils"
+	"github.com/google/uuid"
 	saladclient "github.com/lucklypriyansh-2/salad-client"
+	"github.com/pkg/errors"
 	dto "github.com/prometheus/client_model/go"
 	"github.com/virtual-kubelet/virtual-kubelet/log"
 	nodeapi "github.com/virtual-kubelet/virtual-kubelet/node/api"
@@ -15,6 +17,7 @@ import (
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/resource"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"strconv"
 	"strings"
 	"time"
 )
@@ -83,7 +86,6 @@ func (p *SaladCloudProvider) CreatePod(ctx context.Context, pod *corev1.Pod) err
 	ctx, span := trace.StartSpan(ctx, "CreatePod")
 	defer span.End()
 	log.G(ctx).Debug("creating a CreatePod", pod.Name)
-
 	createContainerObject := p.createContainersObject(pod)
 	createContainerGroup := p.createContainerGroup(createContainerObject, pod)
 
@@ -336,63 +338,65 @@ func (p *SaladCloudProvider) getContainerEnvironment(podMetadata metav1.ObjectMe
 }
 
 func (p *SaladCloudProvider) createContainersObject(pod *corev1.Pod) []saladclient.CreateContainer {
-
 	cpu, memory := utils.GetPodResource(pod.Spec)
-
 	creteContainersArray := make([]saladclient.CreateContainer, 0)
 	for _, container := range pod.Spec.Containers {
-
 		containerResourceRequirement := saladclient.NewContainerResourceRequirements(int32(cpu), int32(memory))
 		createContainer := saladclient.NewCreateContainer(container.Image, *containerResourceRequirement)
 
 		createContainer.SetEnvironmentVariables(p.getContainerEnvironment(pod.ObjectMeta, container))
 		if container.Command != nil {
 			createContainer.SetCommand(container.Command)
 		}
+		gpuClasses, err := p.getGPUClasses(pod)
+		if err == nil && gpuClasses != nil && len(gpuClasses) > 0 {
+			createContainer.Resources.SetGpuClasses(gpuClasses)
+		}
 		creteContainersArray = append(creteContainersArray, *createContainer)
-		// TODO Add support for container Registry auth
 	}
 	return creteContainersArray
-
 }
 
 func (p *SaladCloudProvider) createContainerGroup(createContainerList []saladclient.CreateContainer, pod *corev1.Pod) []saladclient.CreateContainerGroup {
-
 	createContainerGroups := make([]saladclient.CreateContainerGroup, 0)
-
-	if pod.ObjectMeta.GetAnnotations()["countryCodes"] == "" {
-		pod.ObjectMeta.SetAnnotations(map[string]string{
-			"countryCodes": "US",
-		})
-	}
-
-	var countryCodesEnum []saladclient.CountryCode
-	for _, countryCode := range strings.Split(pod.ObjectMeta.GetAnnotations()["countryCodes"], ",") {
-		countryCodeEnum := saladclient.CountryCode(countryCode)
-		countryCodesEnum = append(countryCodesEnum, countryCodeEnum)
-	}
-
 	for _, container := range createContainerList {
 		createContainerGroupRequest := *saladclient.NewCreateContainerGroup(utils.GetPodName(pod.Namespace, pod.Name, pod), container, "always", 1)
-		createContainerGroupRequest.SetCountryCodes(countryCodesEnum)
 		readinessProbe, err := p.getWorkloadContainerProbeFrom(pod.Spec.Containers[0].ReadinessProbe)
-		if err == nil {
+		if err == nil && readinessProbe != nil {
 			createContainerGroupRequest.ReadinessProbe = *readinessProbe
 		} else {
 			log.G(context.Background()).Errorf("Failed to get readinessProbe ", err)
 		}
 		livenessProbe, err := p.getWorkloadContainerProbeFrom(pod.Spec.Containers[0].LivenessProbe)
-		if err == nil {
+		if err == nil && livenessProbe != nil {
 			createContainerGroupRequest.LivenessProbe = *livenessProbe
 		} else {
 			log.G(context.Background()).Errorf("Failed to get livenessProbe ", err)
 		}
 		startupProbe, err := p.getWorkloadContainerProbeFrom(pod.Spec.Containers[0].StartupProbe)
-		if err == nil {
+		if err == nil && startupProbe != nil {
 			createContainerGroupRequest.StartupProbe = *startupProbe
 		} else {
 			log.G(context.Background()).Errorf("Failed to get startupProbe ", err)
 		}
+		countryCodes, err := p.getCountryCodes(pod)
+		if err != nil {
+			log.G(context.Background()).Errorf("Failed to get countryCodes ", err)
+		} else {
+			createContainerGroupRequest.SetCountryCodes(countryCodes)
+		}
+		networking, err := p.getNetworking(pod)
+		if err != nil {
+			log.G(context.Background()).Errorf("Failed to get networking ", err)
+		} else if networking != nil {
+			createContainerGroupRequest.SetNetworking(*networking)
+		}
+		restartPolicy, err := p.getRestartPolicy(pod)
+		if err != nil {
+			log.G(context.Background()).Errorf("Failed to get restartPolicy ", err)
+		} else {
+			createContainerGroupRequest.SetRestartPolicy(*restartPolicy)
+		}
 		createContainerGroups = append(createContainerGroups, createContainerGroupRequest)
 	}
 	return createContainerGroups
@@ -426,3 +430,92 @@ func (p *SaladCloudProvider) getWorkloadContainerProbeFrom(k8sProbe *corev1.Prob
 	}
 	return saladclient.NewNullableContainerGroupProbe(probe), nil
 }
+
+func (p *SaladCloudProvider) getGPUClasses(pod *corev1.Pod) ([]string, error) {
+	gpuRequestedString, ok := pod.ObjectMeta.Annotations["salad.com/gpu-classes"]
+	if !ok {
+		return nil, nil
+	}
+	gpuRequested := strings.Split(gpuRequestedString, ",")
+	saladClientGpuIds := make([]string, 0)
+	var gpuClasses *saladclient.GpuClassesList = nil
+
+	for _, gpu := range gpuRequested {
+		gpuCleaned := strings.TrimSpace(strings.ToLower(gpu))
+		_, uuidErr := uuid.Parse(gpuCleaned)
+		if uuidErr == nil {
+			saladClientGpuIds = append(saladClientGpuIds, gpuCleaned)
+		} else {
+			if gpuClasses == nil {
+				classes, _, err := p.apiClient.OrganizationDataAPI.ListGpuClasses(context.Background(), p.inputVars.OrganizationName).Execute()
+				if err != nil {
+					log.G(context.Background()).Errorf("Failed to get gpuClasses ", err)
+					return nil, err
+				} else {
+					gpuClasses = classes
+				}
+			}
+			for _, gpuClass := range gpuClasses.Items {
+				if strings.TrimSpace(strings.ToLower(gpuClass.Name)) == gpuCleaned {
+					saladClientGpuIds = append(saladClientGpuIds, gpuClass.Id)
+					break
+				}
+			}
+		}
+	}
+	return saladClientGpuIds, nil
+}
+
+func (p *SaladCloudProvider) getCountryCodes(pod *corev1.Pod) ([]saladclient.CountryCode, error) {
+	countryCodes := make([]saladclient.CountryCode, 0)
+	countryCodes = append(countryCodes, "US")
+	countryCodesFromAnnotation, ok := pod.ObjectMeta.Annotations["salad.com/country-codes"]
+	if !ok {
+		return countryCodes, nil
+	}
+	codes := strings.Split(countryCodesFromAnnotation, ",")
+	for _, code := range codes {
+		cc, err := saladclient.NewCountryCodeFromValue(code)
+		if err != nil {
+			return []saladclient.CountryCode{}, errors.WithMessage(err, "Invalid country code provided: "+code)
+		}
+		countryCodes = append(countryCodes, *cc)
+	}
+	return countryCodes, nil
+}
+
+func (p *SaladCloudProvider) getNetworking(pod *corev1.Pod) (*saladclient.CreateContainerGroupNetworking, error) {
+	protocol, hasProtocol := pod.ObjectMeta.Annotations["salad.com/networking-protocol"]
+	port, hasPort := pod.ObjectMeta.Annotations["salad.com/networking-port"]
+	auth, hasAuth := pod.ObjectMeta.Annotations["salad.com/networking-auth"]
+	if !hasProtocol || !hasPort || !hasAuth {
+		return nil, nil
+	}
+	networkingProtocol, err := saladclient.NewContainerNetworkingProtocolFromValue(protocol)
+	if err != nil {
+		return nil, err
+	}
+	parsedPortInt, err := strconv.Atoi(port)
+	if err != nil {
+		return nil, err
+	}
+	parsedAuth := false
+	if strings.ToLower(auth) == "true" {
+		parsedAuth = true
+	}
+	return saladclient.NewCreateContainerGroupNetworking(*networkingProtocol, int32(parsedPortInt), parsedAuth), nil
+}
+
+func (p *SaladCloudProvider) getRestartPolicy(pod *corev1.Pod) (*saladclient.ContainerRestartPolicy, error) {
+	restartPolicy := "never"
+	if pod.Spec.RestartPolicy == corev1.RestartPolicyAlways {
+		restartPolicy = "always"
+	}
+	if pod.Spec.RestartPolicy == corev1.RestartPolicyOnFailure {
+		restartPolicy = "on_failure"
+	}
+	if pod.Spec.RestartPolicy == corev1.RestartPolicyNever {
+		restartPolicy = "never"
+	}
+	return saladclient.NewContainerRestartPolicyFromValue(restartPolicy)
+}

internal/utils/utils.go

@@ -19,7 +19,7 @@ func roundUpToNearest(value int64, list []int64) int64 {
 func GetPodResource(podSpec corev1.PodSpec) (cpu int64, memory int64) {
 	allowedCPUValues := []int64{1, 2, 3, 4, 6, 8, 12, 16}
 
-	allowedMemoryValues := []int64{1024, 2048, 3, 4, 5, 6, 12} // in GB
+	allowedMemoryValues := []int64{1024, 2048, 3, 4, 5, 6, 12, 16, 24, 30} // in GB
 
 	for _, container := range podSpec.Containers {
 		// Convert milliCPU to cores and round to nearest value in the list

sample-deployment.yaml

@@ -2,6 +2,12 @@ apiVersion: apps/v1
 kind: Deployment
 metadata:
   name: demo-deployment
+  annotations:
+    salad.com/country-codes: "ca,us,mx"
+    salad.com/gpu-classes: "GTX 1070 (8 GB), ffc51032-64d2-4df3-855a-f3a649895c0f"
+    salad.com/networking-protocol: "http"
+    salad.com/networking-port: "80"
+    salad.com/networking-auth: "true"
 spec:
   replicas: 1
   selector:
@@ -29,6 +35,7 @@ spec:
               cpu: "1"
       nodeSelector:
         kubernetes.io/role: agent
+      restartPolicy: Always
       tolerations:
         - key: "virtual-kubelet.io/provider"
           operator: "Equal"