feat: Add Lustre support to the Vertex Training Custom Job API

google/cloud/aiplatform_v1/__init__.py

@@ -617,6 +617,7 @@
 from .types.machine_resources import BatchDedicatedResources
 from .types.machine_resources import DedicatedResources
 from .types.machine_resources import DiskSpec
+from .types.machine_resources import LustreMount
 from .types.machine_resources import MachineSpec
 from .types.machine_resources import NfsMount
 from .types.machine_resources import PersistentDiskSpec
@@ -1775,6 +1776,7 @@ def _get_version(dependency_name):
     "LlmUtilityServiceClient",
     "LogprobsResult",
     "LookupStudyRequest",
+    "LustreMount",
     "MachineSpec",
     "ManualBatchTuningParameters",
     "MatchServiceClient",

google/cloud/aiplatform_v1/services/migration_service/client.py

@@ -271,40 +271,40 @@ def parse_annotated_dataset_path(path: str) -> Dict[str, str]:
     @staticmethod
     def dataset_path(
         project: str,
-        location: str,
         dataset: str,
     ) -> str:
         """Returns a fully-qualified dataset string."""
-        return "projects/{project}/locations/{location}/datasets/{dataset}".format(
+        return "projects/{project}/datasets/{dataset}".format(
             project=project,
-            location=location,
             dataset=dataset,
         )
 
     @staticmethod
     def parse_dataset_path(path: str) -> Dict[str, str]:
         """Parses a dataset path into its component segments."""
-        m = re.match(
-            r"^projects/(?P<project>.+?)/locations/(?P<location>.+?)/datasets/(?P<dataset>.+?)$",
-            path,
-        )
+        m = re.match(r"^projects/(?P<project>.+?)/datasets/(?P<dataset>.+?)$", path)
         return m.groupdict() if m else {}
 
     @staticmethod
     def dataset_path(
         project: str,
+        location: str,
         dataset: str,
     ) -> str:
         """Returns a fully-qualified dataset string."""
-        return "projects/{project}/datasets/{dataset}".format(
+        return "projects/{project}/locations/{location}/datasets/{dataset}".format(
             project=project,
+            location=location,
             dataset=dataset,
         )
 
     @staticmethod
     def parse_dataset_path(path: str) -> Dict[str, str]:
         """Parses a dataset path into its component segments."""
-        m = re.match(r"^projects/(?P<project>.+?)/datasets/(?P<dataset>.+?)$", path)
+        m = re.match(
+            r"^projects/(?P<project>.+?)/locations/(?P<location>.+?)/datasets/(?P<dataset>.+?)$",
+            path,
+        )
         return m.groupdict() if m else {}
 
     @staticmethod

google/cloud/aiplatform_v1/types/__init__.py

@@ -630,6 +630,7 @@
     BatchDedicatedResources,
     DedicatedResources,
     DiskSpec,
+    LustreMount,
     MachineSpec,
     NfsMount,
     PersistentDiskSpec,
@@ -1718,6 +1719,7 @@
     "BatchDedicatedResources",
     "DedicatedResources",
     "DiskSpec",
+    "LustreMount",
     "MachineSpec",
     "NfsMount",
     "PersistentDiskSpec",

google/cloud/aiplatform_v1/types/custom_job.py

@@ -420,6 +420,8 @@ class WorkerPoolSpec(proto.Message):
             use for this worker pool.
         nfs_mounts (MutableSequence[google.cloud.aiplatform_v1.types.NfsMount]):
             Optional. List of NFS mount spec.
+        lustre_mounts (MutableSequence[google.cloud.aiplatform_v1.types.LustreMount]):
+            Optional. List of Lustre mounts.
         disk_spec (google.cloud.aiplatform_v1.types.DiskSpec):
             Disk spec.
     """
@@ -450,6 +452,11 @@ class WorkerPoolSpec(proto.Message):
         number=4,
         message=machine_resources.NfsMount,
     )
+    lustre_mounts: MutableSequence[machine_resources.LustreMount] = proto.RepeatedField(
+        proto.MESSAGE,
+        number=9,
+        message=machine_resources.LustreMount,
+    )
     disk_spec: machine_resources.DiskSpec = proto.Field(
         proto.MESSAGE,
         number=5,

google/cloud/aiplatform_v1/types/machine_resources.py

@@ -36,6 +36,7 @@
         "DiskSpec",
         "PersistentDiskSpec",
         "NfsMount",
+        "LustreMount",
         "AutoscalingMetricSpec",
         "ShieldedVmConfig",
     },
@@ -125,33 +126,32 @@ class MachineSpec(proto.Message):
 
 class DedicatedResources(proto.Message):
     r"""A description of resources that are dedicated to a
-    DeployedModel, and that need a higher degree of manual
-    configuration.
+    DeployedModel or DeployedIndex, and that need a higher degree of
+    manual configuration.
 
     Attributes:
         machine_spec (google.cloud.aiplatform_v1.types.MachineSpec):
             Required. Immutable. The specification of a
-            single machine used by the prediction.
+            single machine being used.
         min_replica_count (int):
             Required. Immutable. The minimum number of
-            machine replicas this DeployedModel will be
-            always deployed on. This value must be greater
-            than or equal to 1.
-
-            If traffic against the DeployedModel increases,
-            it may dynamically be deployed onto more
-            replicas, and as traffic decreases, some of
-            these extra replicas may be freed.
+            machine replicas that will be always deployed
+            on. This value must be greater than or equal to
+            1.
+
+            If traffic increases, it may dynamically be
+            deployed onto more replicas, and as traffic
+            decreases, some of these extra replicas may be
+            freed.
         max_replica_count (int):
-            Immutable. The maximum number of replicas this DeployedModel
-            may be deployed on when the traffic against it increases. If
-            the requested value is too large, the deployment will error,
-            but if deployment succeeds then the ability to scale the
-            model to that many replicas is guaranteed (barring service
-            outages). If traffic against the DeployedModel increases
-            beyond what its replicas at maximum may handle, a portion of
-            the traffic will be dropped. If this value is not provided,
-            will use
+            Immutable. The maximum number of replicas that may be
+            deployed on when the traffic against it increases. If the
+            requested value is too large, the deployment will error, but
+            if deployment succeeds then the ability to scale to that
+            many replicas is guaranteed (barring service outages). If
+            traffic increases beyond what its replicas at maximum may
+            handle, a portion of the traffic will be dropped. If this
+            value is not provided, will use
             [min_replica_count][google.cloud.aiplatform.v1.DedicatedResources.min_replica_count]
             as the default value.
 
@@ -163,8 +163,8 @@ class DedicatedResources(proto.Message):
         required_replica_count (int):
             Optional. Number of required available replicas for the
             deployment to succeed. This field is only needed when
-            partial model deployment/mutation is desired. If set, the
-            model deploy/mutate operation will succeed once
+            partial deployment/mutation is desired. If set, the
+            deploy/mutate operation will succeed once
             available_replica_count reaches required_replica_count, and
             the rest of the replicas will be retried. If not set, the
             default required_replica_count will be min_replica_count.
@@ -239,28 +239,27 @@ class AutomaticResources(proto.Message):
 
     Attributes:
         min_replica_count (int):
-            Immutable. The minimum number of replicas this DeployedModel
-            will be always deployed on. If traffic against it increases,
-            it may dynamically be deployed onto more replicas up to
+            Immutable. The minimum number of replicas that will be
+            always deployed on. If traffic against it increases, it may
+            dynamically be deployed onto more replicas up to
             [max_replica_count][google.cloud.aiplatform.v1.AutomaticResources.max_replica_count],
             and as traffic decreases, some of these extra replicas may
             be freed. If the requested value is too large, the
             deployment will error.
         max_replica_count (int):
             Immutable. The maximum number of replicas
-            this DeployedModel may be deployed on when the
-            traffic against it increases. If the requested
-            value is too large, the deployment will error,
-            but if deployment succeeds then the ability to
-            scale the model to that many replicas is
-            guaranteed (barring service outages). If traffic
-            against the DeployedModel increases beyond what
-            its replicas at maximum may handle, a portion of
-            the traffic will be dropped. If this value is
-            not provided, a no upper bound for scaling under
-            heavy traffic will be assume, though Vertex AI
-            may be unable to scale beyond certain replica
-            number.
+            that may be deployed on when the traffic against
+            it increases. If the requested value is too
+            large, the deployment will error, but if
+            deployment succeeds then the ability to scale to
+            that many replicas is guaranteed (barring
+            service outages). If traffic increases beyond
+            what its replicas at maximum may handle, a
+            portion of the traffic will be dropped. If this
+            value is not provided, a no upper bound for
+            scaling under heavy traffic will be assume,
+            though Vertex AI may be unable to scale beyond
+            certain replica number.
     """
 
     min_replica_count: int = proto.Field(
@@ -331,10 +330,12 @@ class DiskSpec(proto.Message):
 
     Attributes:
         boot_disk_type (str):
-            Type of the boot disk (default is "pd-ssd").
-            Valid values: "pd-ssd" (Persistent Disk Solid
-            State Drive) or "pd-standard" (Persistent Disk
-            Hard Disk Drive).
+            Type of the boot disk. For non-A3U machines,
+            the default value is "pd-ssd", for A3U machines,
+            the default value is "hyperdisk-balanced". Valid
+            values: "pd-ssd" (Persistent Disk Solid State
+            Drive), "pd-standard" (Persistent Disk Hard Disk
+            Drive) or "hyperdisk-balanced".
         boot_disk_size_gb (int):
             Size in GB of the boot disk (default is
             100GB).
@@ -407,6 +408,40 @@ class NfsMount(proto.Message):
     )
 
 
+class LustreMount(proto.Message):
+    r"""Represents a mount configuration for Lustre file system.
+
+    Attributes:
+        instance_ip (str):
+            Required. IP address of the Lustre instance.
+        volume_handle (str):
+            Required. The unique identifier of the Lustre
+            volume.
+        filesystem (str):
+            Required. The name of the Lustre filesystem.
+        mount_point (str):
+            Required. Destination mount path. The Lustre file system
+            will be mounted for the user under /mnt/lustre/<mount_point>
+    """
+
+    instance_ip: str = proto.Field(
+        proto.STRING,
+        number=1,
+    )
+    volume_handle: str = proto.Field(
+        proto.STRING,
+        number=2,
+    )
+    filesystem: str = proto.Field(
+        proto.STRING,
+        number=3,
+    )
+    mount_point: str = proto.Field(
+        proto.STRING,
+        number=4,
+    )
+
+
 class AutoscalingMetricSpec(proto.Message):
     r"""The metric specification that defines the target resource
     utilization (CPU utilization, accelerator's duty cycle, and so
@@ -419,6 +454,7 @@ class AutoscalingMetricSpec(proto.Message):
             - For Online Prediction:
             - ``aiplatform.googleapis.com/prediction/online/accelerator/duty_cycle``
             - ``aiplatform.googleapis.com/prediction/online/cpu/utilization``
+            - ``aiplatform.googleapis.com/prediction/online/request_count``
         target (int):
             The target resource utilization in percentage
             (1% - 100%) for the given metric; once the real

samples/generated_samples/snippet_metadata_google.cloud.aiplatform.v1.json

@@ -8,7 +8,7 @@
     ],
     "language": "PYTHON",
     "name": "google-cloud-aiplatform",
-    "version": "1.132.0"
+    "version": "0.0.0"
   },
   "snippets": [
     {

samples/generated_samples/snippet_metadata_google.cloud.aiplatform.v1beta1.json

@@ -8,7 +8,7 @@
     ],
     "language": "PYTHON",
     "name": "google-cloud-aiplatform",
-    "version": "1.132.0"
+    "version": "0.0.0"
   },
   "snippets": [
     {

tests/unit/gapic/aiplatform_v1/test_job_service.py

@@ -24609,6 +24609,14 @@ def test_create_custom_job_rest_call_success(request_type):
                             "mount_point": "mount_point_value",
                         }
                     ],
+                    "lustre_mounts": [
+                        {
+                            "instance_ip": "instance_ip_value",
+                            "volume_handle": "volume_handle_value",
+                            "filesystem": "filesystem_value",
+                            "mount_point": "mount_point_value",
+                        }
+                    ],
                     "disk_spec": {
                         "boot_disk_type": "boot_disk_type_value",
                         "boot_disk_size_gb": 1792,
@@ -26268,6 +26276,14 @@ def test_create_hyperparameter_tuning_job_rest_call_success(request_type):
                             "mount_point": "mount_point_value",
                         }
                     ],
