Add readiness probe to Server Deployment (#183)

Fixes #153
substratusai · Aug 11, 2023 · 546e219 · 546e219
1 parent 09c0b48
commit 546e219
Show file tree

Hide file tree

Showing 5 changed files with 27 additions and 18 deletions.
diff --git a/docs/container-contract.md b/docs/container-contract.md
@@ -10,7 +10,7 @@ The working directory MUST be set to `/content`.
 WORKDIR /content
 ```
 
-### Jupyter
+## Jupyter
 
 This requirement applies to Model, Dataset, and Notebook containers.
 
@@ -22,7 +22,7 @@ The `notebook.sh` script MUST be located in `$PATH`. It is recommended that this
 
 Note: This requirement is satisfied by default when using Substratus base images.
 
-### Directory Structure
+## Directory Structure
 
 ```
 /content/         # Working directory.
@@ -33,7 +33,7 @@ Note: This requirement is satisfied by default when using Substratus base images
   saved-model/    # Location where a previously saved model will be mounted.
 ```
 
-### Parameters
+## Parameters
 
 Substratus provides params as a file (`/content/params.json`) and as environment variables to containers.
 
@@ -48,3 +48,10 @@ spec:
 Parameters get converted to environment variables using the following scheme.
 
 `PARAM_{upper(param_key)}={param_value}`
+
+## Server
+
+Substratus Server containers are expected to:
+
+* Serve HTTP traffic on port `8080`.
+* Serve a 200 OK on the root path `/` when ready to serve traffic.
diff --git a/examples/facebook-opt-125m/base-server.yaml b/examples/facebook-opt-125m/base-server.yaml
@@ -3,9 +3,6 @@ kind: Server
 metadata:
   name: facebook-opt-125m
 spec:
-  build:
-    git:
-      url: https://github.com/substratusai/images
-      path: model-server-basaran
+  image: substratusai/model-server-basaran
   model:
     name: facebook-opt-125m
diff --git a/internal/controller/model_controller.go b/internal/controller/model_controller.go
@@ -297,10 +297,8 @@ func (r *ModelReconciler) modellerJob(ctx context.Context, model, baseModel *api
 					ServiceAccountName: modellerServiceAccountName,
 					Containers: []corev1.Container{
 						{
-							Name:  containerName,
-							Image: model.GetImage(),
-							// NOTE: tini should be installed as the ENTRYPOINT the image and will be used
-							// to execute this script.
+							Name:    containerName,
+							Image:   model.GetImage(),
 							Command: model.Spec.Command,
 							Env:     paramsToEnv(model.Spec.Params),
 						},

diff --git a/internal/controller/notebook_controller.go b/internal/controller/notebook_controller.go
@@ -336,10 +336,8 @@ func (r *NotebookReconciler) notebookPod(notebook *apiv1.Notebook, model *apiv1.
 			ServiceAccountName: notebookServiceAccountName,
 			Containers: []corev1.Container{
 				{
-					Name:  containerName,
-					Image: notebook.GetImage(),
-					// NOTE: tini should be installed as the ENTRYPOINT the image and will be used
-					// to execute this script.
+					Name:    containerName,
+					Image:   notebook.GetImage(),
 					Command: notebook.Spec.Command,
 					//WorkingDir: "/home/jovyan",
 					Ports: []corev1.ContainerPort{

diff --git a/internal/controller/server_controller.go b/internal/controller/server_controller.go
@@ -139,15 +139,24 @@ func (r *ServerReconciler) serverDeployment(server *apiv1.Server, model *apiv1.M
 							Name:            containerName,
 							Image:           server.GetImage(),
 							ImagePullPolicy: "Always",
-							// NOTE: tini should be installed as the ENTRYPOINT the image and will be used
-							// to execute this script.
-							Command: server.Spec.Command,
+							Command:         server.Spec.Command,
 							Ports: []corev1.ContainerPort{
 								{
 									Name:          modelServerHTTPServePortName,
 									ContainerPort: 8080,
 								},
 							},
+							ReadinessProbe: &corev1.Probe{
+								ProbeHandler: corev1.ProbeHandler{
+									HTTPGet: &corev1.HTTPGetAction{
+										// TBD: Not sure if this should ever be something we configure via
+										// the Server API. For now, we'll just hardcode it and add it to the container
+										// contract.
+										Path: "/",
+										Port: intstr.FromString(modelServerHTTPServePortName),
+									},
+								},
+							},
 						},
 					},
 				},