feat(api): add termination grace period to PodOverrides (#244)

Allows users to configure terminationGracePeriodSeconds via spec.server.podOverrides to prevent in-flight requests from being killed during HPA scale-down events.


Approved-by: VaishnaviHire

Approved-by: derekhiggins

Author: mfleader

api/v1alpha1/llamastackdistribution_types.go

@@ -401,6 +401,12 @@ spec:
                           ServiceAccountName allows users to specify their own ServiceAccount
                           If not specified, the operator will use the default ServiceAccount
                         type: string
+                      terminationGracePeriodSeconds:
+                        description: |-
+                          TerminationGracePeriodSeconds is the time allowed for graceful pod shutdown.
+                          If not specified, Kubernetes defaults to 30 seconds.
+                        format: int64
+                        type: integer
                       volumeMounts:
                         items:
                           description: VolumeMount describes a mounting of a Volume

api/v1alpha1/zz_generated.deepcopy.go

@@ -532,6 +532,11 @@ func configurePodOverrides(instance *llamav1alpha1.LlamaStackDistribution, podSp
 				podSpec.Containers[0].VolumeMounts = append(podSpec.Containers[0].VolumeMounts, instance.Spec.Server.PodOverrides.VolumeMounts...)
 			}
 		}
+
+		// Apply termination grace period if specified
+		if instance.Spec.Server.PodOverrides.TerminationGracePeriodSeconds != nil {
+			podSpec.TerminationGracePeriodSeconds = instance.Spec.Server.PodOverrides.TerminationGracePeriodSeconds
+		}
 	}
 }
 

config/crd/bases/llamastack.io_llamastackdistributions.yaml

@@ -650,6 +650,108 @@ func TestPodOverridesWithoutServiceAccount(t *testing.T) {
 	}
 }
 
+func TestPodOverridesWithTerminationGracePeriod(t *testing.T) {
+	gracePeriod := int64(120)
+	instance := &llamav1alpha1.LlamaStackDistribution{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "test-instance",
+			Namespace: "test-namespace",
+		},
+		Spec: llamav1alpha1.LlamaStackDistributionSpec{
+			Server: llamav1alpha1.ServerSpec{
+				PodOverrides: &llamav1alpha1.PodOverrides{
+					TerminationGracePeriodSeconds: &gracePeriod,
+				},
+			},
+		},
+	}
+	deployment := &appsv1.Deployment{
+		Spec: appsv1.DeploymentSpec{
+			Template: corev1.PodTemplateSpec{
+				Spec: corev1.PodSpec{
+					Containers: []corev1.Container{
+						{
+							Name: "test-container",
+						},
+					},
+				},
+			},
+		},
+	}
+
+	configurePodOverrides(instance, &deployment.Spec.Template.Spec)
+
+	require.NotNil(t, deployment.Spec.Template.Spec.TerminationGracePeriodSeconds)
+	assert.Equal(t, int64(120), *deployment.Spec.Template.Spec.TerminationGracePeriodSeconds)
+}
+
+func TestPodOverridesWithTerminationGracePeriodZero(t *testing.T) {
+	// Ensures we distinguish "not set" (nil) from "set to 0" (immediate termination)
+	gracePeriod := int64(0)
+	instance := &llamav1alpha1.LlamaStackDistribution{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "test-instance",
+			Namespace: "test-namespace",
+		},
+		Spec: llamav1alpha1.LlamaStackDistributionSpec{
+			Server: llamav1alpha1.ServerSpec{
+				PodOverrides: &llamav1alpha1.PodOverrides{
+					TerminationGracePeriodSeconds: &gracePeriod,
+				},
+			},
+		},
+	}
+	deployment := &appsv1.Deployment{
+		Spec: appsv1.DeploymentSpec{
+			Template: corev1.PodTemplateSpec{
+				Spec: corev1.PodSpec{
+					Containers: []corev1.Container{
+						{
+							Name: "test-container",
+						},
+					},
+				},
+			},
+		},
+	}
+
+	configurePodOverrides(instance, &deployment.Spec.Template.Spec)
+
+	require.NotNil(t, deployment.Spec.Template.Spec.TerminationGracePeriodSeconds)
+	assert.Equal(t, int64(0), *deployment.Spec.Template.Spec.TerminationGracePeriodSeconds)
+}
+
+func TestPodOverridesWithoutTerminationGracePeriod(t *testing.T) {
+	instance := &llamav1alpha1.LlamaStackDistribution{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "test-instance",
+			Namespace: "test-namespace",
+		},
+		Spec: llamav1alpha1.LlamaStackDistributionSpec{
+			Server: llamav1alpha1.ServerSpec{
+				PodOverrides: &llamav1alpha1.PodOverrides{},
+			},
+		},
+	}
+	deployment := &appsv1.Deployment{
+		Spec: appsv1.DeploymentSpec{
+			Template: corev1.PodTemplateSpec{
+				Spec: corev1.PodSpec{
+					Containers: []corev1.Container{
+						{
+							Name: "test-container",
+						},
+					},
+				},
+			},
+		},
+	}
+
+	configurePodOverrides(instance, &deployment.Spec.Template.Spec)
+
+	assert.Nil(t, deployment.Spec.Template.Spec.TerminationGracePeriodSeconds)
+}
+
 func TestValidateConfigMapKeys(t *testing.T) {
 	tests := []struct {
 		name        string

controllers/resource_helper.go

@@ -203,6 +203,7 @@ _Appears in:_
 | Field | Description | Default | Validation |
 | --- | --- | --- | --- |
 | `serviceAccountName` _string_ | ServiceAccountName allows users to specify their own ServiceAccount<br />If not specified, the operator will use the default ServiceAccount |  |  |
+| `terminationGracePeriodSeconds` _integer_ | TerminationGracePeriodSeconds is the time allowed for graceful pod shutdown.<br />If not specified, Kubernetes defaults to 30 seconds. |  |  |
 | `volumes` _[Volume](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.31/#volume-v1-core) array_ |  |  |  |
 | `volumeMounts` _[VolumeMount](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.31/#volumemount-v1-core) array_ |  |  |  |
 

controllers/resource_helper_test.go

@@ -410,6 +410,12 @@ spec:
                           ServiceAccountName allows users to specify their own ServiceAccount
                           If not specified, the operator will use the default ServiceAccount
                         type: string
+                      terminationGracePeriodSeconds:
+                        description: |-
+                          TerminationGracePeriodSeconds is the time allowed for graceful pod shutdown.
+                          If not specified, Kubernetes defaults to 30 seconds.
+                        format: int64
+                        type: integer
                       volumeMounts:
                         items:
                           description: VolumeMount describes a mounting of a Volume