NVIDIA 机密 NIM 部署

export NGC_API_KEY="<NGC_API_KEY>"

cat <<'EOF' | tee nvidia-nim-llama-3-1-8b-instruct.yaml >/dev/null
apiVersion: v1
kind: Secret
metadata:
  name: ngc-secret-instruct
type: kubernetes.io/dockerconfigjson
stringData:
  .dockerconfigjson: |
    {
      "auths": {
        "nvcr.io": {
          "username": "$oauthtoken",
          "password": "${NGC_API_KEY}"
        }
      }
    }
---
apiVersion: v1
kind: Secret
metadata:
  name: ngc-api-key-instruct
type: Opaque
stringData:
  api-key: "${NGC_API_KEY}"
---
apiVersion: v1
kind: Pod
metadata:
  name: nvidia-nim-llama-3-1-8b-instruct
  labels:
    app: nvidia-nim-llama-3-1-8b-instruct
spec:
  restartPolicy: Never
  runtimeClassName: kata-qemu-nvidia-gpu
  imagePullSecrets:
    - name: ngc-secret-instruct
  containers:
    - name: nvidia-nim-llama-3-1-8b-instruct
      image: nvcr.io/nim/meta/llama-3.1-8b-instruct:1.13.1
      ports:
        - containerPort: 8000
          name: http-openai
      livenessProbe:
        httpGet:
          path: /v1/health/live
          port: http-openai
        initialDelaySeconds: 15
        periodSeconds: 10
        timeoutSeconds: 1
        successThreshold: 1
        failureThreshold: 3
      readinessProbe:
        httpGet:
          path: /v1/health/ready
          port: http-openai
        initialDelaySeconds: 15
        periodSeconds: 10
        timeoutSeconds: 1
        successThreshold: 1
        failureThreshold: 3
      startupProbe:
        httpGet:
          path: /v1/health/ready
          port: http-openai
        initialDelaySeconds: 360
        periodSeconds: 10
        timeoutSeconds: 1
        successThreshold: 1
        failureThreshold: 120
      env:
        - name: NGC_API_KEY
          valueFrom:
            secretKeyRef:
              name: ngc-api-key-instruct
              key: api-key
      resources:
        limits:
          nvidia.com/pgpu: "1"
          cpu: "16"
          memory: "48Gi"
EOF

kubectl label nodes --all \
  nvidia.com/gpu.workload.config=vm-passthrough \
  --overwrite

kubectl label nodes --all nvidia.com/cc.mode=off --overwrite

kubectl wait \
  --for=jsonpath='{.metadata.labels.nvidia\.com/cc\.mode\.state}'=off \
  node --all \
  --timeout=15m

kubectl -n gpu-operator wait \
  --for=condition=Available \
  deployment --all \
  --timeout=10m

kubectl -n gpu-operator rollout status \
  daemonset/nvidia-vfio-manager \
  --timeout=10m

kubectl -n gpu-operator rollout status \
  daemonset/nvidia-sandbox-validator \
  --timeout=10m

kubectl -n gpu-operator rollout status \
  daemonset/nvidia-kata-sandbox-device-plugin-daemonset \
  --timeout=10m

kubectl -n gpu-operator rollout status \
  daemonset/nvidia-cc-manager \
  --timeout=10m

kubectl get nodes \
  -L nvidia.com/gpu.workload.config,nvidia.com/cc.mode,nvidia.com/cc.mode.state

: "${NGC_API_KEY:?set NGC_API_KEY to your NGC API key}"

envsubst '${NGC_API_KEY}' \
  < nvidia-nim-llama-3-1-8b-instruct.yaml \
  | kubectl apply -f -

kubectl wait --for=condition=Ready \
  --timeout=600s \
  pod/nvidia-nim-llama-3-1-8b-instruct

POD_IP="$(kubectl get pod nvidia-nim-llama-3-1-8b-instruct -o jsonpath='{.status.podIP}')"

curl -fsS "http://${POD_IP}:8000/v1/health/ready" | jq .
curl -fsS "http://${POD_IP}:8000/v1/models" | jq -r '.data[].id'

MODEL_NAME="$(curl -fsS "http://${POD_IP}:8000/v1/models" | jq -r '.data[0].id')"

curl -fsS "http://${POD_IP}:8000/v1/chat/completions" \
  -H 'Content-Type: application/json' \
  -d "$(jq -n --arg model "${MODEL_NAME}" '{
    model: $model,
    messages: [{role: "user", content: "Reply with exactly: hello from nim"}],
    max_tokens: 16,
    temperature: 0
  }')" | jq -r '.choices[0].message.content'

kubectl delete pod nvidia-nim-llama-3-1-8b-instruct --ignore-not-found
kubectl delete secret ngc-secret-instruct ngc-api-key-instruct \
  --ignore-not-found

kubectl label nodes --all nvidia.com/cc.mode=on --overwrite

# Local paths.
KBS_WORKDIR="${HOME}/nim-kbs"
KBS_TRUSTEE_DIR="${KBS_WORKDIR}/trustee"
KBS_CONFIG_DIR="${KBS_TRUSTEE_DIR}/kbs/config"
KBS_COMPOSE_CONFIG_DIR="${KBS_CONFIG_DIR}/docker-compose"
KBS_DATA_DIR="${KBS_TRUSTEE_DIR}/kbs/data"
KBS_STORAGE_DIR="${KBS_DATA_DIR}/kbs-storage"

# Artifact selection.
KBS_CLIENT_ARCH="${KBS_CLIENT_ARCH:-x86_64}"

# Network endpoint.
KBS_PORT="${KBS_PORT:-8080}"
KBS_NODE_IP="$(kubectl get nodes \
  -o jsonpath='{.items[0].status.addresses[?(@.type=="InternalIP")].address}')"
KBS_URL="https://${KBS_NODE_IP}:${KBS_PORT}"

