adding boot_disk_type and boot_disk_size_gb (#177)

* adding boot_disk_type and boot_disk_size_gb there params now will be standard on vertexai block, with defaults set to how they were implicitly in the code before. * param fix and linting * test now passes --------- Co-authored-by: Alexander Streed <[email protected]>
PrefectHQ · May 23, 2023 · 56bd219 · 56bd219
1 parent 1a586e0
commit 56bd219
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 3 deletions.
diff --git a/prefect_gcp/aiplatform.py b/prefect_gcp/aiplatform.py
@@ -78,7 +78,7 @@
     )
     from google.cloud.aiplatform_v1.types.job_service import CancelCustomJobRequest
     from google.cloud.aiplatform_v1.types.job_state import JobState
-    from google.cloud.aiplatform_v1.types.machine_resources import MachineSpec
+    from google.cloud.aiplatform_v1.types.machine_resources import DiskSpec, MachineSpec
     from google.protobuf.duration_pb2 import Duration
 except ModuleNotFoundError:
     pass
@@ -141,6 +141,16 @@ class VertexAICustomTrainingJob(Infrastructure):
     accelerator_count: Optional[int] = Field(
         default=None, description="The number of accelerators to attach to the machine."
     )
+    boot_disk_type: str = Field(
+        default="pd-ssd",
+        title="Boot Disk Type",
+        description="The type of boot disk to attach to the machine.",
+    )
+    boot_disk_size_gb: int = Field(
+        default=100,
+        title="Boot Disk Size",
+        description="The size of the boot disk to attach to the machine, in gigabytes.",
+    )
     maximum_run_time: datetime.timedelta = Field(
         default=datetime.timedelta(days=7), description="The maximum job running time."
     )
@@ -223,9 +233,14 @@ def _build_job_spec(self) -> "CustomJobSpec":
             accelerator_count=self.accelerator_count,
         )
         worker_pool_spec = WorkerPoolSpec(
-            container_spec=container_spec, machine_spec=machine_spec, replica_count=1
+            container_spec=container_spec,
+            machine_spec=machine_spec,
+            replica_count=1,
+            disk_spec=DiskSpec(
+                boot_disk_type=self.boot_disk_type,
+                boot_disk_size_gb=self.boot_disk_size_gb,
+            ),
         )
-
         # look for service account
         service_account = (
             self.service_account or self.gcp_credentials._service_account_email

diff --git a/tests/test_aiplatform.py b/tests/test_aiplatform.py
@@ -46,6 +46,10 @@ def test_preview(
                         machine_type: "n1-standard-4"
                     }
                     replica_count: 1
+                    disk_spec {
+                        boot_disk_type: "pd-ssd"
+                        boot_disk_size_gb: 100
+                    }
                 }
                 scheduling {
                 }