koordlet: support gpu share with HAMi

Signed-off-by: wangjianyu.wjy <[email protected]>

apis/extension/device_share.go

@@ -58,6 +58,8 @@ const (
 	LabelGPUModel                   string = NodeDomainPrefix + "/gpu-model"
 	LabelGPUDriverVersion           string = NodeDomainPrefix + "/gpu-driver-version"
 	LabelSecondaryDeviceWellPlanned string = NodeDomainPrefix + "/secondary-device-well-planned"
+
+	LabelGPUIsolationProvider = DomainPrefix + "gpu-isolation-provider"
 )
 
 // DeviceAllocations would be injected into Pod as form of annotation during Pre-bind stage.
@@ -202,6 +204,12 @@ const (
 	GPUPartitionPolicyPrefer GPUPartitionPolicy = "Prefer"
 )
 
+type GPUIsolationProvider string
+
+const (
+	GPUIsolationProviderHAMICore GPUIsolationProvider = "HAMi-core"
+)
+
 func GetDeviceAllocations(podAnnotations map[string]string) (DeviceAllocations, error) {
 	deviceAllocations := DeviceAllocations{}
 	data, ok := podAnnotations[AnnotationDeviceAllocated]

pkg/koordlet/runtimehooks/hooks/gpu/gpu.go

@@ -26,6 +26,7 @@ import (
 	schedulingv1alpha1 "github.com/koordinator-sh/koordinator/apis/scheduling/v1alpha1"
 	"github.com/koordinator-sh/koordinator/pkg/koordlet/runtimehooks/hooks"
 	"github.com/koordinator-sh/koordinator/pkg/koordlet/runtimehooks/protocol"
+	"github.com/koordinator-sh/koordinator/pkg/koordlet/util/system"
 	rmconfig "github.com/koordinator-sh/koordinator/pkg/runtimeproxy/config"
 )
 
@@ -70,5 +71,30 @@ func (p *gpuPlugin) InjectContainerGPUEnv(proto protocol.HooksProtocol) error {
 		containerCtx.Response.AddContainerEnvs = make(map[string]string)
 	}
 	containerCtx.Response.AddContainerEnvs[GpuAllocEnv] = strings.Join(gpuIDs, ",")
+	if containerReq.PodAnnotations[ext.LabelGPUIsolationProvider] == string(ext.GPUIsolationProviderHAMICore) {
+		gpuResources := devices[0].Resources
+		gpuMemoryRatio, ok := gpuResources[ext.ResourceGPUMemoryRatio]
+		if !ok {
+			return fmt.Errorf("gpu memory ratio not found in gpu resource")
+		}
+		if gpuMemoryRatio.Value() < 100 {
+			gpuMemory, ok := gpuResources[ext.ResourceGPUMemory]
+			if !ok {
+				return fmt.Errorf("gpu memory not found in gpu resource")
+			}
+			containerCtx.Response.AddContainerEnvs["CUDA_DEVICE_MEMORY_LIMIT"] = fmt.Sprintf("%d", gpuMemory.Value())
+			gpuCore, ok := gpuResources[ext.ResourceGPUCore]
+			if ok {
+				containerCtx.Response.AddContainerEnvs["CUDA_DEVICE_SM_LIMIT"] = fmt.Sprintf("%d", gpuCore.Value())
+			}
+			containerCtx.Response.AddContainerEnvs["LD_PRELOAD"] = "/libvgpu.so"
+			containerCtx.Response.AddContainerMounts = append(containerCtx.Response.AddContainerMounts, &protocol.Mount{
+				Destination: "/libvgpu.so",
+				Type:        "bind",
+				Source:      system.Conf.HAMICoreLibraryPath,
+			})
+		}
+	}
+
 	return nil
 }

pkg/koordlet/runtimehooks/hooks/gpu/gpu_test.go