mkdir -p "${KBS_WORKDIR}"

KATA_VERSIONS_FILE="${KATA_VERSIONS_FILE:-/opt/kata/versions.yaml}"
KBS_TRUSTEE_REF="${KBS_TRUSTEE_REF:-}"
KBS_IMAGE_REPO="ghcr.io/confidential-containers/staged-images/kbs-grpc-as"
KBS_AS_IMAGE_REPO="ghcr.io/confidential-containers/staged-images/coco-as-grpc"
KBS_RVPS_IMAGE_REPO="ghcr.io/confidential-containers/staged-images/rvps"
KBS_CLIENT_IMAGE_REPO="ghcr.io/confidential-containers/staged-images/kbs-client"

if test -z "${KBS_TRUSTEE_REF}"; then
  if ! test -r "${KATA_VERSIONS_FILE}"; then
    echo "ERROR: cannot resolve KBS_TRUSTEE_REF because ${KATA_VERSIONS_FILE} is not readable." >&2
    echo "Look up the validated Trustee reference for your platform, set KBS_TRUSTEE_REF, and rerun this snippet." >&2
    exit 1
  fi

  KBS_TRUSTEE_REF="$(
    awk '
      $1 == "coco-trustee:" { in_trustee = 1; next }
      in_trustee && $1 == "version:" {
        gsub(/"/, "", $2)
        print $2
        exit
      }
    ' "${KATA_VERSIONS_FILE}"
  )"
fi

if test -z "${KBS_TRUSTEE_REF}"; then
  echo "ERROR: could not find the coco-trustee version in ${KATA_VERSIONS_FILE}." >&2
  echo "Look up the validated Trustee reference for your platform, set KBS_TRUSTEE_REF, and rerun this snippet." >&2
  exit 1
fi

KBS_COMPOSE_IMAGE_TAG="${KBS_COMPOSE_IMAGE_TAG:-${KBS_TRUSTEE_REF}-${KBS_CLIENT_ARCH}}"
KBS_CLIENT_ARTIFACT="${KBS_CLIENT_IMAGE_REPO}:sample_only-${KBS_COMPOSE_IMAGE_TAG}"
KBS_TRUSTEE_TARBALL_URL="${KBS_TRUSTEE_TARBALL_URL:-https://github.com/confidential-containers/trustee/archive/${KBS_TRUSTEE_REF}.tar.gz}"

skopeo inspect "docker://${KBS_IMAGE_REPO}:${KBS_COMPOSE_IMAGE_TAG}" >/dev/null
skopeo inspect "docker://${KBS_AS_IMAGE_REPO}:${KBS_COMPOSE_IMAGE_TAG}" >/dev/null
skopeo inspect "docker://${KBS_RVPS_IMAGE_REPO}:${KBS_COMPOSE_IMAGE_TAG}" >/dev/null
oras manifest fetch "${KBS_CLIENT_ARTIFACT}" >/dev/null

printf 'KBS_TRUSTEE_REF=%s\nKBS_IMAGE=%s:%s\nAS_IMAGE=%s:%s\nRVPS_IMAGE=%s:%s\nKBS_CLIENT_ARTIFACT=%s\nKBS_URL=%s\n' \
  "${KBS_TRUSTEE_REF}" \
  "${KBS_IMAGE_REPO}" \
  "${KBS_COMPOSE_IMAGE_TAG}" \
  "${KBS_AS_IMAGE_REPO}" \
  "${KBS_COMPOSE_IMAGE_TAG}" \
  "${KBS_RVPS_IMAGE_REPO}" \
  "${KBS_COMPOSE_IMAGE_TAG}" \
  "${KBS_CLIENT_ARTIFACT}" \
  "${KBS_URL}"

rm -rf "${KBS_TRUSTEE_DIR}"
mkdir -p "${KBS_TRUSTEE_DIR}"

curl -fsSL "${KBS_TRUSTEE_TARBALL_URL}" |
  tar -xz --strip-components=1 -C "${KBS_TRUSTEE_DIR}"

sed -i \
  -e "s#image: ghcr.io/confidential-containers/staged-images/kbs-grpc-as:latest#image: ${KBS_IMAGE_REPO}:${KBS_COMPOSE_IMAGE_TAG}#" \
  -e "s#image: ghcr.io/confidential-containers/staged-images/coco-as-grpc:latest#image: ${KBS_AS_IMAGE_REPO}:${KBS_COMPOSE_IMAGE_TAG}#" \
  -e "s#image: ghcr.io/confidential-containers/staged-images/rvps:latest#image: ${KBS_RVPS_IMAGE_REPO}:${KBS_COMPOSE_IMAGE_TAG}#" \
  -e "s#\"8080:8080\"#\"${KBS_PORT}:8080\"#" \
  "${KBS_TRUSTEE_DIR}/docker-compose.yml"

openssl req -x509 \
  -out "${KBS_CONFIG_DIR}/kbs-https.crt" \
  -keyout "${KBS_CONFIG_DIR}/kbs-https.key" \
  -newkey rsa:2048 \
  -nodes \
  -sha256 \
  -days 365 \
  -subj "/CN=${KBS_NODE_IP}" \
  -addext "subjectAltName=IP:${KBS_NODE_IP},IP:127.0.0.1,DNS:localhost" \
  -addext "basicConstraints=CA:FALSE"

sed -i \
  -e 's#insecure_http = true#insecure_http = false\
private_key = "/opt/confidential-containers/kbs/user-keys/kbs-https.key"\
certificate = "/opt/confidential-containers/kbs/user-keys/kbs-https.crt"#' \
  "${KBS_COMPOSE_CONFIG_DIR}/kbs-config.toml"

