Add params support for Server resource (#214)

Required for providing additional params to e.g. basaran image.

api/v1/server_types.go

@@ -2,6 +2,7 @@ package v1
 
 import (
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/util/intstr"
 	"k8s.io/utils/ptr"
 )
 
@@ -21,6 +22,9 @@ type ServerSpec struct {
 
 	// Model references the Model object to be served.
 	Model ObjectRef `json:"model,omitempty"`
+
+	// Params will be passed into the loading process as environment variables.
+	Params map[string]intstr.IntOrString `json:"params,omitempty"`
 }
 
 // ServerStatus defines the observed state of Server
@@ -53,6 +57,10 @@ type Server struct {
 	Status ServerStatus `json:"status,omitempty"`
 }
 
+func (s *Server) GetParams() map[string]intstr.IntOrString {
+	return s.Spec.Params
+}
+
 func (s *Server) GetBuild() *Build {
 	return s.Spec.Build
 }

cmd/controllermanager/main.go

@@ -159,6 +159,10 @@ func main() {
 		Scheme: mgr.GetScheme(),
 		Cloud:  cld,
 		SCI:    sciClient,
+		ParamsReconciler: &controller.ParamsReconciler{
+			Scheme: mgr.GetScheme(),
+			Client: mgr.GetClient(),
+		},
 	}).SetupWithManager(mgr); err != nil {
 		setupLog.Error(err, "unable to create controller", "controller", "Server")
 		os.Exit(1)

config/crd/bases/substratus.ai_servers.yaml

@@ -113,6 +113,15 @@ spec:
                 required:
                 - name
                 type: object
+              params:
+                additionalProperties:
+                  anyOf:
+                  - type: integer
+                  - type: string
+                  x-kubernetes-int-or-string: true
+                description: Params will be passed into the loading process as environment
+                  variables.
+                type: object
               resources:
                 description: Resources are the compute resources required by the container.
                 properties:

docs/api/generated.md

@@ -347,6 +347,7 @@ _Appears in:_
 | `build` _[Build](#build)_ | Build specifies how to build an image. |
 | `resources` _[Resources](#resources)_ | Resources are the compute resources required by the container. |
 | `model` _[ObjectRef](#objectref)_ | Model references the Model object to be served. |
+| `params` _object (keys:string, values:IntOrString)_ | Params will be passed into the loading process as environment variables. |
 
 
 ### ServerStatus

examples/falcon-7b-instruct/server-4bit-any-gpu.yaml

@@ -0,0 +1,13 @@
+apiVersion: substratus.ai/v1
+kind: Server
+metadata:
+  name: falcon-7b-instruct
+spec:
+  image: substratusai/model-server-basaran
+  model:
+    name: falcon-7b-instruct
+  params:
+    load_in_4bit: "true"
+  resources:
+    gpu:
+      count: 1

internal/controller/main_test.go

@@ -115,6 +115,10 @@ func TestMain(m *testing.M) {
 		Scheme: mgr.GetScheme(),
 		Cloud:  testCloud,
 		SCI:    sciClient,
+		ParamsReconciler: &controller.ParamsReconciler{
+			Scheme: mgr.GetScheme(),
+			Client: mgr.GetClient(),
+		},
 	}).SetupWithManager(mgr)
 	requireNoError(err)
 	err = (&controller.BuildReconciler{

internal/controller/server_controller.go

@@ -34,6 +34,8 @@ type ServerReconciler struct {
 	Cloud cloud.Cloud
 	SCI   sci.ControllerClient
 
+	*ParamsReconciler
+
 	// log should be used outside the context of Reconcile()
 	log logr.Logger
 }
@@ -61,6 +63,10 @@ func (r *ServerReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctr
 		return ctrl.Result{}, nil
 	}
 
+	if result, err := r.ReconcileParamsConfigMap(ctx, &server); !result.success {
+		return result.Result, err
+	}
+
 	if result, err := r.reconcileServer(ctx, &server); !result.success {
 		return result.Result, err
 	}
@@ -142,6 +148,7 @@ func (r *ServerReconciler) serverDeployment(server *apiv1.Server, model *apiv1.M
 							Image:           server.GetImage(),
 							ImagePullPolicy: "Always",
 							Command:         server.Spec.Command,
+							Env:             paramsToEnv(server.Spec.Params),
 							Ports: []corev1.ContainerPort{
 								{
 									Name:          modelServerHTTPServePortName,
@@ -166,6 +173,10 @@ func (r *ServerReconciler) serverDeployment(server *apiv1.Server, model *apiv1.M
 		},
 	}
 
+	if err := mountParamsConfigMap(&deploy.Spec.Template.Spec, server, containerName); err != nil {
+		return nil, fmt.Errorf("mounting params configmap: %w", err)
+	}
+
 	if err := r.Cloud.MountBucket(&deploy.Spec.Template.ObjectMeta, &deploy.Spec.Template.Spec, model, cloud.MountBucketConfig{
 		Name: "model",
 		Mounts: []cloud.BucketMount{