@@ -31,49 +31,82 @@ func Test_InjectContainerGPUEnv(t *testing.T) {
 		expectedAllocStr string
 		expectedError    bool
 		proto            protocol.HooksProtocol
+		expectedMounts   []*protocol.Mount
+		expectedEnvs     map[string]string
 	}{
 		{
-			"test empty proto",
-			"",
-			true,
-			nil,
+			name:             "test empty proto",
+			expectedAllocStr: "",
+			expectedError:    true,
+			proto:            nil,
 		},
 		{
-			"test normal gpu alloc",
-			"0,1",
-			false,
-			&protocol.ContainerContext{
+			name:             "test normal gpu alloc",
+			expectedAllocStr: "0,1",
+			expectedError:    false,
+			proto: &protocol.ContainerContext{
 				Request: protocol.ContainerRequest{
 					PodAnnotations: map[string]string{
 						ext.AnnotationDeviceAllocated: "{\"gpu\": [{\"minor\": 0},{\"minor\": 1}]}",
 					},
 				},
 			},
+			expectedEnvs: map[string]string{GpuAllocEnv: "0,1"},
 		},
 		{
-			"test empty gpu alloc",
-			"",
-			false,
-			&protocol.ContainerContext{
+			name:             "test empty gpu alloc",
+			expectedAllocStr: "",
+			expectedError:    false,
+			proto: &protocol.ContainerContext{
 				Request: protocol.ContainerRequest{
 					PodAnnotations: map[string]string{
 						ext.AnnotationDeviceAllocated: "{\"fpga\": [{\"minor\": 0},{\"minor\": 1}]}",
 					},
 				},
 			},
 		},
+		{
+			name:             "gpu share with HAMi",
+			expectedAllocStr: "1",
+			expectedError:    false,
+			proto: &protocol.ContainerContext{
+				Request: protocol.ContainerRequest{
+					PodAnnotations: map[string]string{
+						ext.AnnotationDeviceAllocated: `{"gpu":[{"minor":1,"resources":{"koordinator.sh/gpu-core":"50","koordinator.sh/gpu-memory":"16Gi","koordinator.sh/gpu-memory-ratio":"50"}}]}`,
+						ext.LabelGPUIsolationProvider: string(ext.GPUIsolationProviderHAMICore),
+					},
+				},
+			},
+			expectedEnvs: map[string]string{
+				GpuAllocEnv:                "1",
+				"CUDA_DEVICE_MEMORY_LIMIT": "17179869184",
+				"CUDA_DEVICE_SM_LIMIT":     "50",
+				"LD_PRELOAD":               "/libvgpu.so",
+			},
+			expectedMounts: []*protocol.Mount{
+				{
+					Destination: "/libvgpu.so",
+					Type:        "bind",
+					Source:      "/data/bin/libvgpu.so",
+				},
+			},
+		},
 	}
 	plugin := gpuPlugin{}
 	for _, tt := range tests {
-		var containerCtx *protocol.ContainerContext
-		if tt.proto != nil {
-			containerCtx = tt.proto.(*protocol.ContainerContext)
-		}
-		err := plugin.InjectContainerGPUEnv(containerCtx)
-		assert.Equal(t, tt.expectedError, err != nil, tt.name)
-		if tt.proto != nil {
-			containerCtx := tt.proto.(*protocol.ContainerContext)
-			assert.Equal(t, containerCtx.Response.AddContainerEnvs[GpuAllocEnv], tt.expectedAllocStr, tt.name)
-		}
+		t.Run(tt.name, func(t *testing.T) {
+			var containerCtx *protocol.ContainerContext
+			if tt.proto != nil {
+				containerCtx = tt.proto.(*protocol.ContainerContext)
+			}
+			err := plugin.InjectContainerGPUEnv(containerCtx)
+			assert.Equal(t, tt.expectedError, err != nil, tt.name)
+			if tt.proto != nil {
+				containerCtx := tt.proto.(*protocol.ContainerContext)
+				assert.Equal(t, containerCtx.Response.AddContainerEnvs[GpuAllocEnv], tt.expectedAllocStr, tt.name)
+				assert.Equal(t, containerCtx.Response.AddContainerEnvs, tt.expectedEnvs, tt.name)
+				assert.Equal(t, containerCtx.Response.AddContainerMounts, tt.expectedMounts, tt.name)
+			}
+		})
 	}
 }

pkg/koordlet/runtimehooks/protocol/container_context.go

@@ -181,8 +181,9 @@ func (c *ContainerRequest) FromReconciler(podMeta *statesinformer.PodMeta, conta
 }
 
 type ContainerResponse struct {
-	Resources        Resources
-	AddContainerEnvs map[string]string
+	Resources          Resources
+	AddContainerEnvs   map[string]string
+	AddContainerMounts []*Mount
 }
 
 func (c *ContainerResponse) ProxyDone(resp *runtimeapi.ContainerResourceHookResponse) {
@@ -278,6 +279,15 @@ func (c *ContainerContext) NriDone(executor resourceexecutor.ResourceUpdateExecu
 			adjust.AddEnv(k, v)
 		}
 	}
+	if len(c.Response.AddContainerMounts) != 0 {
+		for _, m := range c.Response.AddContainerMounts {
+			adjust.AddMount(&api.Mount{
+				Destination: m.Destination,
+				Type:        m.Type,
+				Source:      m.Source,
+			})
+		}
+	}
 
 	c.Update()
 

pkg/koordlet/runtimehooks/protocol/protocol.go

@@ -135,6 +135,12 @@ func (r *Resources) FromContainer(container *corev1.Container) {
 	}
 }
 
