feat: kubernetes overrides for the entrypoint and cmd (#1396)

ATTRIBUTIONS-Go.md

@@ -1484,7 +1484,7 @@ Subdependencies:
 * `golang.org/x/net`
 * `golang.org/x/text`
 * `gopkg.in/inf.v0`
-* `gopkg.in/yaml.v2`
+* `github.com/goccy/go-yaml`
 * `k8s.io/klog/v2`
 * `k8s.io/utils`
 * `sigs.k8s.io/json`
@@ -3179,7 +3179,7 @@ THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 ```
-### gopkg.in/yaml.v2
+### github.com/goccy/go-yaml
 License Identifier: Apache-2.0
 License Text:
@@ -4864,7 +4864,7 @@ Subdependencies:
 * `google.golang.org/genproto/googleapis/api`
 * `google.golang.org/genproto/googleapis/rpc`
 * `google.golang.org/protobuf`
-* `gopkg.in/yaml.v2`
+* `github.com/goccy/go-yaml`
 ### github.com/dustin/go-humanize
@@ -8180,7 +8180,7 @@ Subdependencies:
 * `golang.org/x/text`
 * `google.golang.org/genproto/googleapis/rpc`
 * `gopkg.in/inf.v0`
-* `gopkg.in/yaml.v2`
+* `github.com/goccy/go-yaml`
 * `k8s.io/klog/v2`
 * `k8s.io/utils`
 * `sigs.k8s.io/json`
@@ -9056,7 +9056,7 @@ Subdependencies:
 * `google.golang.org/genproto/googleapis/api`
 * `google.golang.org/protobuf`
 * `gopkg.in/inf.v0`
-* `gopkg.in/yaml.v2`
+* `github.com/goccy/go-yaml`
 * `gopkg.in/yaml.v3`
 * `k8s.io/api`
 * `k8s.io/klog/v2`
@@ -22970,7 +22970,7 @@ Subdependencies:
 * `google.golang.org/protobuf`
 * `gopkg.in/evanphx/json-patch.v4`
 * `gopkg.in/inf.v0`
-* `gopkg.in/yaml.v2`
+* `github.com/goccy/go-yaml`
 * `gopkg.in/yaml.v3`
 * `k8s.io/component-base`
 * `k8s.io/gengo/v2`

deploy/cloud/helm/README.md

@@ -86,7 +86,7 @@ kubectl get storageclass
 
 1. Set the required environment variables:
 ```bash
-export PROJECT_ROOT=($pwd)
+export PROJECT_ROOT=$(pwd)
 export DOCKER_USERNAME=<your-docker-username>
 export DOCKER_PASSWORD=<your-docker-password>
 export DOCKER_SERVER=<your-docker-server>

deploy/cloud/operator/api/dynamo/common/common.go

@@ -62,4 +62,5 @@ type ExtraPodSpec struct {
 	Containers                []corev1.Container                `json:"containers,omitempty"`
 	ServiceAccountName        string                            `json:"serviceAccountName,omitempty"`
 	PriorityClassName         string                            `json:"priorityClassName,omitempty"`
+	MainContainer             *corev1.Container                 `json:"mainContainer,omitempty"`
 }

deploy/cloud/operator/api/dynamo/common/zz_generated.deepcopy.go

@@ -151,6 +151,11 @@ func (in *ExtraPodSpec) DeepCopyInto(out *ExtraPodSpec) {
 			(*in)[i].DeepCopyInto(&(*out)[i])
 		}
 	}
+	if in.MainContainer != nil {
+		in, out := &in.MainContainer, &out.MainContainer
+		*out = new(v1.Container)
+		(*in).DeepCopyInto(*out)
+	}
 }
 
 // DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new ExtraPodSpec.

deploy/cloud/operator/go.mod

