fix: mpi flow and add resourceClaim (#3446)

Signed-off-by: Rohan Varma <[email protected]>

Author: nvrohanv

deploy/cloud/helm/crds/templates/nvidia.com_dynamocomponentdeployments.yaml

@@ -10173,6 +10173,26 @@ spec:
                     Resources requested and limits for this component, including CPU, memory,
                     GPUs/devices, and any runtime-specific resources.
                   properties:
+                    claims:
+                      items:
+                        description: ResourceClaim references one entry in PodSpec.ResourceClaims.
+                        properties:
+                          name:
+                            description: |-
+                              Name must match the name of one entry in pod.spec.resourceClaims of
+                              the Pod where this field is used. It makes that resource available
+                              inside a container.
+                            type: string
+                          request:
+                            description: |-
+                              Request is the name chosen for a request in the referenced claim.
+                              If empty, everything from the claim is made available, otherwise
+                              only the result of this request.
+                            type: string
+                        required:
+                          - name
+                        type: object
+                      type: array
                     limits:
                       properties:
                         cpu:

deploy/cloud/helm/crds/templates/nvidia.com_dynamographdeployments.yaml

@@ -10307,6 +10307,26 @@ spec:
                           Resources requested and limits for this component, including CPU, memory,
                           GPUs/devices, and any runtime-specific resources.
                         properties:
+                          claims:
+                            items:
+                              description: ResourceClaim references one entry in PodSpec.ResourceClaims.
+                              properties:
+                                name:
+                                  description: |-
+                                    Name must match the name of one entry in pod.spec.resourceClaims of
+                                    the Pod where this field is used. It makes that resource available
+                                    inside a container.
+                                  type: string
+                                request:
+                                  description: |-
+                                    Request is the name chosen for a request in the referenced claim.
+                                    If empty, everything from the claim is made available, otherwise
+                                    only the result of this request.
+                                  type: string
+                              required:
+                                - name
+                              type: object
+                            type: array
                           limits:
                             properties:
                               cpu:

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

@@ -32,8 +32,9 @@ type ResourceItem struct {
 }
 
 type Resources struct {
-	Requests *ResourceItem `json:"requests,omitempty"`
-	Limits   *ResourceItem `json:"limits,omitempty"`
+	Requests *ResourceItem          `json:"requests,omitempty"`
+	Limits   *ResourceItem          `json:"limits,omitempty"`
+	Claims   []corev1.ResourceClaim `json:"claims,omitempty"`
 }
 
 type DeploymentTargetHPAConf struct {

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

@@ -10173,6 +10173,26 @@ spec:
                     Resources requested and limits for this component, including CPU, memory,
                     GPUs/devices, and any runtime-specific resources.
                   properties:
+                    claims:
+                      items:
+                        description: ResourceClaim references one entry in PodSpec.ResourceClaims.
+                        properties:
+                          name:
+                            description: |-
+                              Name must match the name of one entry in pod.spec.resourceClaims of
+                              the Pod where this field is used. It makes that resource available
+                              inside a container.
+                            type: string
+                          request:
+                            description: |-
+                              Request is the name chosen for a request in the referenced claim.
+                              If empty, everything from the claim is made available, otherwise
+                              only the result of this request.
+                            type: string
+                        required:
+                          - name
+                        type: object
+                      type: array
                     limits:
                       properties:
                         cpu:

deploy/cloud/operator/config/crd/bases/nvidia.com_dynamocomponentdeployments.yaml

@@ -10307,6 +10307,26 @@ spec:
                           Resources requested and limits for this component, including CPU, memory,
                           GPUs/devices, and any runtime-specific resources.
                         properties:
+                          claims:
+                            items:
+                              description: ResourceClaim references one entry in PodSpec.ResourceClaims.
+                              properties:
+                                name:
+                                  description: |-
+                                    Name must match the name of one entry in pod.spec.resourceClaims of
+                                    the Pod where this field is used. It makes that resource available
+                                    inside a container.
+                                  type: string
+                                request:
+                                  description: |-
+                                    Request is the name chosen for a request in the referenced claim.
+                                    If empty, everything from the claim is made available, otherwise
+                                    only the result of this request.
+                                  type: string
+                              required:
+                                - name
+                              type: object
+                            type: array
                           limits:
                             properties:
                               cpu:

deploy/cloud/operator/config/crd/bases/nvidia.com_dynamographdeployments.yaml

@@ -468,6 +468,12 @@ func GetResourcesConfig(resources *common.Resources) (*corev1.ResourceRequiremen
 			currentResources.Requests[corev1.ResourceName(k)] = q
 		}
 	}
+	if resources.Claims != nil {
+		if currentResources.Claims == nil {
+			currentResources.Claims = make([]corev1.ResourceClaim, 0)
+		}
+		currentResources.Claims = append(currentResources.Claims, resources.Claims...)
+	}
 	return currentResources, nil
 }
 

deploy/cloud/operator/internal/controller_common/resource.go

@@ -143,12 +143,12 @@ func (b *TRTLLMBackend) setupLeaderContainer(container *corev1.Container, number
 
 	// Build mpirun command with explicit SSH configuration and environment variables
 	// Wrap the entire command (trtllm-llmapi-launch + original command) in bash -c for proper shell interpretation
-	wrappedCommand := fmt.Sprintf("bash -c 'source /opt/dynamo/venv/bin/activate && trtllm-llmapi-launch %s'", originalCommand)
+	wrappedCommand := fmt.Sprintf("bash -c 'trtllm-llmapi-launch %s'", originalCommand)
 
 	// Generate environment variable flags for mpirun
 	envVarsStr := generateEnvVarFlags(container.Env)
 
-	mpirunCmd := fmt.Sprintf("mpirun --oversubscribe -n %d -H %s --mca pml ob1 --mca plm_rsh_args \"-p %d -o StrictHostKeyChecking=no -i ~/.ssh/id_rsa\" %s %s",
+	mpirunCmd := fmt.Sprintf("mpirun --allow-run-as-root --oversubscribe -n %d -H %s --mca pml ob1 --mca plm_rsh_args \"-p %d -o StrictHostKeyChecking=no -i ~/.ssh/id_rsa\" %s %s",
 		totalGPUs,
 		workerHosts,
 		commonconsts.MpiRunSshPort,