+type Mount struct {
+	Destination string `protobuf:"bytes,1,opt,name=destination,proto3" json:"destination,omitempty"`
+	Type        string `protobuf:"bytes,2,opt,name=type,proto3" json:"type,omitempty"`
+	Source      string `protobuf:"bytes,3,opt,name=source,proto3" json:"source,omitempty"`
+}
+
 func injectCPUShares(cgroupParent string, cpuShares int64, a *audit.EventHelper, e resourceexecutor.ResourceUpdateExecutor) (resourceexecutor.ResourceUpdater, error) {
 	cpuShareStr := strconv.FormatInt(cpuShares, 10)
 	updater, err := resourceexecutor.DefaultCgroupUpdaterFactory.New(sysutil.CPUSharesName, cgroupParent, cpuShareStr, a)

pkg/koordlet/util/system/config.go

@@ -44,11 +44,12 @@ type Config struct {
 	RunRootDir            string
 	RuntimeHooksConfigDir string
 
-	ContainerdEndPoint string
-	PouchEndpoint      string
-	DockerEndPoint     string
-	CrioEndPoint       string
-	DefaultRuntimeType string
+	ContainerdEndPoint  string
+	PouchEndpoint       string
+	DockerEndPoint      string
+	CrioEndPoint        string
+	DefaultRuntimeType  string
+	HAMICoreLibraryPath string
 }
 
 func init() {
@@ -86,6 +87,7 @@ func NewHostModeConfig() *Config {
 		RunRootDir:            "/run/",
 		RuntimeHooksConfigDir: "/etc/runtime/hookserver.d",
 		DefaultRuntimeType:    "containerd",
+		HAMICoreLibraryPath:   "/data/bin/libvgpu.so",
 	}
 }
 
@@ -101,6 +103,7 @@ func NewDsModeConfig() *Config {
 		RunRootDir:            "/host-run/",
 		RuntimeHooksConfigDir: "/host-etc-hookserver/",
 		DefaultRuntimeType:    "containerd",
+		HAMICoreLibraryPath:   "/data/bin/libvgpu.so",
 	}
 }
 
@@ -122,4 +125,5 @@ func (c *Config) InitFlags(fs *flag.FlagSet) {
 	fs.StringVar(&c.PouchEndpoint, "pouch-endpoint", c.PouchEndpoint, "pouch endPoint")
 
 	fs.StringVar(&c.DefaultRuntimeType, "default-runtime-type", c.DefaultRuntimeType, "default runtime type during runtime hooks handle request, candidates are containerd/docker/pouch.")
+	fs.StringVar(&c.HAMICoreLibraryPath, "hami-core-library-path", c.HAMICoreLibraryPath, "path of hami core library")
 }

pkg/koordlet/util/system/config_test.go

@@ -33,6 +33,7 @@ func Test_NewDsModeConfig(t *testing.T) {
 		RunRootDir:            "/host-run/",
 		RuntimeHooksConfigDir: "/host-etc-hookserver/",
 		DefaultRuntimeType:    "containerd",
+		HAMICoreLibraryPath:   "/data/bin/libvgpu.so",
 	}
 	defaultConfig := NewDsModeConfig()
 	assert.Equal(t, expectConfig, defaultConfig)
@@ -49,6 +50,7 @@ func Test_NewHostModeConfig(t *testing.T) {
 		RunRootDir:            "/run/",
 		RuntimeHooksConfigDir: "/etc/runtime/hookserver.d",
 		DefaultRuntimeType:    "containerd",
+		HAMICoreLibraryPath:   "/data/bin/libvgpu.so",
 	}
 	defaultConfig := NewHostModeConfig()
 	assert.Equal(t, expectConfig, defaultConfig)