feat: expose network policy interfaces (#28)

Signed-off-by: vsoch <vsoch@users.noreply.github.com>
Co-authored-by: vsoch <vsoch@users.noreply.github.com>

Author: vsoch

docs/further.md

@@ -65,33 +65,6 @@ the ratelimit is hit very easily. We are currently assessing strategies to deliv
 without hitting issues with this rate limit. It looks possible to create "[sinks](https://cloud.google.com/logging/docs/routing/overview#sinks)" using
 Pub Sub, however this would be adding an extra API dependency (and cost). 
 
-### Arguments
-
-And custom arguments can be any of the following, either on the command line or provided in the environment.
-
-| Name | Description | Flag | Type | Environment Variable | Required | Default |
-|------|-------------|------|------|----------------------|----------|---------|
-| project | The name of the Google Project | `--googlebatch-project` | str | `SNAKEMAKE_GOOGLEBATCH_PROJECT` | True |  unset |
-| region | The name of the Google Project region (e.g., us-central1) | str | `--googlebatch-region` |`SNAKEMAKE_GOOGLEBATCH_REGION` | True | unset |
-| machine_type | Google Cloud machine type or VM (mpitune configurations are on c2 and c2d family) | str | `--googlebatch-machine-type` | | False | c2-standard-4 |
-| image_family | Google Cloud image family (defaults to hpc-centos-7) | `--googlebatch-image-family` | str | | False | hpc-centos-7 |
-| image_project | The project the selected image belongs to (defaults to cloud-hpc-image-public) | `--googlebatch-image-project` | str | | False | cloud-hpc-image-public |
-| bucket | A bucket to mount with snakemake data | `--googlebatch-bucket` | str | `SNAKEMAKE_GOOGLEBATCH_BUCKET` | True |  unset |
-| mount_path | The mount path for a bucket (if provided) | `--googlebatch-mount-path` | str | | False | /mnt/share |
-| work_tasks | The default number of work tasks (these are NOT MPI ranks) | `--googlebatch-work-tasks` | int | | False | 1 |
-| cpu_milli | Milliseconds per cpu-second | `--googlebatch-cpu-milli` | int | | False | 1000 |
-| work_tasks_per_node | The default number of work tasks per node (Google Batch calls these tasks) | `--googlebatch-work-tasks-per-node` | int | | False | 1 |
-| memory | Memory in MiB | `--googlebatch-memory` | int | | False | 1000 |
-| retry_count | Retry count (default to 1) | `--googlebatch-retry-count` | int | | False | 1 |
-| max_run_duration | Maximum run duration, string (e.g., 3600s) | `--googlebatch-max-run-duration` | str | | False | "3600s" |
-| labels | Comma separated key value pairs to label job (e.g., model=a3,stage=test) |`--googlebatch-labels` | str | | False | unset|
-| container | Container to use (only when image_family is batch-cos*) [see here](https://cloud.google.com/batch/docs/vm-os-environment-overview#supported_vm_os_images) for families/projects | `--googlebatch-container` | str | | False | unset|
-| keep_source_cache | Cache workflows in your Google Cloud Storage Bucket | `--googlebatch-keep-source-cache` | bool | | False | False |
-| snippet | A comma separated list of one or more snippets to add to your setup | `--googelbatch-snippets` | str | | False | unset |
-
-For machine type, note that for MPI workloads, mpitune configurations are validated on c2 and c2d instances only.
-Also note that you can customize the machine type on the level of the step (see [Step Options](#step-options) below).
-
 #### Choosing an Image
 
 You can read about how to choose an image [here](https://cloud.google.com/batch/docs/view-os-images). Note that
@@ -152,6 +125,9 @@ rule hello_world:
         "..."
 ```
 
+Note that for MPI workloads, mpitune configurations are validated on c2 and c2d instances only.
+
+
 #### googlebatch_image_family
 
 This will define the image family for a particular step, overriding the default from the command line.
@@ -225,6 +201,34 @@ rule hello_world:
         "..."
 ```
 
+#### googlebatch_network
+
+The URL of an existing network resource (e.g., `projects/{project}/global/networks/{network}`)
+
+```console
+rule hello_world:
+	output:
+		"...",
+	resources: 
+		googlebatch_network="projects/{project}/global/networks/{network}"
+	shell:
+        "..."
+```
+
+#### googlebatch_subnetwork
+
+The URL of an existing subnetwork resource (e.g., `projects/{project}/regions/{region}/subnetworks/{subnetwork}`)
+
+```console
+rule hello_world:
+	output:
+		"...",
+	resources: 
+		googlebatch_subnetwork="projects/{project}/regions/{region}/subnetworks/{subnetwork}"
+	shell:
+        "..."
+```
+
 #### googlebatch_cpu_milli
 
 This will define the milliseconds per cpu-second for a particular step, overriding the default from the command line.

snakemake_executor_plugin_googlebatch/__init__.py

@@ -110,6 +110,24 @@ class ExecutorSettings(ExecutorSettingsBase):
         },
     )
 
+    network: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": "The URL of an existing network resource",
+            "env_var": False,
+            "required": False,
+        },
+    )
+
+    subnetwork: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": "The URL of an existing subnetwork resource",
+            "env_var": False,
+            "required": False,
+        },
+    )
+
     # local SSD uses type "local-ssd".
     # Also "pd-balanced", "pd-extreme", "pd-ssd", "pd-standard"
     boot_disk_type: Optional[str] = field(

snakemake_executor_plugin_googlebatch/executor.py

@@ -306,16 +306,39 @@ def get_allocation_policy(self, job):
 
         instances.policy = policy
         allocation_policy.instances = [instances]
+
+        # Add custom network interfaces
+        network_policy = self.get_network_policy(job)
+        if network_policy is not None:
+            allocation_policy.network = network_policy
         return allocation_policy
 
+    def get_network_policy(self, job):
+        """
+        Given a job request, get the network policy
+        """
+        network = self.get_param(job, "network")
+        subnetwork = self.get_param(job, "subnetwork")
+        if all(x is None for x in [network, subnetwork]):
+            return
+
+        policy = batch_v1.AllocationPolicy.NetworkPolicy()
+        interface = batch_v1.AllocationPolicy.NetworkInterface()
+        if network is not None:
+            interface.network = network
+        if subnetwork is not None:
+            interface.subnetwork = subnetwork
+        policy.network_interfaces = [interface]
+        return policy
+
     def get_boot_disk(self, job):
         """
         Given a job request, add a customized boot disk.
         """
         # Reference disk, boot disk type, and size
-        image = job.resources.get("googlebatch_boot_disk_image")
-        size = job.resources.get("googlebatch_boot_disk_gb")
-        typ = job.resources.get("googlebatch_boot_disk_type")
+        image = self.get_param(job, "boot_disk_image")
+        size = self.get_param(job, "boot_disk_gb")
+        typ = self.get_param(job, "boot_disk_type")
 
         # Cut out early if no customization
         if all(x is None for x in [image, size, typ]):