grep -q 'insecure_http = false' "${KBS_COMPOSE_CONFIG_DIR}/kbs-config.toml"
grep -q 'type = "Simple"' "${KBS_COMPOSE_CONFIG_DIR}/kbs-config.toml"
jq -e '.verifier_config.nvidia_verifier.type == "Remote"' \
  "${KBS_CONFIG_DIR}/as-config.json" >/dev/null

docker compose \
  --project-directory "${KBS_TRUSTEE_DIR}" \
  --project-name nim-trustee \
  pull kbs as rvps setup

docker compose \
  --project-directory "${KBS_TRUSTEE_DIR}" \
  --project-name nim-trustee \
  up -d --no-build

docker compose \
  --project-directory "${KBS_TRUSTEE_DIR}" \
  --project-name nim-trustee \
  ps

KBS_ADMIN_TOKEN_FILE="${KBS_CONFIG_DIR}/admin-token"
KBS_CERT_FILE="${KBS_CONFIG_DIR}/kbs-https.crt"
test -s "${KBS_ADMIN_TOKEN_FILE}"
test -s "${KBS_CERT_FILE}"

command -v oras >/dev/null

mkdir -p "${KBS_WORKDIR}/bin"

oras pull \
  --output "${KBS_WORKDIR}/bin" \
  "${KBS_CLIENT_ARTIFACT}"

chmod +x "${KBS_WORKDIR}/bin/kbs-client"
KBS_CLIENT="${KBS_WORKDIR}/bin/kbs-client"

"${KBS_CLIENT}" --version

printf 'checkpoint2-probe\n' > "${KBS_WORKDIR}/checkpoint2-probe.txt"

"${KBS_CLIENT}" \
  --cert-file "${KBS_CERT_FILE}" \
  --url "${KBS_URL}" \
  config \
  --admin-token-file "${KBS_ADMIN_TOKEN_FILE}" \
  set-resource \
  --path default/checkpoint2/probe \
  --resource-file "${KBS_WORKDIR}/checkpoint2-probe.txt"

"${KBS_CLIENT}" \
  --cert-file "${KBS_CERT_FILE}" \
  --url "${KBS_URL}" \
  config \
  --admin-token-file "${KBS_ADMIN_TOKEN_FILE}" \
  delete-resource \
  --path default/checkpoint2/probe

KBS_WORKDIR="${KBS_WORKDIR}" python3 - <<'EOF'
import os
from jwcrypto import jwk

workdir = os.environ["KBS_WORKDIR"]
k = jwk.JWK.generate(
    kty='EC', crv='P-256', alg='ES256',
    use='sig', kid='sealed-secret-nim-key')
open(f'{workdir}/signing-key-private.jwk', 'w').write(k.export_private())
open(f'{workdir}/signing-key-public.jwk', 'w').write(k.export_public())
EOF

KBS_WORKDIR="${KBS_WORKDIR}" python3 - <<'EOF'
import base64
import json
import os
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import ec, utils

workdir = os.environ["KBS_WORKDIR"]
signing_jwk = json.loads(open(f"{workdir}/signing-key-private.jwk").read())

def b64url_decode(value):
    return base64.urlsafe_b64decode(value + "=" * (-len(value) % 4))

def b64url_encode(value):
    return base64.urlsafe_b64encode(value).rstrip(b"=").decode()

payload = {
    "version": "0.1.0",
    "type": "vault",
    "name": "kbs:///default/ngc-api-key/instruct",
    "provider": "kbs",
    "provider_settings": {},
    "annotations": {},
}

protected = {
    "b64": True,
    "alg": "ES256",
    "kid": "kbs:///default/signing-key/sealed-secret",
}

private_value = int.from_bytes(b64url_decode(signing_jwk["d"]), "big")
private_key = ec.derive_private_key(private_value, ec.SECP256R1())

protected_b64 = b64url_encode(json.dumps(
    protected, separators=(",", ":")).encode())
payload_b64 = b64url_encode(json.dumps(payload, separators=(",", ":")).encode())
signing_input = f"{protected_b64}.{payload_b64}".encode()
signature_der = private_key.sign(signing_input, ec.ECDSA(hashes.SHA256()))
r, s = utils.decode_dss_signature(signature_der)
signature = r.to_bytes(32, "big") + s.to_bytes(32, "big")

with open(f"{workdir}/ngc-api-key-instruct.sealed", "w") as f:
    f.write(f"sealed.{protected_b64}.{payload_b64}.{b64url_encode(signature)}")
EOF

NIM_TEE_MANIFEST="${KBS_WORKDIR}/nvidia-nim-llama-3-1-8b-instruct-tee.yaml"

SEALED_NGC_API_KEY_BASE64="$(
  base64 -w0 "${KBS_WORKDIR}/ngc-api-key-instruct.sealed"
)"

cat > "${NIM_TEE_MANIFEST}" <<EOF
apiVersion: v1
kind: Secret
metadata:
  name: ngc-api-key-sealed-instruct
type: Opaque
data:
  api-key: "${SEALED_NGC_API_KEY_BASE64}"
---
apiVersion: v1
kind: Pod
metadata:
  name: nvidia-nim-llama-3-1-8b-instruct
  labels:
    app: nvidia-nim-llama-3-1-8b-instruct
