# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
apiVersion: batch/v1
kind: Job
metadata:
  name: oss-gpt120b-bench
spec:
  backoffLimit: 1
  completions: 1
  parallelism: 1
  template:
    metadata:
      labels:
        app: oss-gpt120b
    spec:
      restartPolicy: Never
      containers:
      - name: perf
        image: my-registry/vllm-runtime:my-tag
        workingDir: /workspace/components/backends/vllm
        env:
          - name: TARGET_MODEL
            value: openai/gpt-oss-120b
          - name: ENDPOINT
            value: gpt-oss-agg-trtllmworker:8000
          - name: CONCURRENCIES
            value: "13000 13500 1400"
          - name: ISL
            value: "16"
          - name: OSL
            value: "1000"
          - name: DEPLOYMENT_MODE
            value: "agg"
          - name: DEPLOYMENT_GPU_COUNT
            value: "32"
          - name: JOB_NAME
            valueFrom:
              fieldRef:
                fieldPath: metadata.labels['job-name']
          - name: ROOT_ARTIFACT_DIR
            value: /root/.cache/huggingface/hub/perf
        command:
        - /bin/sh
        - -c
        - |
          #TODO: this can be baked into the aiperf image
          apt-get update && apt-get install -y curl jq
          export COLUMNS=200
          EPOCH=$(date +%s)
          ## utility functions -- can be moved to a bash script / configmap
          wait_for_model_ready() {
            echo "Waiting for model '$TARGET_MODEL' at $ENDPOINT/v1/models (checking every 5s)..."
            while ! curl -s "http://$ENDPOINT/v1/models" | jq -e --arg model "$TARGET_MODEL" '.data[]? | select(.id == $model)' >/dev/null 2>&1; do
                echo "[$(date '+%H:%M:%S')] Model not ready yet, waiting 5s..."
                sleep 5
            done
            echo "✅ Model '$TARGET_MODEL' is now available!"
            echo "Model '$TARGET_MODEL' is now available!"
            curl -s "http://$ENDPOINT/v1/models" | jq .
          }
          run_perf() {
            local concurrency=$1
            local isl=$2
            local osl=$3
            key=concurrency_${concurrency}
            export ARTIFACT_DIR="${ROOT_ARTIFACT_DIR}/${EPOCH}_${JOB_NAME}/${key}"
            mkdir -p "$ARTIFACT_DIR"
            aiperf profile --artifact-dir $ARTIFACT_DIR \
                --model $TARGET_MODEL \
                --tokenizer ~/.cache/huggingface/hub/models--openai--gpt-oss-120b/snapshots/b5c939de8f754692c1647ca79fbf85e8c1e70f8a  \
                --endpoint-type chat \
                --endpoint /v1/chat/completions \
                --streaming \
                --url http://$ENDPOINT \
                --synthetic-input-tokens-mean $isl \
                --synthetic-input-tokens-stddev 0 \
                --output-tokens-mean $osl \
                --output-tokens-stddev 0 \
                --extra-inputs "{\"max_tokens\":$osl}" \
                --extra-inputs "{\"min_tokens\":$osl}" \
                --extra-inputs "{\"ignore_eos\":true}" \
                --extra-inputs "{\"nvext\":{\"ignore_eos\":true}}" \
                --concurrency $concurrency \
                --request-count $((3*concurrency)) \
                --warmup-request-count $concurrency \
                --conversation-num 1 \
                --random-seed 100 \
                --request-rate 100000 \
                --workers-max 128 \
                -H 'Authorization: Bearer NOT USED' \
                -H 'Accept: text/event-stream'\
                --record-processors 32 \
                --ui simple
            echo "ARTIFACT_DIR: $ARTIFACT_DIR"
            ls -la $ARTIFACT_DIR
          }
          #### Actual execution ####
          wait_for_model_ready
          mkdir -p "${ROOT_ARTIFACT_DIR}/${EPOCH}_${JOB_NAME}"
          # Write input_config.json
          cat > "${ROOT_ARTIFACT_DIR}/${EPOCH}_${JOB_NAME}/input_config.json" <<EOF
          {
            "gpu_count": $DEPLOYMENT_GPU_COUNT,
            "mode": "$DEPLOYMENT_MODE",
            "isl": $ISL,
            "osl": $OSL,
            "endpoint": "$ENDPOINT",
            "model endpoint": "$TARGET_MODEL"
          }
          EOF
          # Run perf for each concurrency
          for concurrency in $CONCURRENCIES; do
            run_perf $concurrency $ISL $OSL
            sleep 10
          done
        volumeMounts:
        - name: model-cache
          mountPath: /root/.cache/huggingface
      imagePullSecrets:
        - name: nvcrimagepullsecret
      volumes:
      - name: model-cache
        persistentVolumeClaim:
          claimName: model-cache