@@ -11,6 +11,7 @@ require (
 	github.com/awslabs/amazon-ecr-credential-helper/ecr-login v0.9.1
 	github.com/bsm/gomega v1.27.10
 	github.com/chrismellard/docker-credential-acr-env v0.0.0-20230304212654-82a0ddb27589
+	github.com/goccy/go-yaml v1.18.0
 	github.com/google/go-cmp v0.7.0
 	github.com/google/go-containerregistry v0.20.5
 	github.com/huandu/xstrings v1.4.0
@@ -22,7 +23,6 @@ require (
 	github.com/sergeymakinen/go-quote v1.1.0
 	github.com/sirupsen/logrus v1.9.3
 	go.etcd.io/etcd/client/v3 v3.5.16
-	gopkg.in/yaml.v2 v2.4.0
 	istio.io/api v1.23.1
 	istio.io/client-go v1.23.1
 	k8s.io/api v0.32.3

deploy/cloud/operator/go.sum

@@ -111,6 +111,8 @@ github.com/go-openapi/swag v0.23.0 h1:vsEVJDUo2hPJ2tu0/Xc+4noaxyEffXNIs3cOULZ+Gr
 github.com/go-openapi/swag v0.23.0/go.mod h1:esZ8ITTYEsH1V2trKHjAN8Ai7xHb8RV+YSZ577vPjgQ=
 github.com/go-task/slim-sprig/v3 v3.0.0 h1:sUs3vkvUymDpBKi3qH1YSqBQk9+9D/8M2mN1vB6EwHI=
 github.com/go-task/slim-sprig/v3 v3.0.0/go.mod h1:W848ghGpv3Qj3dhTPRyJypKRiqCdHZiAzKg9hl15HA8=
+github.com/goccy/go-yaml v1.18.0 h1:8W7wMFS12Pcas7KU+VVkaiCng+kG8QiFeFwzFb+rwuw=
+github.com/goccy/go-yaml v1.18.0/go.mod h1:XBurs7gK8ATbW4ZPGKgcbrY1Br56PdM69F7LkFRi1kA=
 github.com/godbus/dbus/v5 v5.0.4/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA=
 github.com/gogo/protobuf v1.3.2 h1:Ov1cvc58UF3b5XjBnZv7+opcTcQFZebYjWzi34vdm4Q=
 github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69NZV8Q=
@@ -310,8 +312,6 @@ gopkg.in/evanphx/json-patch.v4 v4.12.0 h1:n6jtcsulIzXPJaxegRbvFNNrZDjbij7ny3gmSP
 gopkg.in/evanphx/json-patch.v4 v4.12.0/go.mod h1:p8EYWUEYMpynmqDbY58zCKCFZw8pRWMG4EsWvDvM72M=
 gopkg.in/inf.v0 v0.9.1 h1:73M5CoZyi3ZLMOyDlQh031Cx6N9NDJ2Vvfl76EDAgDc=
 gopkg.in/inf.v0 v0.9.1/go.mod h1:cWUDdTG/fYaXco+Dcufb5Vnc6Gp2YChqWtbxRZE0mXw=
-gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY=
-gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
 gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

deploy/cloud/operator/internal/controller/dynamocomponent_controller.go

@@ -40,6 +40,7 @@ import (
 	"github.com/apparentlymart/go-shquot/shquot"
 	"github.com/awslabs/amazon-ecr-credential-helper/ecr-login"
 	"github.com/chrismellard/docker-credential-acr-env/pkg/credhelper"
+	"github.com/goccy/go-yaml"
 	"github.com/google/go-containerregistry/pkg/authn"
 	"github.com/google/go-containerregistry/pkg/name"
 	"github.com/google/go-containerregistry/pkg/v1/google"
@@ -49,7 +50,6 @@ import (
 	"github.com/rs/xid"
 	"github.com/sergeymakinen/go-quote/unix"
 	"github.com/sirupsen/logrus"
-	"gopkg.in/yaml.v2"
 	batchv1 "k8s.io/api/batch/v1"
 	corev1 "k8s.io/api/core/v1"
 	k8serrors "k8s.io/apimachinery/pkg/api/errors"

deploy/cloud/operator/internal/controller/dynamocomponentdeployment_controller.go

@@ -1345,6 +1345,7 @@ func getDynamoComponentRepositoryNameAndDynamoComponentVersion(dynamoComponent *
 
 //nolint:gocyclo,nakedret
 func (r *DynamoComponentDeploymentReconciler) generatePodTemplateSpec(ctx context.Context, opt generateResourceOption) (podTemplateSpec *corev1.PodTemplateSpec, err error) {
+	logs := log.FromContext(ctx)
 	podLabels := r.getKubeLabels(opt.dynamoComponentDeployment, opt.dynamoComponent)
 	if opt.isStealingTrafficDebugModeEnabled {
 		podLabels[commonconsts.KubeLabelDynamoDeploymentTargetType] = DeploymentTargetTypeDebug
@@ -1659,6 +1660,28 @@ func (r *DynamoComponentDeploymentReconciler) generatePodTemplateSpec(ctx contex
 		container.SecurityContext.RunAsUser = &[]int64{0}[0]
 	}
 
+	// For now only overwrite the command and args.
+	if opt.dynamoComponentDeployment.Spec.ExtraPodSpec != nil {
+		extraPodSpecMainContainer := opt.dynamoComponentDeployment.Spec.ExtraPodSpec.MainContainer
+		if extraPodSpecMainContainer != nil {
+			if len(extraPodSpecMainContainer.Command) > 0 {
+				logs.Info("Overriding container '" + container.Name + "' Command with: " + strings.Join(extraPodSpecMainContainer.Command, " "))
+				container.Command = extraPodSpecMainContainer.Command
+			}
+			if len(extraPodSpecMainContainer.Args) > 0 {
+				// Special case: if command is "sh -c", we must collapse args into a single string
+				if len(container.Command) == 2 && container.Command[0] == "sh" && container.Command[1] == "-c" {
+					joinedArgs := strings.Join(extraPodSpecMainContainer.Args, " ")
+					logs.Info("Special case detected for container '" + container.Name + "': Command is 'sh -c'; collapsing Args to: " + joinedArgs)
+					container.Args = []string{joinedArgs}
+				} else {
+					logs.Info("Overriding container '" + container.Name + "' Args with: " + strings.Join(extraPodSpecMainContainer.Args, " "))
+					container.Args = extraPodSpecMainContainer.Args
+				}
+			}
+		}
+	}
+
 	containers = append(containers, container)
 
 	debuggerImage := "python:3.12-slim"

deploy/cloud/operator/internal/dynamo/graph.go

@@ -42,7 +42,7 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/log"
 
 	"github.com/ai-dynamo/dynamo/deploy/cloud/operator/internal/archive"
-	"gopkg.in/yaml.v2"
+	"github.com/goccy/go-yaml"
 )
 
 const (
@@ -83,6 +83,7 @@ type Config struct {
 	ApiEndpoints []string             `yaml:"api_endpoints,omitempty"`
 	Workers      *int32               `yaml:"workers,omitempty"`
 	TotalGpus    *int32               `yaml:"total_gpus,omitempty"`
+	ExtraPodSpec *common.ExtraPodSpec `yaml:"extraPodSpec,omitempty"`
 }
 
 type ServiceConfig struct {
@@ -360,6 +361,7 @@ func GenerateDynamoComponentsDeployments(ctx context.Context, parentDynamoGraphD
 				return nil, err
 			}
 		}
+
 		deployment.Spec.Autoscaling = &v1alpha1.Autoscaling{
 			Enabled: false,
 		}
@@ -368,6 +370,12 @@ func GenerateDynamoComponentsDeployments(ctx context.Context, parentDynamoGraphD
 			deployment.Spec.Autoscaling.MinReplicas = service.Config.Autoscaling.MinReplicas
 			deployment.Spec.Autoscaling.MaxReplicas = service.Config.Autoscaling.MaxReplicas
 		}
+
+		// Override properties from the ExtraPodSpec (i.e. command and args) if provided.
+		if err := mergeExtraPodSpec(deployment, &service.Config); err != nil {
+			return nil, err
+		}
+
 		// override the component config with the component config that is in the parent deployment
 		if configOverride, ok := parentDynamoGraphDeployment.Spec.Services[service.Name]; ok {
 			err := mergo.Merge(&deployment.Spec.DynamoComponentDeploymentSharedSpec, configOverride.DynamoComponentDeploymentSharedSpec, mergo.WithOverride)
@@ -527,3 +535,22 @@ func mergeEnvs(common, specific []corev1.EnvVar) []corev1.EnvVar {
 	}
 	return merged
 }
+
+// mergeExtraPodSpec merges the ExtraPodSpec from service config into the deployment spec
+func mergeExtraPodSpec(deployment *v1alpha1.DynamoComponentDeployment, serviceConfig *Config) error {
+	if serviceConfig.ExtraPodSpec != nil && serviceConfig.ExtraPodSpec.MainContainer != nil {
+		if deployment.Spec.DynamoComponentDeploymentSharedSpec.ExtraPodSpec == nil {
+			deployment.Spec.DynamoComponentDeploymentSharedSpec.ExtraPodSpec = new(common.ExtraPodSpec)
+		}
+		err := mergo.Merge(
+			deployment.Spec.DynamoComponentDeploymentSharedSpec.ExtraPodSpec,
+			serviceConfig.ExtraPodSpec,
+			mergo.WithOverride,
+			mergo.WithOverwriteWithEmptyValue,
+		)
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}

deploy/cloud/operator/internal/dynamo/graph_test.go

@@ -1402,6 +1402,69 @@ func TestGenerateDynamoComponentsDeployments(t *testing.T) {
 			},
 			wantErr: false,
 		},
+		{
+			name: "Test GenerateDynamoComponentsDeployments with ExtraPodSpec.MainContainer Command and Args",
+			args: args{
+				parentDynamoGraphDeployment: &v1alpha1.DynamoGraphDeployment{
+					ObjectMeta: metav1.ObjectMeta{
+						Name:      "test-dynamographdeployment",
+						Namespace: "default",
+					},
+					Spec: v1alpha1.DynamoGraphDeploymentSpec{
+						DynamoGraph: "dynamocomponent:ac4e234",
+					},
+				},
+				config: &DynamoGraphConfig{
+					DynamoTag: "dynamocomponent:MyServiceWithOverrides",
+					Services: []ServiceConfig{
+						{
+							Name:         "service1",
+							Dependencies: []map[string]string{},
+							Config: Config{
+								ExtraPodSpec: &compounaiCommon.ExtraPodSpec{
+									MainContainer: &corev1.Container{
+										Command: []string{"sh", "-c"},
+										Args:    []string{"echo hello world", "sleep 99999"},
+									},
+								},
+							},
+						},
+					},
+				},
+				ingressSpec: &v1alpha1.IngressSpec{},
+			},
+			want: map[string]*v1alpha1.DynamoComponentDeployment{
+				"service1": {
+					ObjectMeta: metav1.ObjectMeta{
+						Name:      "test-dynamographdeployment-service1",
+						Namespace: "default",
+						Labels: map[string]string{
+							commonconsts.KubeLabelDynamoComponent: "service1",
+						},
+					},
+					Spec: v1alpha1.DynamoComponentDeploymentSpec{
+						DynamoComponent: "dynamocomponent:ac4e234",
+						DynamoTag:       "dynamocomponent:MyServiceWithOverrides",
+						DynamoComponentDeploymentSharedSpec: v1alpha1.DynamoComponentDeploymentSharedSpec{
+							ServiceName: "service1",
+							Autoscaling: &v1alpha1.Autoscaling{
+								Enabled: false,
+							},
+							Labels: map[string]string{
+								commonconsts.KubeLabelDynamoComponent: "service1",
+							},
+							ExtraPodSpec: &compounaiCommon.ExtraPodSpec{
+								MainContainer: &corev1.Container{
+									Command: []string{"sh", "-c"},
+									Args:    []string{"echo hello world", "sleep 99999"},
+								},
+							},
+						},
+					},
+				},
+			},
+			wantErr: false,
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {

deploy/sdk/src/dynamo/sdk/cli/build.py

@@ -22,6 +22,7 @@ import importlib.util
 import inspect
 import logging
 import os
+import shlex
 import shutil
 import subprocess
 import sys
@@ -42,6 +43,7 @@ from dynamo.sdk.core.protocol.deployment import Service
 from dynamo.sdk.core.protocol.interface import (
     DynamoConfig,
     DynamoTransport,
+    KubernetesOverrides,
     LinkedServices,
     ServiceInterface,
 )
@@ -110,6 +112,7 @@ class ServiceConfig(BaseModel):
     dynamo: DynamoConfig = Field(default_factory=DynamoConfig)
     http_exposed: bool = False
     api_endpoints: t.List[str] = Field(default_factory=list)
+    kubernetes_overrides: KubernetesOverrides | None = None
 
 
 class ServiceInfo(BaseModel):
@@ -147,6 +150,7 @@ class ServiceInfo(BaseModel):
             dynamo=DynamoConfig(**service.config.dynamo.model_dump()),
             http_exposed=len(api_endpoints) > 0,
             api_endpoints=api_endpoints,
+            kubernetes_overrides=service.config.kubernetes_overrides,
         )
 
         return cls(
@@ -221,6 +225,33 @@ class ManifestInfo(BaseModel):
                     "dynamo": service["config"]["dynamo"],
                 },
             }
+            # Add kubernetes_overrides if present.
+            if (
+                "kubernetes_overrides" in service["config"]
+                and service["config"]["kubernetes_overrides"]
+            ):
+                # Map kubernetes_overrides fields to mainContainer if present.
+                kube_overrides = service["config"]["kubernetes_overrides"]
+                main_container = {}
+
+                entrypoint = kube_overrides.get("entrypoint")
+                if entrypoint:
+                    if isinstance(entrypoint, str):
+                        main_container["command"] = shlex.split(entrypoint)
+                    else:
+                        main_container["command"] = shlex.split(entrypoint[0])
+
+                cmd = kube_overrides.get("cmd")
+                if cmd:
+                    if isinstance(cmd, str):
+                        main_container["args"] = shlex.split(cmd)
+                    else:
+                        main_container["args"] = shlex.split(cmd[0])
+
+                if main_container:
+                    service_dict["config"]["extraPodSpec"] = {
+                        "mainContainer": main_container
+                    }
 
             # Add HTTP configuration if exposed
             if service["config"]["http_exposed"]:
@@ -309,6 +340,7 @@ class Package:
         build_ctx: str,
         version: t.Optional[str] = None,
     ) -> Package:
+        """Create a package from a build config."""
         dyn_svc = cls.dynamo_service(build_config, build_ctx)
 
         # Get service name for package

deploy/sdk/src/dynamo/sdk/core/protocol/interface.py

@@ -21,7 +21,7 @@ from enum import Enum, auto
 from typing import Any, Dict, Generic, List, Optional, Set, Tuple, Type, TypeVar
 
 from fastapi import FastAPI
-from pydantic import BaseModel, ConfigDict, Field
+from pydantic import BaseModel, ConfigDict, Field, field_validator
 
 from .deployment import Env
 
@@ -74,6 +74,24 @@ class ResourceConfig(BaseModel):
     gpu: str = Field(default="0")
 
 
+class KubernetesOverrides(BaseModel):
+    """Class for kubernetes overrides from the sdk to limit to supported fields."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    entrypoint: List[str] | None = None
+    cmd: List[str] | None = None
+
+    @field_validator("entrypoint", "cmd", mode="before")
+    @classmethod
+    def _coerce_str_to_list(cls, v):
+        if v is None or isinstance(v, list):
+            return v
+        if isinstance(v, str):
+            return [v]
+        raise TypeError("Must be str or list[str]")
+
+
 class ServiceConfig(BaseModel):
     """Base service configuration that can be extended by adapters"""
 
@@ -83,6 +101,7 @@ class ServiceConfig(BaseModel):
     image: str | None = None
     envs: List[Env] | None = None
     labels: Dict[str, str] | None = None
+    kubernetes_overrides: KubernetesOverrides | None = None
 
 
 class DynamoEndpointInterface(ABC):

docs/examples/hello_world.md

@@ -95,12 +95,19 @@ This example can be deployed to a Kubernetes cluster using [Dynamo Cloud](../../
 
 You must have first followed the instructions in [deploy/cloud/helm/README.md](https://github.com/ai-dynamo/dynamo/blob/main/deploy/cloud/helm/README.md) to create your Dynamo cloud deployment.
 
-### Deployment Steps
+### Deployment Steps For your Hello World graph.
 
 For detailed deployment instructions, please refer to the [Operator Deployment Guide](../../docs/guides/dynamo_deploy/operator_deployment.md). The following are the specific commands for the hello world example:
 
 ```bash
-# Set your project root directory
+Make sure your dynamo cloud deploy.sh script from the prior step finished successfully and setup port forwaring in another window
+per its suggestion.
+
+kubectl port-forward svc/...-dynamo-api-store <local-port>:80 -n $NAMESPACE
+
+
+# Set your dynamo root directory
+cd <root dynamo folder>
 export PROJECT_ROOT=$(pwd)
 
 # Configure environment variables (see operator_deployment.md for details)

examples/hello_world/hello_world.py

@@ -125,6 +125,11 @@ class Middle:
 @service(
     dynamo={"namespace": "inference"},
     image=DYNAMO_IMAGE,
+    # Example of kubernetes overrides if needed.
+    # kubernetes_overrides={
+    #     "entrypoint": ["sh -c"],
+    #     "cmd": ["echo hello from FrontEnd!"],
+    # },
 )
 class Frontend:
     """A simple frontend HTTP API that forwards requests to the dynamo graph."""