feat: add CustomEnvironmentSpec to NotebookExecutionJob

Google APIs · copybara-github · commit 7b30cb57f3ce · 2024-11-18T11:14:04.000-08:00
PiperOrigin-RevId: 697694162
diff --git a/google/cloud/aiplatform/v1/notebook_execution_job.proto b/google/cloud/aiplatform/v1/notebook_execution_job.proto
@@ -20,6 +20,8 @@ import "google/api/field_behavior.proto";
 import "google/api/resource.proto";
 import "google/cloud/aiplatform/v1/encryption_spec.proto";
 import "google/cloud/aiplatform/v1/job_state.proto";
+import "google/cloud/aiplatform/v1/machine_resources.proto";
+import "google/cloud/aiplatform/v1/network_spec.proto";
 import "google/protobuf/duration.proto";
 import "google/protobuf/timestamp.proto";
 import "google/rpc/status.proto";
@@ -70,6 +72,18 @@ message NotebookExecutionJob {
     bytes content = 1;
   }
 
+  // Compute configuration to use for an execution job.
+  message CustomEnvironmentSpec {
+    // The specification of a single machine for the execution job.
+    MachineSpec machine_spec = 1;
+
+    // The specification of a persistent disk to attach for the execution job.
+    PersistentDiskSpec persistent_disk_spec = 2;
+
+    // The network configuration to use for the execution job.
+    NetworkSpec network_spec = 3;
+  }
+
   // The input notebook.
   oneof notebook_source {
     // The Dataform Repository pointing to a single file notebook repository.
@@ -90,6 +104,9 @@ message NotebookExecutionJob {
         [(google.api.resource_reference) = {
           type: "aiplatform.googleapis.com/NotebookRuntimeTemplate"
         }];
+
+    // The custom compute configuration for an execution job.
+    CustomEnvironmentSpec custom_environment_spec = 16;
   }
 
   // The location to store the notebook execution result.