+                    "lustre_mounts": [
+                        {
+                            "instance_ip": "instance_ip_value",
+                            "volume_handle": "volume_handle_value",
+                            "filesystem": "filesystem_value",
+                            "mount_point": "mount_point_value",
+                        }
+                    ],
                     "disk_spec": {
                         "boot_disk_type": "boot_disk_type_value",
                         "boot_disk_size_gb": 1792,
@@ -27155,6 +27171,14 @@ def test_create_nas_job_rest_call_success(request_type):
                                         "mount_point": "mount_point_value",
                                     }
                                 ],
+                                "lustre_mounts": [
+                                    {
+                                        "instance_ip": "instance_ip_value",
+                                        "volume_handle": "volume_handle_value",
+                                        "filesystem": "filesystem_value",
+                                        "mount_point": "mount_point_value",
+                                    }
+                                ],
                                 "disk_spec": {
                                     "boot_disk_type": "boot_disk_type_value",
                                     "boot_disk_size_gb": 1792,
@@ -32047,6 +32071,14 @@ async def test_create_custom_job_rest_asyncio_call_success(request_type):
                             "mount_point": "mount_point_value",
                         }
                     ],
+                    "lustre_mounts": [
+                        {
+                            "instance_ip": "instance_ip_value",
+                            "volume_handle": "volume_handle_value",
+                            "filesystem": "filesystem_value",
+                            "mount_point": "mount_point_value",
+                        }
+                    ],
                     "disk_spec": {
                         "boot_disk_type": "boot_disk_type_value",
                         "boot_disk_size_gb": 1792,
@@ -33871,6 +33903,14 @@ async def test_create_hyperparameter_tuning_job_rest_asyncio_call_success(reques
                             "mount_point": "mount_point_value",
                         }
                     ],
+                    "lustre_mounts": [
+                        {
+                            "instance_ip": "instance_ip_value",
+                            "volume_handle": "volume_handle_value",
+                            "filesystem": "filesystem_value",
+                            "mount_point": "mount_point_value",
+                        }
+                    ],
                     "disk_spec": {
                         "boot_disk_type": "boot_disk_type_value",
                         "boot_disk_size_gb": 1792,
@@ -34856,6 +34896,14 @@ async def test_create_nas_job_rest_asyncio_call_success(request_type):
                                         "mount_point": "mount_point_value",
                                     }
                                 ],
+                                "lustre_mounts": [
+                                    {
+                                        "instance_ip": "instance_ip_value",
+                                        "volume_handle": "volume_handle_value",
+                                        "filesystem": "filesystem_value",
+                                        "mount_point": "mount_point_value",
+                                    }
+                                ],
                                 "disk_spec": {
                                     "boot_disk_type": "boot_disk_type_value",
                                     "boot_disk_size_gb": 1792,