spec:
  securityContext:
    runAsUser: 1000
    runAsGroup: 1000
    fsGroup: 1000
    supplementalGroups: [4, 20, 24, 25, 27, 29, 30, 44, 46]
  restartPolicy: Never
  runtimeClassName: kata-qemu-nvidia-gpu-snp
  imagePullSecrets:
    - name: ngc-secret-instruct
  containers:
    - name: nvidia-nim-llama-3-1-8b-instruct
      image: nvcr.io/nim/meta/llama-3.1-8b-instruct:1.13.1
      ports:
        - containerPort: 8000
          name: http-openai
      livenessProbe:
        httpGet:
          path: /v1/health/live
          port: http-openai
        initialDelaySeconds: 15
        periodSeconds: 10
        timeoutSeconds: 1
        successThreshold: 1
        failureThreshold: 3
      readinessProbe:
        httpGet:
          path: /v1/health/ready
          port: http-openai
        initialDelaySeconds: 15
        periodSeconds: 10
        timeoutSeconds: 1
        successThreshold: 1
        failureThreshold: 3
      startupProbe:
        httpGet:
          path: /v1/health/ready
          port: http-openai
        initialDelaySeconds: 360
        periodSeconds: 10
        timeoutSeconds: 1
        successThreshold: 1
        failureThreshold: 120
      env:
        - name: NGC_API_KEY
          valueFrom:
            secretKeyRef:
              name: ngc-api-key-sealed-instruct
              key: api-key
      resources:
        limits:
          nvidia.com/pgpu: "1"
          cpu: "16"
          memory: "48Gi"
      volumeMounts:
        - name: nim-trusted-cache
          mountPath: /opt/nim/.cache
      volumeDevices:
        - name: trusted-storage
          devicePath: /dev/trusted_store
  volumes:
    - name: nim-trusted-cache
      emptyDir:
        sizeLimit: 64Gi
    - name: trusted-storage
      persistentVolumeClaim:
        claimName: trusted-pvc-instruct
EOF

kubectl apply --dry-run=server -f "${NIM_TEE_MANIFEST}"

INITDATA_FILE="${KBS_WORKDIR}/nim-initdata.toml"

{
  cat <<EOF
version = "0.1.0"
algorithm = "sha256"

[data]
"aa.toml" = """
[token_configs]
[token_configs.kbs]
url = "${KBS_URL}"
cert = '''
EOF
  cat "${KBS_CERT_FILE}"
  cat <<EOF
'''
"""

"cdh.toml" = """
[kbc]
name = "cc_kbc"
url = "${KBS_URL}"
kbs_cert = '''
EOF
  cat "${KBS_CERT_FILE}"
  cat <<'EOF'
'''

[image]
max_concurrent_layer_downloads_per_image = 1
authenticated_registry_credentials_uri = "kbs:///default/credentials/nvcr"
image_security_policy_uri = "kbs:///default/security-policy/nim"
"""
EOF
} > "${INITDATA_FILE}"

GENPOLICY_WORKDIR="${KBS_WORKDIR}/genpolicy"
NIM_TEE_MANIFEST="${NIM_TEE_MANIFEST:-${KBS_WORKDIR}/nvidia-nim-llama-3-1-8b-instruct-tee.yaml}"
INITDATA_FILE="${INITDATA_FILE:-${KBS_WORKDIR}/nim-initdata.toml}"
NIM_POLICY_MANIFEST="${GENPOLICY_WORKDIR}/nvidia-nim-llama-3-1-8b-instruct-tee-policy.yaml"
GENPOLICY_INITDATA="${GENPOLICY_WORKDIR}/nim-initdata.toml"
GENPOLICY_DOCKER_CONFIG="${GENPOLICY_WORKDIR}/docker-config"
KATA_RUNTIME_VERSION="$(/opt/kata/bin/kata-runtime --version \
  | awk '/kata-runtime/ {print $3}')"
KATA_TOOLS_ARCH="amd64"
KATA_TOOLS_NAME="kata-tools-static-${KATA_RUNTIME_VERSION}-${KATA_TOOLS_ARCH}"
KATA_TOOLS_TARBALL="${KATA_TOOLS_NAME}.tar.zst"
KATA_RELEASE_BASE="https://github.com/kata-containers/kata-containers/releases/download"
KATA_TOOLS_URL="${KATA_RELEASE_BASE}/${KATA_RUNTIME_VERSION}/${KATA_TOOLS_TARBALL}"
KATA_TOOLS_EXTRACT_DIR="${GENPOLICY_WORKDIR}/kata-tools"
KATA_TOOLS_DIR="${KATA_TOOLS_EXTRACT_DIR}/opt/kata"
GENPOLICY_BIN="${KATA_TOOLS_DIR}/bin/genpolicy"
GENPOLICY_DEFAULTS="${KATA_TOOLS_DIR}/share/defaults/kata-containers"

mkdir -p "${GENPOLICY_WORKDIR}/genpolicy-settings.d" \
  "${GENPOLICY_DOCKER_CONFIG}" \
  "${KATA_TOOLS_EXTRACT_DIR}"

curl -fL \
  -o "${GENPOLICY_WORKDIR}/${KATA_TOOLS_TARBALL}" \
  "${KATA_TOOLS_URL}"

tar --zstd \
  -xf "${GENPOLICY_WORKDIR}/${KATA_TOOLS_TARBALL}" \
  -C "${KATA_TOOLS_EXTRACT_DIR}"

"${GENPOLICY_BIN}" -v \
  -j "${GENPOLICY_DEFAULTS}"

: "${NGC_API_KEY:?set NGC_API_KEY to an NGC API key that can pull the NIM image}"
export NGC_API_KEY

GENPOLICY_AUTH="$(printf '$oauthtoken:%s' "${NGC_API_KEY}" | base64 -w0)"
jq -n \
  --arg username '$oauthtoken' \
  --arg password "${NGC_API_KEY}" \
  --arg auth "${GENPOLICY_AUTH}" \
  '{auths: {"nvcr.io": {username: $username, password: $password, auth: $auth}}}' \
  > "${GENPOLICY_DOCKER_CONFIG}/config.json"

cp "${GENPOLICY_DEFAULTS}/rules.rego" \
  "${GENPOLICY_WORKDIR}/rules.rego"

cp "${GENPOLICY_DEFAULTS}/genpolicy-settings.json" \
  "${GENPOLICY_WORKDIR}/genpolicy-settings.json"

