[Feat] Support scaling LoRA adapters to multiple replicas (#1132)

* Remove unused labels import
* update sample files to support scaling
* Update the lora scaling to be based on the replica count set as part of the ModelAdapter
* fix linting

---------

Signed-off-by: dittops <[email protected]>

Author: dittops

pkg/controller/modeladapter/modeladapter_controller.go

@@ -25,15 +25,16 @@ import (
 	"k8s.io/utils/ptr"
 )
 
-func buildModelAdapterEndpointSlice(instance *modelv1alpha1.ModelAdapter, pod *corev1.Pod) *discoveryv1.EndpointSlice {
+func buildModelAdapterEndpointSlice(instance *modelv1alpha1.ModelAdapter, pods []corev1.Pod) *discoveryv1.EndpointSlice {
 	serviceLabels := map[string]string{
 		"kubernetes.io/service-name": instance.Name,
 	}
 
-	addresses := []discoveryv1.Endpoint{
-		{
+	addresses := make([]discoveryv1.Endpoint, 0, len(pods))
+	for _, pod := range pods {
+		addresses = append(addresses, discoveryv1.Endpoint{
 			Addresses: []string{pod.Status.PodIP},
-		},
+		})
 	}
 
 	ports := []discoveryv1.EndpointPort{

pkg/controller/modeladapter/resources.go

@@ -37,15 +37,16 @@ func TestBuildModelAdapterEndpointSlice(t *testing.T) {
 		},
 	}
 
-	// Mock input for Pod
-	pod := &corev1.Pod{
+	// Mock input for Pods
+	pod := corev1.Pod{
 		Status: corev1.PodStatus{
 			PodIP: "192.168.1.1",
 		},
 	}
+	pods := []corev1.Pod{pod}
 
 	// Call the function to test
-	endpointSlice := buildModelAdapterEndpointSlice(instance, pod)
+	endpointSlice := buildModelAdapterEndpointSlice(instance, pods)
 
 	// Check EndpointSlice metadata
 	assert.Equal(t, "test-instance", endpointSlice.Name)

pkg/controller/modeladapter/resources_test.go

@@ -0,0 +1,20 @@
+apiVersion: model.aibrix.ai/v1alpha1
+kind: ModelAdapter
+metadata:
+  name: sample-lora-multi-replica
+  namespace: default
+spec:
+  # Specify the number of replicas for the adapter
+  # The adapter will be loaded on this many pods matching the selector
+  replicas: 3
+  # Pod selector to identify which pods can host this adapter
+  podSelector:
+    matchLabels:
+      model.aibrix.ai/name: qwen-coder-1-5b-instruct
+      adapter.model.aibrix.ai/enabled: "true"
+  # URL for the LoRA adapter artifact
+  artifactURL: "huggingface://SomethingNew/lora-adapter-demo"
+  # Optional: Additional configuration
+  additionalConfig:
+    rank: "16"
+    alpha: "32"

samples/adapter/adapter-multi-replica.yaml

@@ -7,6 +7,9 @@ metadata:
     model.aibrix.ai/name: "qwen-code-lora"
     model.aibrix.ai/port: "8000"
 spec:
+  # Optional: Number of replicas for the adapter (default: 1)
+  # Uncomment to load adapter on multiple pods
+  # replicas: 3
   baseModel: qwen-coder-1-5b-instruct
   podSelector:
     matchLabels:

samples/adapter/adapter.yaml

@@ -16,6 +16,7 @@ spec:
     metadata:
       labels:
         model.aibrix.ai/name: qwen-coder-1-5b-instruct
+        adapter.model.aibrix.ai/enabled: "true"
     spec:
       containers:
         - command:

samples/adapter/base-api-key.yaml

@@ -16,6 +16,7 @@ spec:
     metadata:
       labels:
         model.aibrix.ai/name: qwen-coder-1-5b-instruct
+        adapter.model.aibrix.ai/enabled: "true"
     spec:
       containers:
         - command:

samples/adapter/base-without-runtime.yaml

@@ -16,6 +16,7 @@ spec:
     metadata:
       labels:
         model.aibrix.ai/name: qwen-coder-1-5b-instruct
+        adapter.model.aibrix.ai/enabled: "true"
     spec:
       containers:
         - command: