Merge branch 'xgboost_security' of https://github.com/GoogleCloudPlatform/vertex-ai-samples into xgboost_security

notebooks/official/training/xgboost_data_parallel_training_on_cpu_using_dask.ipynb

@@ -67,7 +67,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "AksIKBzZ-nre"
@@ -552,6 +551,8 @@
         "    'nthread': 1\n",
         "}\n",
         "\n",
+        "ROUNDS = 2\n",
+        "\n",
         "\n",
         "def square(x):\n",
         "    return x ** 2\n",
@@ -615,7 +616,7 @@
         "      wait(y)\n",
         "      dtrain = DaskDMatrix(client, X, y)\n",
         "\n",
-        "      output = xgb.dask.train(client, XGB_PARAMS, dtrain,  num_boost_round=100, evals=[(dtrain, 'train')])\n",
+        "      output = xgb.dask.train(client, XGB_PARAMS, dtrain,  num_boost_round=ROUNDS, evals=[(dtrain, 'train')])\n",
         "      print(\"Output: {}\".format(output), flush=True)\n",
         "      print(\"Saving file to: {}\".format(MODEL_FILE), flush=True)\n",
         "      output['booster'].save_model(MODEL_FILE)\n",
@@ -902,14 +903,15 @@
         "    container_uri=TRAIN_IMAGE,\n",
         ")\n",
         "\n",
-        "custom_container_training_job.run(\n",
-        "    base_output_dir=gcs_output_uri_prefix,\n",
-        "    replica_count=replica_count,\n",
-        "    machine_type=machine_type,\n",
-        "    enable_dashboard_access=True,\n",
-        "    enable_web_access=True,\n",
-        "    sync=False,\n",
-        ")"
+        "if not os.getenv(\"IS_TESTING\"):\n",
+        "    custom_container_training_job.run(\n",
+        "        base_output_dir=gcs_output_uri_prefix,\n",
+        "        replica_count=replica_count,\n",
+        "        machine_type=machine_type,\n",
+        "        enable_dashboard_access=True,\n",
+        "        enable_web_access=True,\n",
+        "        sync=False,\n",
+        "    )"
       ]
     },
     {
@@ -942,8 +944,9 @@
       },
       "outputs": [],
       "source": [
-        "print(f\"Custom Training Job Name: {custom_container_training_job.resource_name}\")\n",
-        "print(f\"GCS Output URI Prefix: {gcs_output_uri_prefix}\")"
+        "if not os.getenv(\"IS_TESTING\"):\n",
+        "    print(f\"Custom Training Job Name: {custom_container_training_job.resource_name}\")\n",
+        "    print(f\"GCS Output URI Prefix: {gcs_output_uri_prefix}\")"
       ]
     },
     {
@@ -963,9 +966,10 @@
       },
       "outputs": [],
       "source": [
-        "print(\n",
-        "    f\"Custom Training Job URI: {custom_container_training_job._custom_job_console_uri()}\"\n",
-        ")"
+        "if not os.getenv(\"IS_TESTING\"):\n",
+        "    print(\n",
+        "        f\"Custom Training Job URI: {custom_container_training_job._custom_job_console_uri()}\"\n",
+        "    )"
       ]
     },
     {
@@ -985,7 +989,10 @@
       },
       "outputs": [],
       "source": [
-        "print(f\"Web Access and Dashboard URIs: {custom_container_training_job.web_access_uris}\")"
+        "if not os.getenv(\"IS_TESTING\"):\n",
+        "    print(\n",
+        "        f\"Web Access and Dashboard URIs: {custom_container_training_job.web_access_uris}\"\n",
+        "    )"
       ]
     },
     {
@@ -1066,7 +1073,8 @@
       },
       "outputs": [],
       "source": [
-        "! gcloud ai custom-jobs create --region=us-central1 --config=config.yaml --display-name={display_name}"
+        "if not os.getenv(\"IS_TESTING\"):\n",
+        "    ! gcloud ai custom-jobs create --region=us-central1 --config=config.yaml --display-name={display_name}"
       ]
     },
     {
@@ -1187,10 +1195,7 @@
         "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n",
         "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
         "\n",
-        "Otherwise, you can delete the individual resources you created in this tutorial:\n",
-        "\n",
-        "- Cloud Storage Bucket\n",
-        "- Cloud Vertex Training Job"
+        "Otherwise, you can delete the individual resources you created in this tutorial."
       ]
     },
     {
@@ -1201,9 +1206,6 @@
       },
       "outputs": [],
       "source": [
-        "import logging\n",
-        "import traceback\n",
-        "\n",
         "# Set this to true only if you'd like to delete your bucket\n",
         "delete_bucket = False\n",
         "\n",
@@ -1215,8 +1217,10 @@
         "try:\n",
         "    custom_container_training_job.delete()\n",
         "except Exception as e:\n",
-        "    logging.error(traceback.format_exc())\n",
-        "    print(e)"
+        "    print(e)\n",
+        "\n",
+        "# delete the docker image\n",
+        "! gcloud artifacts repositories delete  --location {REGION}  {TRAIN_IMAGE}"
       ]
     }
   ],