cp "${GENPOLICY_DEFAULTS}/drop-in-examples/20-oci-1.3.0-drop-in.json" \
  "${GENPOLICY_WORKDIR}/genpolicy-settings.d/20-oci-1.3.0-drop-in.json"

cp "${NIM_TEE_MANIFEST}" "${NIM_POLICY_MANIFEST}"
cp "${INITDATA_FILE}" "${GENPOLICY_INITDATA}"

cat > "${GENPOLICY_WORKDIR}/genpolicy-settings.d/99-observation-read-stream.json" <<'EOF'
[
  {
    "op": "replace",
    "path": "/request_defaults/ReadStreamRequest",
    "value": true
  }
]
EOF

(
  cd "${GENPOLICY_WORKDIR}"

  DOCKER_CONFIG="${GENPOLICY_DOCKER_CONFIG}" \
  RUST_LOG=info \
  "${GENPOLICY_BIN}" \
    -u \
    -y "${NIM_POLICY_MANIFEST}" \
    -p "${GENPOLICY_WORKDIR}/rules.rego" \
    -j "${GENPOLICY_WORKDIR}" \
    --initdata-path="${GENPOLICY_INITDATA}"
)

kubectl apply --dry-run=server -f "${NIM_POLICY_MANIFEST}"

kubectl create --dry-run=client \
  -f "${NIM_POLICY_MANIFEST}" \
  -o json \
  | jq -r '
      select(.kind == "Pod")
      | .metadata.annotations["io.katacontainers.config.hypervisor.cc_init_data"]
    ' \
  | base64 -d \
  | gzip -d > "${GENPOLICY_WORKDIR}/generated-initdata.toml"

grep -n '^"aa\.toml"[[:space:]]*=' "${GENPOLICY_WORKDIR}/generated-initdata.toml"
grep -n '^"cdh\.toml"[[:space:]]*=' "${GENPOLICY_WORKDIR}/generated-initdata.toml"
grep -n '^"policy\.rego"[[:space:]]*=' "${GENPOLICY_WORKDIR}/generated-initdata.toml"

kubectl get nodes \
  -L nvidia.com/cc.mode,nvidia.com/cc.mode.state,nvidia.com/cc.ready.state

MEASUREMENT_POD_NAME="nim-snp-measurement"
MEASUREMENT_POD_MANIFEST="${KBS_WORKDIR}/nim-snp-measurement-pod.yaml"
NODE_NAME="$(
  kubectl get nodes \
    -l nvidia.com/cc.mode=on,nvidia.com/cc.mode.state=on,nvidia.com/cc.ready.state=true \
    -o jsonpath='{.items[0].metadata.name}'
)"

CC_INIT_DATA="$(
  kubectl create --dry-run=client \
    -f "${NIM_POLICY_MANIFEST}" \
    -o json \
    | jq -r '
        select(.kind == "Pod")
        | .metadata.annotations["io.katacontainers.config.hypervisor.cc_init_data"]
      '
)"

test -n "${NODE_NAME}"
test -n "${CC_INIT_DATA}"
test "${CC_INIT_DATA}" != "null"

cat > "${MEASUREMENT_POD_MANIFEST}" <<EOF
apiVersion: v1
kind: Pod
metadata:
  name: ${MEASUREMENT_POD_NAME}
  annotations:
    io.katacontainers.config.hypervisor.cc_init_data: "${CC_INIT_DATA}"
spec:
  restartPolicy: Never
  runtimeClassName: kata-qemu-nvidia-gpu-snp
  nodeName: ${NODE_NAME}
  containers:
    - name: sleep
      image: quay.io/prometheus/busybox:latest
      imagePullPolicy: IfNotPresent
      command: ["sh", "-c", "sleep 600"]
      resources:
        limits:
          nvidia.com/pgpu: "1"
          cpu: "16"
          memory: "48Gi"
EOF

kubectl delete pod "${MEASUREMENT_POD_NAME}" --ignore-not-found
kubectl apply -f "${MEASUREMENT_POD_MANIFEST}"

export PATH="${HOME}/.local/bin:${HOME}/.cargo/bin:${PATH}"

python3 -m pip install --user sev-snp-measure
cargo install --locked snphost

command -v sev-snp-measure
sev-snp-measure --help >/dev/null
command -v snphost
snphost --help >/dev/null

POD_UID="$(
  kubectl get pod "${MEASUREMENT_POD_NAME}" \
    -o jsonpath='{.metadata.uid}'
)"
POD_UID_UNDERSCORE="${POD_UID//-/_}"

QEMU_PID=""
for _ in $(seq 1 120); do
  QEMU_PID="$(
    for pid in $(pgrep -f 'qemu-system-' || true); do
      if sudo grep -Eq "pod${POD_UID_UNDERSCORE}|${POD_UID}" \
        "/proc/${pid}/cgroup" 2>/dev/null; then
        echo "${pid}"
        break
      fi
    done
  )"
  if test -n "${QEMU_PID}"; then
    break
  fi
  sleep 2
done

test -n "${QEMU_PID}"
echo "${QEMU_PID}"

PATH="${PATH}:${HOME}/.local/bin:${HOME}/.cargo/bin" \
QEMU_PID="${QEMU_PID}" \
KBS_WORKDIR="${KBS_WORKDIR}" \
python3 - <<'PY'
import json
import os
from pathlib import Path
import re
import subprocess

pid = os.environ["QEMU_PID"]
cmd = [
    item.decode()
    for item in Path(f"/proc/{pid}/cmdline").read_bytes().split(b"\0")
    if item
]

def value(flag):
    index = cmd.index(flag)
    return cmd[index + 1]

kernel = value("-kernel")
firmware = value("-bios")
append = value("-append").replace(r"\"", '"')
cpu_model = value("-cpu").split(",", 1)[0]
vcpu_count = re.match(r"\d+", value("-smp")).group(0)
initrd = value("-initrd") if "-initrd" in cmd else None
snp_object = next(
    cmd[index + 1]
    for index, item in enumerate(cmd)
    if item == "-object" and cmd[index + 1].startswith("sev-snp-guest")
)
host_data = re.search(r"(?:^|,)host-data=([^,]+)", snp_object).group(1)

