@@ -1364,6 +1364,7 @@
         "from google.cloud import aiplatform\n",
         "\n",
         "epochs = 100\n",
+        "timeout_min = 60  # 1 hour\n",
         "\n",
         "# Cloud Storage paths.\n",
         "data_path = f\"/gcs/{bucket}/weather/data\"\n",
@@ -1376,7 +1377,7 @@
         "    display_name=\"weather-forecasting\",\n",
         "    python_package_gcs_uri=f\"gs://{bucket}/weather/weather-model-1.0.0.tar.gz\",\n",
         "    python_module_name=\"weather.trainer\",\n",
-        "    container_uri=\"us-docker.pkg.dev/vertex-ai/training/pytorch-gpu.1-11:latest\",\n",
+        "    container_uri=\"us-docker.pkg.dev/vertex-ai/training/pytorch-gpu.1-13:latest\",\n",
         ")\n",
         "job.run(\n",
         "    machine_type=\"n1-highmem-8\",\n",
@@ -1387,6 +1388,7 @@
         "        f\"--model-path={model_path}\",\n",
         "        f\"--epochs={epochs}\",\n",
         "    ],\n",
+        "    timeout=timeout_min * 60,  # in seconds\n",
         ")"
       ],
       "id": "Ny4x99GiS2Lm"
 
@@ -18,6 +18,7 @@ name = "weather-model"
 version = "1.0.0"
 dependencies = [
     "datasets==2.8.0",
+    "torch==1.13.1",  # make sure this matches the `container_uri` in `notebooks/3-training.ipynb`
     "transformers==4.25.1",
 ]
 
 
@@ -1,4 +1,3 @@
 ../../serving/weather-model
 build==0.10.0
 google-cloud-aiplatform==1.21.0
-torch==1.13.1
@@ -82,6 +82,7 @@ def test_train_model(
                     "data_path": data_path_gcs.replace("gs://", "/gcs/"),
                     "model_path": f"/gcs/{bucket_name}/test/weather/model-vertex",
                     "epochs": 2,
+                    "timeout_min": 5,  # should take a couple seconds, so timeout in 5 minutes
                 }
             },
         },
Original file line number	Diff line number	Diff line change
`@@ -18,6 +18,7 @@ name = "weather-model"`
`18`	`18`	`version = "1.0.0"`
`19`	`19`	`dependencies = [`
`20`	`20`	`"datasets==2.8.0",`
	`21`	+ "torch==1.13.1", # make sure this matches the `container_uri` in `notebooks/3-training.ipynb`
`21`	`22`	`"transformers==4.25.1",`
`22`	`23`	`]`
`23`	`24`
Original file line number	Diff line number	Diff line change
`@@ -82,6 +82,7 @@ def test_train_model(`
`82`	`82`	`"data_path": data_path_gcs.replace("gs://", "/gcs/"),`
`83`	`83`	`"model_path": f"/gcs/{bucket_name}/test/weather/model-vertex",`
`84`	`84`	`"epochs": 2,`
	`85`	`+ "timeout_min": 5, # should take a couple seconds, so timeout in 5 minutes`
`85`	`86`	`}`
`86`	`87`	`},`
`87`	`88`	`},`