fix: import tf

notebooks/official/tensorboard/tensorboard_profiler_custom_training.ipynb

@@ -139,7 +139,7 @@
       },
       "outputs": [],
       "source": [
-        "! pip3 install --user --upgrade google-cloud-aiplatform --quiet"
+        "! pip3 install --upgrade --quiet google-cloud-aiplatform "
       ]
     },
     {
@@ -318,7 +318,27 @@
       },
       "outputs": [],
       "source": [
-        "SERVICE_ACCOUNT = \"[your-service-account]\"  # @param {type:\"string\"}"
+        "import sys\n",
+        "\n",
+        "SERVICE_ACCOUNT = \"[your-service-account]\"  # @param {type:\"string\"}\n",
+        "\n",
+        "IS_COLAB = \"google.colab\" in sys.modules\n",
+        "if (\n",
+        "    SERVICE_ACCOUNT == \"\"\n",
+        "    or SERVICE_ACCOUNT is None\n",
+        "    or SERVICE_ACCOUNT == \"[your-service-account]\"\n",
+        "):\n",
+        "    # Get your service account from gcloud\n",
+        "    if not IS_COLAB:\n",
+        "        shell_output = !gcloud auth list 2>/dev/null\n",
+        "        SERVICE_ACCOUNT = shell_output[2].replace(\"*\", \"\").strip()\n",
+        "\n",
+        "    if IS_COLAB:\n",
+        "        shell_output = ! gcloud projects describe  $PROJECT_ID\n",
+        "        project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
+        "        SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
+        "\n",
+        "    print(\"Service Account:\", SERVICE_ACCOUNT)"
       ]
     },
     {
@@ -364,7 +384,7 @@
       },
       "outputs": [],
       "source": [
-        "BUCKET_URI = \"gs://your-bucket-name-unique\"  # @param {type:\"string\"}"
+        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
       ]
     },
     {
@@ -404,6 +424,8 @@
       },
       "outputs": [],
       "source": [
+        "import os\n",
+        "\n",
         "from google.cloud import aiplatform"
       ]
     },
@@ -450,7 +472,11 @@
       },
       "outputs": [],
       "source": [
-        "! gcloud services enable artifactregistry.googleapis.com --quiet"
+        "! gcloud services enable artifactregistry.googleapis.com --quiet\n",
+        "\n",
+        "if os.getenv(\"IS_TESTING\"):\n",
+        "    ! sudo apt-get update --yes && sudo apt-get --only-upgrade --yes install google-cloud-sdk-cloud-run-proxy google-cloud-sdk-harbourbridge google-cloud-sdk-cbt google-cloud-sdk-gke-gcloud-auth-plugin google-cloud-sdk-kpt google-cloud-sdk-local-extract google-cloud-sdk-minikube google-cloud-sdk-app-engine-java google-cloud-sdk-app-engine-go google-cloud-sdk-app-engine-python google-cloud-sdk-spanner-emulator google-cloud-sdk-bigtable-emulator google-cloud-sdk-nomos google-cloud-sdk-package-go-module google-cloud-sdk-firestore-emulator kubectl google-cloud-sdk-datastore-emulator google-cloud-sdk-app-engine-python-extras google-cloud-sdk-cloud-build-local google-cloud-sdk-kubectl-oidc google-cloud-sdk-anthos-auth google-cloud-sdk-app-engine-grpc google-cloud-sdk-pubsub-emulator google-cloud-sdk-datalab google-cloud-sdk-skaffold google-cloud-sdk google-cloud-sdk-terraform-tools google-cloud-sdk-config-connector\n",
+        "    ! gcloud components update --quiet"
       ]
     },
     {
@@ -481,7 +507,7 @@
       },
       "outputs": [],
       "source": [
-        "TENSORBOARD_NAME = \"your-tensorboard-unique\"  # @param {type:\"string\"}"
+        "TENSORBOARD_NAME = f\"your-tensorboard-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
       ]
     },
     {
@@ -548,7 +574,7 @@
       },
       "outputs": [],
       "source": [
-        "DOCKER_REPOSITORY = f\"{PROJECT_ID}-repo-unique\"\n",
+        "DOCKER_REPOSITORY = f\"my-docker-repo-{PROJECT_ID}-unique\"\n",
         "\n",
         "! gcloud artifacts repositories create {DOCKER_REPOSITORY} \\\n",
         "    --repository-format=docker \\\n",
@@ -580,8 +606,6 @@
       "source": [
         "import sys\n",
         "\n",
-        "IS_COLAB = \"google.colab\" in sys.modules\n",
-        "\n",
         "if not IS_COLAB:\n",
         "    ! gcloud auth configure-docker {REGION}-docker.pkg.dev --quiet"
       ]
@@ -609,7 +633,7 @@
       "source": [
         "PYTHON_PACKAGE_APPLICATION_DIR = \"trainer\"\n",
         "\n",
-        "!mkdir -p $PYTHON_PACKAGE_APPLICATION_DIR"
+        "! mkdir -p $PYTHON_PACKAGE_APPLICATION_DIR"
       ]
     },
     {
@@ -765,6 +789,7 @@
         "\n",
         "# Installs additional packages as you need.\n",
         "RUN pip3 install \"google-cloud-aiplatform[cloud_profiler]>=1.20.0\"\n",
+        "RUN pip3 install \"protobuf==3.20.3\"\n",
         "\n",
         "# Copies the trainer code to the docker image.\n",
         "RUN mkdir /root/trainer\n",
@@ -913,12 +938,15 @@
         "# Delete docker repository.\n",
         "! gcloud artifacts repositories delete $DOCKER_REPOSITORY --project {PROJECT_ID} --location {REGION} --quiet\n",
         "\n",
-        "job.delete()\n",
+        "try:\n",
+        "    job.delete()\n",
+        "except Exception as e:\n",
+        "    print(e)\n",
         "\n",
         "if delete_tensorboard:\n",
         "    tensorboard.delete()\n",
         "\n",
-        "if delete_bucket and \"BUCKET_URI\" in globals():\n",
+        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
         "    ! gsutil -m rm -r $BUCKET_URI"
       ]
     }