out_dir = Path(os.environ["KBS_WORKDIR"])
inputs = {
    "qemu_pid": pid,
    "kernel": kernel,
    "firmware": firmware,
    "initrd": initrd,
    "cpu_model": cpu_model,
    "vcpu_count": int(vcpu_count),
    "append": append,
    "sev_snp_object": snp_object,
    "host_data": host_data,
}
(out_dir / "nim-qemu-launch-inputs.json").write_text(
    json.dumps(inputs, indent=2) + "\n"
)

measure_args = [
    "sev-snp-measure",
    "--mode=snp",
    f"--vcpus={vcpu_count}",
    f"--vcpu-type={cpu_model}",
    "--output-format=hex",
    f"--ovmf={firmware}",
    f"--kernel={kernel}",
    f"--append={append}",
]
if initrd:
    measure_args.append(f"--initrd={initrd}")

measurement = subprocess.check_output(measure_args, text=True).strip()
(out_dir / "nim-snp-launch-measurement.txt").write_text(measurement + "\n")
print(measurement)
print(f"SNP_HOST_DATA={host_data}")
PY

sudo snphost show tcb | tee "${KBS_WORKDIR}/nim-snp-tcb.txt"

awk '
  /Reported TCB/ { reported = 1; next }
  /Platform TCB/ { reported = 0 }
  reported && /Microcode:/ { print "SNP_MICROCODE=" $2 }
  reported && /SNP:/ { print "SNP_SNP_SVN=" $2 }
  reported && /TEE:/ { print "SNP_TEE_SVN=" $2 }
  reported && /Boot Loader:/ { print "SNP_BOOTLOADER=" $3 }
' "${KBS_WORKDIR}/nim-snp-tcb.txt" > "${KBS_WORKDIR}/nim-snp-tcb.env"

{
  SNP_HOST_DATA_BASE64="$(
    jq -r .host_data "${KBS_WORKDIR}/nim-qemu-launch-inputs.json"
  )"
  SNP_HOST_DATA_HEX="$(
    printf '%s' "${SNP_HOST_DATA_BASE64}" \
      | base64 -d \
      | od -An -tx1 -v \
      | tr -d ' \n'
  )"

  printf 'SNP_LAUNCH_MEASUREMENT=%s\n' "$(
    cat "${KBS_WORKDIR}/nim-snp-launch-measurement.txt"
  )"
  printf 'SNP_HOST_DATA_BASE64=%s\n' "${SNP_HOST_DATA_BASE64}"
  printf 'SNP_HOST_DATA_HEX=%s\n' "${SNP_HOST_DATA_HEX}"
  cat "${KBS_WORKDIR}/nim-snp-tcb.env"
} > "${KBS_WORKDIR}/nim-snp-reference-values.env"

cat "${KBS_WORKDIR}/nim-snp-reference-values.env"

kubectl get pod "${MEASUREMENT_POD_NAME}" \
  -o jsonpath='{.metadata.annotations.io\.katacontainers\.config\.hypervisor\.cc_init_data}' \
  | base64 -d \
  | gzip -d > "${KBS_WORKDIR}/live-generated-initdata.toml"

openssl dgst -sha256 -binary "${KBS_WORKDIR}/live-generated-initdata.toml" \
  | base64 -w0
echo

jq -r .host_data "${KBS_WORKDIR}/nim-qemu-launch-inputs.json"

kubectl delete pod "${MEASUREMENT_POD_NAME}" --ignore-not-found

set -a
. "${KBS_WORKDIR}/nim-snp-reference-values.env"
set +a

export NGC_API_KEY="<NGC_API_KEY>"

AUTH_VALUE="$(printf '$oauthtoken:%s' "${NGC_API_KEY}" | base64 -w0)"

jq -n --arg auth "${AUTH_VALUE}" '{
  auths: {
    "nvcr.io": {
      auth: $auth
    }
  }
}' > "${KBS_WORKDIR}/nvcr-auth.json"

printf '%s' "${NGC_API_KEY}" > "${KBS_WORKDIR}/ngc-api-key-instruct"

cat > "${KBS_WORKDIR}/nim-image-policy.json" <<'EOF'
{
  "default": [
    {
      "type": "reject"
    }
  ],
  "transports": {
    "docker": {
      "nvcr.io/nim/meta": [
        {
          "type": "sigstoreSigned",
          "keyPath": "kbs:///default/cosign-public-key/nim",
          "signedIdentity": {
            "type": "matchRepository"
          }
        }
      ],
      "nvcr.io/nim/nvidia": [
        {
          "type": "sigstoreSigned",
          "keyPath": "kbs:///default/cosign-public-key/nim",
          "signedIdentity": {
            "type": "matchRepository"
          }
        }
      ]
    }
  }
}
EOF

curl -fsSL \
  -o "${KBS_WORKDIR}/nvidia-cosign.pub" \
  https://api.ngc.nvidia.com/v2/catalog/containers/public-key

"${KBS_CLIENT}" \
  --cert-file "${KBS_CERT_FILE}" \
  --url "${KBS_URL}" \
  config \
  --admin-token-file "${KBS_ADMIN_TOKEN_FILE}" \
  set-resource \
  --path default/credentials/nvcr \
  --resource-file "${KBS_WORKDIR}/nvcr-auth.json" \
  >/dev/null

"${KBS_CLIENT}" \
  --cert-file "${KBS_CERT_FILE}" \
  --url "${KBS_URL}" \
  config \
  --admin-token-file "${KBS_ADMIN_TOKEN_FILE}" \
  set-resource \
  --path default/ngc-api-key/instruct \
  --resource-file "${KBS_WORKDIR}/ngc-api-key-instruct" \
  >/dev/null

