feat: add support for enable_dashboard_access field for Training jobs…

… in SDK PiperOrigin-RevId: 512169773
googleapis · Feb 24, 2023 · 3500eab · 3500eab
1 parent 581939b
commit 3500eab
Show file tree

Hide file tree

Showing 2 changed files with 285 additions and 1 deletion.
diff --git a/google/cloud/aiplatform/training_jobs.py b/google/cloud/aiplatform/training_jobs.py
@@ -1479,6 +1479,7 @@ def _prepare_training_task_inputs_and_output_dir(
         timeout: Optional[int] = None,
         restart_job_on_worker_restart: bool = False,
         enable_web_access: bool = False,
+        enable_dashboard_access: bool = False,
         tensorboard: Optional[str] = None,
     ) -> Tuple[Dict, str]:
         """Prepares training task inputs and output directory for custom job.
@@ -1508,6 +1509,9 @@ def _prepare_training_task_inputs_and_output_dir(
                 Whether you want Vertex AI to enable interactive shell access
                 to training containers.
                 https://cloud.google.com/vertex-ai/docs/training/monitor-debug-interactive-shell
+            enable_dashboard_access (bool):
+                Whether you want Vertex AI to enable access to the customized dashboard
+                to training containers.
             tensorboard (str):
                 Optional. The name of a Vertex AI
                 [Tensorboard][google.cloud.aiplatform.v1beta1.Tensorboard]
@@ -1547,6 +1551,8 @@ def _prepare_training_task_inputs_and_output_dir(
             training_task_inputs["tensorboard"] = tensorboard
         if enable_web_access:
             training_task_inputs["enable_web_access"] = enable_web_access
+        if enable_dashboard_access:
+            training_task_inputs["enable_dashboard_access"] = enable_dashboard_access
 
         if timeout or restart_job_on_worker_restart:
             timeout = f"{timeout}s" if timeout else None
@@ -1608,7 +1614,9 @@ def _wait_callback(self):
 
             self._has_logged_custom_job = True
 
-        if self._gca_resource.training_task_inputs.get("enable_web_access"):
+        if self._gca_resource.training_task_inputs.get(
+            "enable_web_access"
+        ) or self._gca_resource.training_task_inputs.get("enable_dashboard_access"):
             self._log_web_access_uris()
 
     def _custom_job_console_uri(self) -> str:
@@ -2902,6 +2910,7 @@ def run(
         timeout: Optional[int] = None,
         restart_job_on_worker_restart: bool = False,
         enable_web_access: bool = False,
+        enable_dashboard_access: bool = False,
         tensorboard: Optional[str] = None,
         sync=True,
         create_request_timeout: Optional[float] = None,
@@ -3164,6 +3173,9 @@ def run(
                 Whether you want Vertex AI to enable interactive shell access
                 to training containers.
                 https://cloud.google.com/vertex-ai/docs/training/monitor-debug-interactive-shell
+            enable_dashboard_access (bool):
+                Whether you want Vertex AI to enable access to the customized dashboard
+                to training containers.
             tensorboard (str):
                 Optional. The name of a Vertex AI
                 [Tensorboard][google.cloud.aiplatform.v1beta1.Tensorboard]
@@ -3238,6 +3250,7 @@ def run(
             timeout=timeout,
             restart_job_on_worker_restart=restart_job_on_worker_restart,
             enable_web_access=enable_web_access,
+            enable_dashboard_access=enable_dashboard_access,
             tensorboard=tensorboard,
             reduction_server_container_uri=reduction_server_container_uri
             if reduction_server_replica_count > 0
@@ -3283,6 +3296,7 @@ def _run(
         timeout: Optional[int] = None,
         restart_job_on_worker_restart: bool = False,
         enable_web_access: bool = False,
+        enable_dashboard_access: bool = False,
         tensorboard: Optional[str] = None,
         reduction_server_container_uri: Optional[str] = None,
         sync=True,
@@ -3444,6 +3458,9 @@ def _run(
                 Whether you want Vertex AI to enable interactive shell access
                 to training containers.
                 https://cloud.google.com/vertex-ai/docs/training/monitor-debug-interactive-shell
+            enable_dashboard_access (bool):
+                Whether you want Vertex AI to enable access to the customized dashboard
+                to training containers.
             tensorboard (str):
                 Optional. The name of a Vertex AI
                 [Tensorboard][google.cloud.aiplatform.v1beta1.Tensorboard]
@@ -3517,6 +3534,7 @@ def _run(
             timeout=timeout,
             restart_job_on_worker_restart=restart_job_on_worker_restart,
             enable_web_access=enable_web_access,
+            enable_dashboard_access=enable_dashboard_access,
             tensorboard=tensorboard,
         )
 
@@ -3834,6 +3852,7 @@ def run(
         timeout: Optional[int] = None,
         restart_job_on_worker_restart: bool = False,
         enable_web_access: bool = False,
+        enable_dashboard_access: bool = False,
         tensorboard: Optional[str] = None,
         sync=True,
         create_request_timeout: Optional[float] = None,
@@ -4089,6 +4108,9 @@ def run(
                 Whether you want Vertex AI to enable interactive shell access
                 to training containers.
                 https://cloud.google.com/vertex-ai/docs/training/monitor-debug-interactive-shell
+            enable_dashboard_access (bool):
+                Whether you want Vertex AI to enable access to the customized dashboard
+                to training containers.
             tensorboard (str):
                 Optional. The name of a Vertex AI
                 [Tensorboard][google.cloud.aiplatform.v1beta1.Tensorboard]
@@ -4162,6 +4184,7 @@ def run(
             timeout=timeout,
             restart_job_on_worker_restart=restart_job_on_worker_restart,
             enable_web_access=enable_web_access,
+            enable_dashboard_access=enable_dashboard_access,
             tensorboard=tensorboard,
             reduction_server_container_uri=reduction_server_container_uri
             if reduction_server_replica_count > 0
@@ -4206,6 +4229,7 @@ def _run(
         timeout: Optional[int] = None,
         restart_job_on_worker_restart: bool = False,
         enable_web_access: bool = False,
+        enable_dashboard_access: bool = False,
         tensorboard: Optional[str] = None,
         reduction_server_container_uri: Optional[str] = None,
         sync=True,
@@ -4363,6 +4387,9 @@ def _run(
                 Whether you want Vertex AI to enable interactive shell access
                 to training containers.
                 https://cloud.google.com/vertex-ai/docs/training/monitor-debug-interactive-shell
+            enable_dashboard_access (bool):
+                Whether you want Vertex AI to enable access to the customized dashboard
+                to training containers.
             tensorboard (str):
                 Optional. The name of a Vertex AI
                 [Tensorboard][google.cloud.aiplatform.v1beta1.Tensorboard]
@@ -4430,6 +4457,7 @@ def _run(
             timeout=timeout,
             restart_job_on_worker_restart=restart_job_on_worker_restart,
             enable_web_access=enable_web_access,
+            enable_dashboard_access=enable_dashboard_access,
             tensorboard=tensorboard,
         )
 
@@ -6124,6 +6152,7 @@ def run(
         timeout: Optional[int] = None,
         restart_job_on_worker_restart: bool = False,
         enable_web_access: bool = False,
+        enable_dashboard_access: bool = False,
         tensorboard: Optional[str] = None,
         sync=True,
         create_request_timeout: Optional[float] = None,
@@ -6379,6 +6408,9 @@ def run(
                 Whether you want Vertex AI to enable interactive shell access
                 to training containers.
                 https://cloud.google.com/vertex-ai/docs/training/monitor-debug-interactive-shell
+            enable_dashboard_access (bool):
+                Whether you want Vertex AI to enable access to the customized dashboard
+                to training containers.
             tensorboard (str):
                 Optional. The name of a Vertex AI
                 [Tensorboard][google.cloud.aiplatform.v1beta1.Tensorboard]
@@ -6447,6 +6479,7 @@ def run(
             timeout=timeout,
             restart_job_on_worker_restart=restart_job_on_worker_restart,
             enable_web_access=enable_web_access,
+            enable_dashboard_access=enable_dashboard_access,
             tensorboard=tensorboard,
             reduction_server_container_uri=reduction_server_container_uri
             if reduction_server_replica_count > 0
@@ -6491,6 +6524,7 @@ def _run(
         timeout: Optional[int] = None,
         restart_job_on_worker_restart: bool = False,
         enable_web_access: bool = False,
+        enable_dashboard_access: bool = False,
         tensorboard: Optional[str] = None,
         reduction_server_container_uri: Optional[str] = None,
         sync=True,
@@ -6635,6 +6669,9 @@ def _run(
                 Whether you want Vertex AI to enable interactive shell access
                 to training containers.
                 https://cloud.google.com/vertex-ai/docs/training/monitor-debug-interactive-shell
+            enable_dashboard_access (bool):
+                Whether you want Vertex AI to enable access to the customized dashboard
+                to training containers.
             tensorboard (str):
                 Optional. The name of a Vertex AI
                 [Tensorboard][google.cloud.aiplatform.v1beta1.Tensorboard]
@@ -6702,6 +6739,7 @@ def _run(
             timeout=timeout,
             restart_job_on_worker_restart=restart_job_on_worker_restart,
             enable_web_access=enable_web_access,
+            enable_dashboard_access=enable_dashboard_access,
             tensorboard=tensorboard,
         )