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

Google APIs · copybara-github · commit a8e146aa1669 · 2025-12-16T11:07:32.000-08:00
docs: A comment for field `timeout` in message `.google.cloud.aiplatform.v1beta1.Scheduling` is changed
docs: A comment for field `restart_job_on_worker_restart` in message `.google.cloud.aiplatform.v1beta1.Scheduling` is changed

PiperOrigin-RevId: 845354898
diff --git a/google/cloud/aiplatform/v1beta1/custom_job.proto b/google/cloud/aiplatform/v1beta1/custom_job.proto
@@ -312,6 +312,10 @@ message WorkerPoolSpec {
   // Optional. List of NFS mount spec.
   repeated NfsMount nfs_mounts = 4 [(google.api.field_behavior) = OPTIONAL];
 
+  // Optional. List of Lustre mounts.
+  repeated LustreMount lustre_mounts = 9
+      [(google.api.field_behavior) = OPTIONAL];
+
   // Disk spec.
   DiskSpec disk_spec = 5;
 }
@@ -387,10 +391,10 @@ message Scheduling {
     FLEX_START = 6;
   }
 
-  // The maximum job running time. The default is 7 days.
+  // Optional. The maximum job running time. The default is 7 days.
   google.protobuf.Duration timeout = 1;
 
-  // Restarts the entire CustomJob if a worker gets restarted.
+  // Optional. Restarts the entire CustomJob if a worker gets restarted.
   // This feature can be used by distributed training jobs that are not
   // resilient to workers leaving and joining a job.
   bool restart_job_on_worker_restart = 3;
diff --git a/google/cloud/aiplatform/v1beta1/machine_resources.proto b/google/cloud/aiplatform/v1beta1/machine_resources.proto
@@ -355,6 +355,22 @@ message NfsMount {
   string mount_point = 3 [(google.api.field_behavior) = REQUIRED];
 }
 
+// Represents a mount configuration for Lustre file system.
+message LustreMount {
+  // Required. IP address of the Lustre instance.
+  string instance_ip = 1 [(google.api.field_behavior) = REQUIRED];
+
+  // Required. The unique identifier of the Lustre volume.
+  string volume_handle = 2 [(google.api.field_behavior) = REQUIRED];
+
+  // Required. The name of the Lustre filesystem.
+  string filesystem = 3 [(google.api.field_behavior) = REQUIRED];
+
+  // Required. Destination mount path. The Lustre file system will be mounted
+  // for the user under /mnt/lustre/<mount_point>
+  string mount_point = 4 [(google.api.field_behavior) = REQUIRED];
+}
+
 // The metric specification that defines the target resource utilization
 // (CPU utilization, accelerator's duty cycle, and so on) for calculating the
 // desired replica count.