"${KBS_CLIENT}" \
  --cert-file "${KBS_CERT_FILE}" \
  --url "${KBS_URL}" \
  config \
  --admin-token-file "${KBS_ADMIN_TOKEN_FILE}" \
  set-resource \
  --path default/security-policy/nim \
  --resource-file "${KBS_WORKDIR}/nim-image-policy.json"

"${KBS_CLIENT}" \
  --cert-file "${KBS_CERT_FILE}" \
  --url "${KBS_URL}" \
  config \
  --admin-token-file "${KBS_ADMIN_TOKEN_FILE}" \
  set-resource \
  --path default/cosign-public-key/nim \
  --resource-file "${KBS_WORKDIR}/nvidia-cosign.pub"

"${KBS_CLIENT}" \
  --cert-file "${KBS_CERT_FILE}" \
  --url "${KBS_URL}" \
  config \
  --admin-token-file "${KBS_ADMIN_TOKEN_FILE}" \
  set-resource \
  --path default/signing-key/sealed-secret \
  --resource-file "${KBS_WORKDIR}/signing-key-public.jwk"

"${KBS_CLIENT}" \
  --cert-file "${KBS_CERT_FILE}" \
  --url "${KBS_URL}" \
  config \
  --admin-token-file "${KBS_ADMIN_TOKEN_FILE}" \
  set-sample-reference-value \
  snp_launch_measurement "${SNP_LAUNCH_MEASUREMENT}"

"${KBS_CLIENT}" \
  --cert-file "${KBS_CERT_FILE}" \
  --url "${KBS_URL}" \
  config \
  --admin-token-file "${KBS_ADMIN_TOKEN_FILE}" \
  set-sample-reference-value \
  --as-integer \
  snp_bootloader "${SNP_BOOTLOADER}"

"${KBS_CLIENT}" \
  --cert-file "${KBS_CERT_FILE}" \
  --url "${KBS_URL}" \
  config \
  --admin-token-file "${KBS_ADMIN_TOKEN_FILE}" \
  set-sample-reference-value \
  --as-integer \
  snp_microcode "${SNP_MICROCODE}"

"${KBS_CLIENT}" \
  --cert-file "${KBS_CERT_FILE}" \
  --url "${KBS_URL}" \
  config \
  --admin-token-file "${KBS_ADMIN_TOKEN_FILE}" \
  set-sample-reference-value \
  --as-integer \
  snp_snp_svn "${SNP_SNP_SVN}"

"${KBS_CLIENT}" \
  --cert-file "${KBS_CERT_FILE}" \
  --url "${KBS_URL}" \
  config \
  --admin-token-file "${KBS_ADMIN_TOKEN_FILE}" \
  set-sample-reference-value \
  --as-integer \
  snp_tee_svn "${SNP_TEE_SVN}"

cat > "${KBS_WORKDIR}/nim-kbs-resource-policy.rego" <<EOF
package policy
import rego.v1

default allow = false

cpu0 := input["submods"]["cpu0"]
gpu0 := input["submods"]["gpu0"]
annotated_evidence := cpu0["ear.veraison.annotated-evidence"]
expected_init_data := "${SNP_HOST_DATA_HEX}"

allow if {
    cpu0["ear.status"] == "affirming"
    gpu0["ear.status"] == "affirming"

    annotated_evidence["init_data"] == expected_init_data
}
EOF

"${KBS_CLIENT}" \
  --cert-file "${KBS_CERT_FILE}" \
  --url "${KBS_URL}" \
  config \
  --admin-token-file "${KBS_ADMIN_TOKEN_FILE}" \
  set-resource-policy \
  --policy-file "${KBS_WORKDIR}/nim-kbs-resource-policy.rego"

NODE_NAME="$(kubectl get nodes -o jsonpath='{.items[0].metadata.name}')"
LOOP_FILE="/tmp/trusted-image-storage-instruct.img"
STORAGE_SIZE_MIB="57344"

df -h "$(dirname "${LOOP_FILE}")"

# Match Kata's NVIDIA NIM CI setup: remove any stale loop device and fully
# allocate the backing file before the workload starts. A sparse tmpfs file
# created with truncate can charge pages to the pod cgroup as QEMU writes them.
if LOOP_DEVICE="$(sudo losetup -j "${LOOP_FILE}" | awk -F: 'NR == 1 {print $1}')"; \
   test -n "${LOOP_DEVICE}"; then
  sudo losetup --detach "${LOOP_DEVICE}"
fi

rm -f "${LOOP_FILE}"
dd if=/dev/zero of="${LOOP_FILE}" bs=1M count="${STORAGE_SIZE_MIB}" status=progress
LOOP_DEVICE="$(sudo losetup --find --show "${LOOP_FILE}")"

echo "NODE_NAME=${NODE_NAME}"
echo "LOOP_DEVICE=${LOOP_DEVICE}"

cat > trusted-storage-instruct.yaml <<EOF
apiVersion: storage.k8s.io/v1
kind: StorageClass
metadata:
  name: local-storage
provisioner: kubernetes.io/no-provisioner
volumeBindingMode: WaitForFirstConsumer
---
apiVersion: v1
kind: PersistentVolume
metadata:
  name: trusted-block-pv-instruct
spec:
  capacity:
    storage: 57344Mi
  volumeMode: Block
  accessModes:
    - ReadWriteOnce
  persistentVolumeReclaimPolicy: Retain
  storageClassName: local-storage
  local:
    path: ${LOOP_DEVICE}
  nodeAffinity:
    required:
      nodeSelectorTerms:
        - matchExpressions:
            - key: kubernetes.io/hostname
              operator: In
              values:
                - ${NODE_NAME}
