Add inference client setup script for trusted space codelab

PiperOrigin-RevId: 719182154

codelabs/trusted_space_codelab/scripts/config_env.sh

@@ -21,4 +21,8 @@ export PRIMUS_PROJECT_REPOSITORY_REGION=${PRIMUS_PROJECT_REPOSITORY_REGION:-'us'
 export WORKLOAD_SERVICEACCOUNT=${WORKLOAD_SERVICEACCOUNT:-workload-sa}
 export WORKLOAD_IMAGE_NAME=${WORKLOAD_IMAGE_NAME:-workload-container}
 export WORKLOAD_IMAGE_TAG=${WORKLOAD_IMAGE_TAG:-latest}
-export WORKLOAD_VM=${WORKLOAD_VM:-workload-vm}
+export WORKLOAD_VM=${WORKLOAD_VM:-workload-vm}
+
+# Workload client related variables
+export CLIENT_VM=${CLIENT_VM:-client-vm}
+export CLIENT_SERVICEACCOUNT=${CLIENT_SERVICEACCOUNT:-client-sa}

codelabs/trusted_space_codelab/scripts/setup_client.sh

@@ -0,0 +1,57 @@
+#!/bin/bash
+#
+# Creates required cloud resources for workload client.
+
+source config_env.sh
+source common.sh
+
+PARENT_DIR=$(dirname "${PWD}")
+
+echo "Creating workload client's service-account ${CLIENT_SERVICEACCOUNT} ..."
+create_service_account "${CLIENT_SERVICEACCOUNT}"
+
+echo "Granting KMS decryptor role to the service-account ${CLIENT_SERVICEACCOUNT} ..."
+gcloud kms keys add-iam-policy-binding \
+  projects/"${PRIMUS_PROJECT_ID}"/locations/"${PRIMUS_PROJECT_LOCATION}"/keyRings/"${PRIMUS_ENC_KEYRING}"/cryptoKeys/"${PRIMUS_ENC_KEY}" \
+    --member=serviceAccount:"${CLIENT_SERVICEACCOUNT}"@"${PRIMUS_PROJECT_ID}".iam.gserviceaccount.com \
+    --role=roles/cloudkms.cryptoKeyDecrypter
+
+echo "Granting KMS encryptor role to the service-account ${CLIENT_SERVICEACCOUNT} ..."
+gcloud kms keys add-iam-policy-binding \
+  projects/"${PRIMUS_PROJECT_ID}"/locations/"${PRIMUS_PROJECT_LOCATION}"/keyRings/"${PRIMUS_ENC_KEYRING}"/cryptoKeys/"${PRIMUS_ENC_KEY}" \
+    --member=serviceAccount:"${CLIENT_SERVICEACCOUNT}"@"${PRIMUS_PROJECT_ID}".iam.gserviceaccount.com \
+    --role=roles/cloudkms.cryptoKeyEncrypter
+
+echo "Creating workload client VM ${CLIENT_VM} ..."
+gcloud compute instances create "${CLIENT_VM}" \
+  --image-family=ubuntu-2204-lts \
+  --image-project=ubuntu-os-cloud \
+  --zone="${PRIMUS_PROJECT_ZONE}" \
+  --boot-disk-size=100GB \
+  --scopes=cloud-platform \
+  --service-account=${CLIENT_SERVICEACCOUNT}@${PRIMUS_PROJECT_ID}.iam.gserviceaccount.com
+gcloud compute ssh "${CLIENT_VM}" --zone="${PRIMUS_PROJECT_ZONE}" --command="echo 'Client VM is ready'"
+
+echo "Updating client code with required resource names ..."
+cp "${PARENT_DIR}"/src/client/sample_inference_client.py "${PARENT_DIR}"/src/client/inference_client.py
+INFERENCE_SERVER_IP=$(gcloud compute instances describe "${WORKLOAD_VM}" --format='get(networkInterfaces[0].networkIP)' --zone="${PRIMUS_PROJECT_ZONE}")
+sed -i'' "s/INFERENCE_SERVER_IP_VALUE/"${INFERENCE_SERVER_IP}"/" "${PARENT_DIR}"/src/client/inference_client.py
+sed -i'' "s/PRIMUS_PROJECT_ID_VALUE/"${PRIMUS_PROJECT_ID}"/" "${PARENT_DIR}"/src/client/inference_client.py
+sed -i'' "s/PRIMUS_KEY_ID_VALUE/"${PRIMUS_ENC_KEY}"/" "${PARENT_DIR}"/src/client/inference_client.py
+sed -i'' "s/PRIMUS_KEYRING_VALUE/"${PRIMUS_ENC_KEYRING}"/" "${PARENT_DIR}"/src/client/inference_client.py
+sed -i'' "s/PRIMUS_PROJECT_LOCATION_VALUE/"${PRIMUS_PROJECT_LOCATION}"/" "${PARENT_DIR}"/src/client/inference_client.py
+
+echo "Copying client code to client VM ..."
+gcloud compute scp "${PARENT_DIR}"/src/client/inference_client.py "${CLIENT_VM}":~/ --zone="${PRIMUS_PROJECT_ZONE}"
+gcloud compute scp "${PARENT_DIR}"/src/client/requirements.txt "${CLIENT_VM}":~/ --zone="${PRIMUS_PROJECT_ZONE}"
+
+echo "Installing required dependencies for client ..."
+gcloud compute ssh "${CLIENT_VM}" --zone="${PRIMUS_PROJECT_ZONE}" --command="
+  sudo apt-get update
+  sudo apt-get install -y python3 python3-venv
+  python3 -m venv venv
+  source venv/bin/activate
+  pip install -r requirements.txt
+"
+
+echo "Client VM is created and setup is complete. You can now SSH into the client VM: gcloud compute ssh ${CLIENT_VM} --zone=${PRIMUS_PROJECT_ZONE}"

codelabs/trusted_space_codelab/src/client/requirements.txt

@@ -0,0 +1,3 @@
+pyjwt
+requests
+google-cloud-kms==2.21.1

codelabs/trusted_space_codelab/src/client/sample_inference_client.py

@@ -0,0 +1,44 @@
+import base64
+import os
+import sys
+from google.cloud import kms
+import requests
+
+os.environ["PRIMUS_PROJECT_ID"] = "PRIMUS_PROJECT_ID_VALUE"
+os.environ["PRIMUS_PROJECT_LOCATION"] = "PRIMUS_PROJECT_LOCATION_VALUE"
+os.environ["PRIMUS_KEY_ID"] = "PRIMUS_KEY_ID_VALUE"
+os.environ["PRIMUS_KEYRING"] = "PRIMUS_KEYRING_VALUE"
+os.environ["INFERENCE_SERVER_URL"] = "http://INFERENCE_SERVER_IP_VALUE:8080/generate"
+
+# Initialising KMS Client.
+kms_client = kms.KeyManagementServiceClient()
+key_name = kms_client.crypto_key_path(
+    os.environ["PRIMUS_PROJECT_ID"],
+    os.environ["PRIMUS_PROJECT_LOCATION"],
+    os.environ["PRIMUS_KEYRING"],
+    os.environ["PRIMUS_KEY_ID"],
+)
+
+
+def data_exchange():
+  text = input("Enter your prompt: ")
+  encrypted_prompt = kms_client.encrypt(name=key_name, plaintext=bytes(text, "utf-8"))
+  ciphertext = base64.b64encode(encrypted_prompt.ciphertext).decode("utf-8")
+  payload = {
+      "ciphertext": ciphertext,
+  }
+  print("sending encrypted payload: ", payload)
+  response = requests.post(os.environ["INFERENCE_SERVER_URL"], json=payload)
+  data = response.json()
+  print("received encrypted response", data)
+  ciphertext = base64.b64decode(data["generated_code_ciphertext"])
+  decrypted_response = kms_client.decrypt(name=key_name, ciphertext=ciphertext)
+  print("decrypted response: ", decrypted_response.plaintext)
+
+
+def main():
+  data_exchange()
+
+
+if __name__ == "__main__":
+  main()

codelabs/trusted_space_codelab/src/sample_inference_server.py

@@ -97,7 +97,7 @@ def generate():
     ciphertext = response.ciphertext
     ciphertext_base64 = base64.b64encode(ciphertext).decode("utf-8")
     response = {"generated_code_ciphertext": ciphertext_base64}
-    return response
+    return jsonify(response)
 
   except (ValueError, TypeError, KeyError) as e:
     return jsonify({"error": str(e)}), 500