ai-ml/t5-model-serving/kubernetes/loadgenerator.yaml

# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

apiVersion: v1
kind: ConfigMap
metadata:
  name: loadgenerator-files
data:
  locustfile.py: |
    import os

    from locust import FastHttpUser, between, task


    class T5User(FastHttpUser):
        wait_time = between(2.0, 2.5)

        def on_start(self):
            model_name = os.getenv('MODEL_NAME', 't5-small')
            model_version = os.getenv('MODEL_VERSION', '1.0')

            self.infer_url = f'{self.environment.host}/predictions/{model_name}/{model_version}'
            self.payload = {
                "text": "this is a test sentence",
                "from": "en",
                "to": "fr"
            }

        @task()
        def t5(self):
            with self.rest('POST', self.infer_url, json=self.payload, stream=False) as response:
                if response.status_code == 200:
                    response.success()
                else:
                    response.failure(response.js)

---
apiVersion: v1
kind: ConfigMap
metadata:
  name: loadgenerator-env
data:
  LOCUST_HOST: http://t5-inference:8080
  LOCUST_USERS: "10000"
  LOCUST_SPAWN_RATE: "250"
  LOCUST_LOGLEVEL: error
  LOCUST_WEB_PORT: "8080"
  MODEL_NAME: "t5-small"
  MODEL_VERSION: "1.0"
---
apiVersion: apps/v1
kind: Deployment
metadata:
  name: loadgenerator
  labels:
    app: loadgenerator
spec:
  replicas: 1
  selector:
    matchLabels:
      app: loadgenerator
  template:
    metadata:
      labels:
        app: loadgenerator
    spec:
      volumes:
        - name: locustfile
          configMap:
            name: loadgenerator-files
      containers:
        - name: loadgenerator
          image: locustio/locust
          command: ["locust", "--autostart"]
          envFrom:
            - configMapRef:
                name: loadgenerator-env
          volumeMounts:
            - mountPath: /home/locust
              name: locustfile
          resources:
            requests:
              memory: "512Mi"
              cpu: "250m"
              ephemeral-storage: "1Gi"
            limits:
              memory: "512Mi"
              cpu: "250m"
              ephemeral-storage: "1Gi"
          readinessProbe:
            httpGet:
              scheme: HTTP
              path: /
              port: 8080
          ports:
            - containerPort: 8080
---
apiVersion: v1
kind: Service
metadata:
  name: loadgenerator
spec:
  type: ClusterIP
  selector:
    app: loadgenerator
  ports:
    - name: http
      port: 8080
      targetPort: 8080