---
apiVersion: v1
kind: PersistentVolumeClaim
metadata:
  name: trusted-pvc-instruct
spec:
  accessModes:
    - ReadWriteOnce
  resources:
    requests:
      storage: 57344Mi
  volumeMode: Block
  storageClassName: local-storage
EOF

kubectl apply -f trusted-storage-instruct.yaml

kubectl get pv trusted-block-pv-instruct
kubectl get pvc trusted-pvc-instruct

: "${NGC_API_KEY:?set NGC_API_KEY to an NGC API key}"
NIM_POLICY_MANIFEST="${NIM_POLICY_MANIFEST:-${KBS_WORKDIR}/genpolicy/nvidia-nim-llama-3-1-8b-instruct-tee-policy.yaml}"

kubectl create secret docker-registry ngc-secret-instruct \
  --docker-server=nvcr.io \
  --docker-username='$oauthtoken' \
  --docker-password="${NGC_API_KEY}" \
  --dry-run=client \
  -o yaml | kubectl apply -f -

kubectl apply -f "${NIM_POLICY_MANIFEST}"

kubectl get secret ngc-api-key-sealed-instruct
kubectl get secret ngc-api-key-sealed-instruct \
  -o jsonpath='{.data.api-key}' | base64 -d | cut -c1-7
echo

kubectl wait --for=condition=Ready \
  --timeout=1000s \
  pod/nvidia-nim-llama-3-1-8b-instruct

POD_IP="$(kubectl get pod nvidia-nim-llama-3-1-8b-instruct -o jsonpath='{.status.podIP}')"

curl -fsS "http://${POD_IP}:8000/v1/health/ready" | jq .
curl -fsS "http://${POD_IP}:8000/v1/models" | jq -r '.data[].id'

MODEL_NAME="$(curl -fsS "http://${POD_IP}:8000/v1/models" | jq -r '.data[0].id')"

curl -fsS "http://${POD_IP}:8000/v1/chat/completions" \
  -H 'Content-Type: application/json' \
  -d "$(jq -n --arg model "${MODEL_NAME}" '{
    model: $model,
    messages: [{role: "user", content: "Reply with exactly: hello from tee nim"}],
    max_tokens: 16,
    temperature: 0
  }')" | jq -r '.choices[0].message.content'

kubectl logs nvidia-nim-llama-3-1-8b-instruct

kubectl get pod nvidia-nim-llama-3-1-8b-instruct -o wide
kubectl describe pod nvidia-nim-llama-3-1-8b-instruct
kubectl get events \
  --field-selector involvedObject.name=nvidia-nim-llama-3-1-8b-instruct \
  --sort-by=.lastTimestamp

docker compose \
  --project-directory "${KBS_TRUSTEE_DIR}" \
  --project-name nim-trustee \
  ps

docker compose \
  --project-directory "${KBS_TRUSTEE_DIR}" \
  --project-name nim-trustee \
  logs --since 30m kbs as rvps

export RUST_LOG=info,kbs=debug,attestation_service=debug,reference_value_provider_service=debug,policy_engine=debug

export RUST_LOG=info,kbs=debug,attestation_service=debug,reference_value_provider_service=debug,policy_engine=debug

docker compose \
  --project-directory "${KBS_TRUSTEE_DIR}" \
  --project-name nim-trustee \
  up -d --no-build --force-recreate kbs as rvps

export RUST_LOG=info

docker compose \
  --project-directory "${KBS_TRUSTEE_DIR}" \
  --project-name nim-trustee \
  up -d --no-build --force-recreate kbs as rvps

KBS_WORKDIR="${KBS_WORKDIR:-${HOME}/nim-kbs}"
KBS_TRUSTEE_DIR="${KBS_TRUSTEE_DIR:-${KBS_WORKDIR}/trustee}"
GENPOLICY_WORKDIR="${GENPOLICY_WORKDIR:-${KBS_WORKDIR}/genpolicy}"
LOOP_FILE="${LOOP_FILE:-/tmp/trusted-image-storage-instruct.img}"

if test "${KBS_WORKDIR}" = "/" \
   || test "${KBS_TRUSTEE_DIR}" = "/" \
   || test "${GENPOLICY_WORKDIR}" = "/"; then
  echo "Refusing to remove /"
  exit 1
fi

if test -d "${KBS_TRUSTEE_DIR}"; then
  docker compose \
    --project-directory "${KBS_TRUSTEE_DIR}" \
    --project-name nim-trustee \
    down --remove-orphans
fi

kubectl delete pod nvidia-nim-llama-3-1-8b-instruct \
  --ignore-not-found
kubectl delete pod nim-snp-measurement --ignore-not-found

kubectl delete secret \
  ngc-secret-instruct \
  ngc-api-key-instruct \
  ngc-api-key-sealed-instruct \
  --ignore-not-found

kubectl delete pvc trusted-pvc-instruct --ignore-not-found
kubectl delete pv trusted-block-pv-instruct --ignore-not-found
kubectl delete storageclass local-storage --ignore-not-found

if LOOP_DEVICE="$(sudo losetup -j "${LOOP_FILE}" | awk -F: 'NR == 1 {print $1}')" \
   && test -n "${LOOP_DEVICE}"; then
  sudo losetup --detach "${LOOP_DEVICE}"
fi

rm -f "${LOOP_FILE}"

rm -rf "${GENPOLICY_WORKDIR}" || sudo rm -rf "${GENPOLICY_WORKDIR}"
rm -rf "${KBS_WORKDIR}" || sudo rm -rf "${KBS_WORKDIR}"
rm -f trusted-storage-instruct.yaml \
  nvidia-nim-llama-3-1-8b-instruct.yaml \
  nvidia-nim-llama-3-1-8b-instruct-tee.yaml

sudo ls -la \
  /opt/kata/share/defaults/kata-containers/runtimes/qemu-nvidia-gpu-snp/config.d