Add Vertex Step Operator network parameter (#2398)

* feat: add extra parameters (e.g. network) to Vertex Step Operator * fix: use snake case for the reserved_ip_ranges parameter * fix: linting * Update docs/book/stacks-and-components/component-guide/step-operators/vertex.md Co-authored-by: Alex Strick van Linschoten <[email protected]> * fix: split reserved_ip_ranges string when populating cutom_job --------- Co-authored-by: Alex Strick van Linschoten <[email protected]>
zenml-io · Feb 27, 2024 · 23416e8 · 23416e8
1 parent 3932cd0
commit 23416e8
Show file tree

Hide file tree

Showing 3 changed files with 52 additions and 3 deletions.
diff --git a/docs/book/stacks-and-components/component-guide/step-operators/vertex.md b/docs/book/stacks-and-components/component-guide/step-operators/vertex.md
@@ -122,9 +122,38 @@ more about how ZenML builds these images and how you can customize them.
 
 #### Additional configuration
 
+You can specify the service account, network and reserved IP ranges to use for the VertexAI `CustomJob` by passing the `service_account`, `network` and `reserved_ip_ranges` parameters to the `step-operator register` command:
+
+```shell
+    zenml step-operator register <STEP_OPERATOR_NAME> \
+        --flavor=vertex \
+        --project=<GCP_PROJECT> \
+        --region=<REGION> \
+        --service_account=<SERVICE_ACCOUNT> # optionally specify the service account to use for the VertexAI CustomJob
+        --network=<NETWORK> # optionally specify the network to use for the VertexAI CustomJob
+        --reserved_ip_ranges=<RESERVED_IP_RANGES> # optionally specify the reserved IP range to use for the VertexAI CustomJob
+```
+
 For additional configuration of the Vertex step operator, you can pass `VertexStepOperatorSettings` when defining or
-running your pipeline. Check out
-the [SDK docs](https://sdkdocs.zenml.io/latest/integration\_code\_docs/integrations-gcp/#zenml.integrations.gcp.flavors.vertex\_step\_operator\_flavor.VertexStepOperatorSettings)
+running your pipeline.
+
+```python
+from zenml import step
+from zenml.integrations.gcp.flavors.vertex_step_operator_flavor import VertexStepOperatorSettings
+
+@step(step_operator= <NAME>, settings=settings= {"step_operator.vertex": vertex_operator_settings = VertexStepOperatorSettings(
+    accelerator_type  = "NVIDIA_TESLA_T4" # see https://cloud.google.com/vertex-ai/docs/reference/rest/v1/MachineSpec#AcceleratorType
+    accelerator_count = 1
+    machine_type = "n1-standard-2"        # see https://cloud.google.com/vertex-ai/docs/training/configure-compute#machine-types
+    disk_type = "pd-ssd"                  # see https://cloud.google.com/vertex-ai/docs/training/configure-storage#disk-types
+    disk_size_gb = 100                    # see https://cloud.google.com/vertex-ai/docs/training/configure-storage#disk-size
+    )})
+def trainer(...) -> ...:
+    """Train a model."""
+    # This step will be executed in Vertex.
+```
+
+Check out the [SDK docs](https://sdkdocs.zenml.io/latest/integration\_code\_docs/integrations-gcp/#zenml.integrations.gcp.flavors.vertex\_step\_operator\_flavor.VertexStepOperatorSettings)
 for a full list of available attributes and [this docs page](/docs/book/user-guide/advanced-guide/pipelining-features/pipeline-settings.md) for
 more information on how to specify settings.
 

diff --git a/src/zenml/integrations/gcp/flavors/vertex_step_operator_flavor.py b/src/zenml/integrations/gcp/flavors/vertex_step_operator_flavor.py
@@ -71,6 +71,13 @@ class VertexStepOperatorConfig(  # type: ignore[misc] # https://github.com/pydan
     Attributes:
         region: Region name, e.g., `europe-west1`.
         encryption_spec_key_name: Encryption spec key name.
+        network: The full name of the Compute Engine network to which the Job should be peered.
+            For example, projects/12345/global/networks/myVPC
+        reserved_ip_ranges: A list of names for the reserved ip ranges under the VPC network that can be used
+            for this job. If set, we will deploy the job within the provided ip ranges. Otherwise, the job
+            will be deployed to any ip ranges under the provided VPC network.
+        service_account: Specifies the service account for workload run-as account. Users submitting jobs
+            must have act-as permission on this run-as account.
     """
 
     region: str
@@ -79,6 +86,12 @@ class VertexStepOperatorConfig(  # type: ignore[misc] # https://github.com/pydan
     # will be applied to all Vertex AI resources if set
     encryption_spec_key_name: Optional[str] = None
 
+    network: Optional[str] = None
+
+    reserved_ip_ranges: Optional[str] = None
+
+    service_account: Optional[str] = None
+
     @property
     def is_remote(self) -> bool:
         """Checks if this stack component is running remotely.

diff --git a/src/zenml/integrations/gcp/step_operators/vertex_step_operator.py b/src/zenml/integrations/gcp/step_operators/vertex_step_operator.py
@@ -250,7 +250,14 @@ def launch(
                             "boot_disk_size_gb": settings.boot_disk_size_gb,
                         },
                     }
-                ]
+                ],
+                "service_account": self.config.service_account,
+                "network": self.config.network,
+                "reserved_ip_ranges": (
+                    self.config.reserved_ip_ranges.split(",")
+                    if self.config.reserved_ip_ranges
+                    else []
+                ),
             },
             "labels": job_labels,
             "encryption_spec": {