From ce8a0b65fa68590e8ca90c1df66e324d779e91ea Mon Sep 17 00:00:00 2001
From: Chandra Devarakonda <chandrasekhard@google.com>
Date: Fri, 11 Aug 2023 21:00:39 +0000
Subject: [PATCH] Fix wrapt issue

---
 tests/tensorflow/nightly-se/common.libsonnet | 2 ++
 tests/tensorflow/nightly/common.libsonnet    | 1 +
 2 files changed, 3 insertions(+)

diff --git a/tests/tensorflow/nightly-se/common.libsonnet b/tests/tensorflow/nightly-se/common.libsonnet
index fabf4734e..2c9b85f3a 100644
--- a/tests/tensorflow/nightly-se/common.libsonnet
+++ b/tests/tensorflow/nightly-se/common.libsonnet
@@ -176,6 +176,8 @@ local volumes = import 'templates/volumes.libsonnet';
               sleep %(sleepTime)d
 
               softwareVersion=%(softwareVersion)s
+
+              gcloud alpha compute tpus tpu-vm ssh ${tpu_name}  --zone=${zone} --project=${project}  --internal-ip --ssh-key-file=/scripts/id_rsa --worker=all --command "echo 'WRAPT_DISABLE_EXTENSIONS=true' | sudo tee -a /etc/environment"
               if [[ ${softwareVersion: -3} == "pod" ]]; then
                  gcloud alpha compute tpus tpu-vm ssh ${tpu_name}  --zone=${zone} --project=${project}  --internal-ip --ssh-key-file=/scripts/id_rsa --worker=all --command "sudo sed -i 's/TF_DOCKER_URL=.*/TF_DOCKER_URL=gcr.io\/cloud-tpu-v2-images-dev\/grpc_tpu_worker:se-nightly\"/' /etc/systemd/system/tpu-runtime.service"
                  gcloud alpha compute tpus tpu-vm ssh ${tpu_name}  --zone=${zone} --project=${project}  --internal-ip --ssh-key-file=/scripts/id_rsa --worker=all --command "sudo systemctl daemon-reload && sudo systemctl restart tpu-runtime"
diff --git a/tests/tensorflow/nightly/common.libsonnet b/tests/tensorflow/nightly/common.libsonnet
index 19254089e..6db15e038 100644
--- a/tests/tensorflow/nightly/common.libsonnet
+++ b/tests/tensorflow/nightly/common.libsonnet
@@ -177,6 +177,7 @@ local volumes = import 'templates/volumes.libsonnet';
               sleep %(sleepTime)d
 
               softwareVersion=%(softwareVersion)s
+              gcloud alpha compute tpus tpu-vm ssh ${tpu_name}  --zone=${zone} --project=${project}  --internal-ip --ssh-key-file=/scripts/id_rsa --worker=all --command "echo 'WRAPT_DISABLE_EXTENSIONS=true' | sudo tee -a /etc/environment"
               if [[ ${softwareVersion: -3} == "pod" ]]; then
                  gcloud alpha compute tpus tpu-vm ssh ${tpu_name}  --zone=${zone} --project=${project}  --internal-ip --ssh-key-file=/scripts/id_rsa --worker=all --command "sudo sed -i 's/TF_DOCKER_URL=.*/TF_DOCKER_URL=gcr.io\/cloud-tpu-v2-images-dev\/grpc_tpu_worker:nightly\"/' /etc/systemd/system/tpu-runtime.service"
                  gcloud alpha compute tpus tpu-vm ssh ${tpu_name}  --zone=${zone} --project=${project}  --internal-ip --ssh-key-file=/scripts/id_rsa --worker=all --command "sudo systemctl daemon-reload && sudo systemctl restart tpu-runtime"