feat: add support for placement policies (#110)

* feat: add support for placement policies

feat: per-Runnable labels
PiperOrigin-RevId: 532437427

Source-Link: googleapis/googleapis@288aa7f

Source-Link: https://github.com/googleapis/googleapis-gen/commit/615d60f8e98ef120a79410c029b4c72448b53bde
Copy-Tag: eyJwIjoiLmdpdGh1Yi8uT3dsQm90LnlhbWwiLCJoIjoiNjE1ZDYwZjhlOThlZjEyMGE3OTQxMGMwMjliNGM3MjQ0OGI1M2JkZSJ9

* 🦉 Updates from OwlBot post-processor

See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md

---------

Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>

Author: gcf-owl-bot[bot]

google/cloud/batch_v1/types/job.py

@@ -403,6 +403,8 @@ class AllocationPolicy(proto.Message):
             reserved.
         network (google.cloud.batch_v1.types.AllocationPolicy.NetworkPolicy):
             The network policy.
+        placement (google.cloud.batch_v1.types.AllocationPolicy.PlacementPolicy):
+            The placement policy.
     """
 
     class ProvisioningModel(proto.Enum):
@@ -435,13 +437,16 @@ class LocationPolicy(proto.Message):
         Attributes:
             allowed_locations (MutableSequence[str]):
                 A list of allowed location names represented by internal
-                URLs. Each location can be a region or a zone. Only one
-                region or multiple zones in one region is supported now. For
-                example, ["regions/us-central1"] allow VMs in any zones in
-                region us-central1. ["zones/us-central1-a",
-                "zones/us-central1-c"] only allow VMs in zones us-central1-a
-                and us-central1-c. All locations end up in different regions
-                would cause errors. For example, ["regions/us-central1",
+                URLs.
+
+                Each location can be a region or a zone. Only one region or
+                multiple zones in one region is supported now. For example,
+                ["regions/us-central1"] allow VMs in any zones in region
+                us-central1. ["zones/us-central1-a", "zones/us-central1-c"]
+                only allow VMs in zones us-central1-a and us-central1-c.
+
+                All locations end up in different regions would cause
+                errors. For example, ["regions/us-central1",
                 "zones/us-central1-a", "zones/us-central1-b",
                 "zones/us-west1-a"] contains 2 regions "us-central1" and
                 "us-west1". An error is expected in this case.
@@ -455,7 +460,8 @@ class LocationPolicy(proto.Message):
     class Disk(proto.Message):
         r"""A new persistent disk or a local ssd.
         A VM can only have one local SSD setting but multiple local SSD
-        partitions. https://cloud.google.com/compute/docs/disks#pdspecs.
+        partitions. See
+        https://cloud.google.com/compute/docs/disks#pdspecs and
         https://cloud.google.com/compute/docs/disks#localssds.
 
         This message has `oneof`_ fields (mutually exclusive fields).
@@ -468,20 +474,24 @@ class Disk(proto.Message):
         Attributes:
             image (str):
                 Name of a public or custom image used as the data source.
-                For example, the following are all valid URLs: (1) Specify
-                the image by its family name:
-                projects/{project}/global/images/family/{image_family} (2)
-                Specify the image version:
-                projects/{project}/global/images/{image_version} You can
-                also use Batch customized image in short names. The
+                For example, the following are all valid URLs:
+
+                -  Specify the image by its family name:
+                   projects/{project}/global/images/family/{image_family}
+                -  Specify the image version:
+                   projects/{project}/global/images/{image_version}
+
+                You can also use Batch customized image in short names. The
                 following image values are supported for a boot disk:
-                "batch-debian": use Batch Debian images. "batch-centos": use
-                Batch CentOS images. "batch-cos": use Batch
-                Container-Optimized images.
+
+                -  "batch-debian": use Batch Debian images.
+                -  "batch-centos": use Batch CentOS images.
+                -  "batch-cos": use Batch Container-Optimized images.
 
                 This field is a member of `oneof`_ ``data_source``.
             snapshot (str):
                 Name of a snapshot used as the data source.
+                Snapshot is not supported as boot disk now.
 
                 This field is a member of `oneof`_ ``data_source``.
             type_ (str):
@@ -490,11 +500,13 @@ class Disk(proto.Message):
                 disks and boot disks use "pd-balanced", "pd-extreme",
                 "pd-ssd" or "pd-standard".
             size_gb (int):
-                Disk size in GB. For persistent disk, this field is ignored
-                if ``data_source`` is ``image`` or ``snapshot``. For local
-                SSD, size_gb should be a multiple of 375GB, otherwise, the
-                final size will be the next greater multiple of 375 GB. For
-                boot disk, Batch will calculate the boot disk size based on
+                Disk size in GB.
+
+                For persistent disk, this field is ignored if
+                ``data_source`` is ``image`` or ``snapshot``. For local SSD,
+                size_gb should be a multiple of 375GB, otherwise, the final
+                size will be the next greater multiple of 375 GB. For boot
+                disk, Batch will calculate the boot disk size based on
                 source image and task requirements if you do not speicify
                 the size. If both this field and the boot_disk_mib field in
                 task spec's compute_resource are defined, Batch will only
@@ -607,18 +619,20 @@ class InstancePolicy(proto.Message):
             machine_type (str):
                 The Compute Engine machine type.
             min_cpu_platform (str):
-                The minimum CPU platform. See
-                ``https://cloud.google.com/compute/docs/instances/specify-min-cpu-platform``.
+                The minimum CPU platform.
+                See
+                https://cloud.google.com/compute/docs/instances/specify-min-cpu-platform.
                 Not yet implemented.
             provisioning_model (google.cloud.batch_v1.types.AllocationPolicy.ProvisioningModel):
                 The provisioning model.
             accelerators (MutableSequence[google.cloud.batch_v1.types.AllocationPolicy.Accelerator]):
                 The accelerators attached to each VM
                 instance.
             boot_disk (google.cloud.batch_v1.types.AllocationPolicy.Disk):
-                Book disk to be created and attached to each
+                Boot disk to be created and attached to each
                 VM by this InstancePolicy. Boot disk will be
-                deleted when the VM is deleted.
+                deleted when the VM is deleted. Batch API now
+                only supports booting from image.
             disks (MutableSequence[google.cloud.batch_v1.types.AllocationPolicy.AttachedDisk]):
                 Non-boot disks to be attached for each VM
                 created by this InstancePolicy. New disks will
@@ -705,21 +719,23 @@ class NetworkInterface(proto.Message):
 
         Attributes:
             network (str):
-                The URL of an existing network resource.
-                You can specify the network as a full or partial
-                URL. For example, the following are all valid
-                URLs:
-                https://www.googleapis.com/compute/v1/projects/{project}/global/networks/{network}
-                projects/{project}/global/networks/{network}
-                global/networks/{network}
+                The URL of an existing network resource. You can specify the
+                network as a full or partial URL.
+
+                For example, the following are all valid URLs:
+
+                -  https://www.googleapis.com/compute/v1/projects/{project}/global/networks/{network}
+                -  projects/{project}/global/networks/{network}
+                -  global/networks/{network}
             subnetwork (str):
-                The URL of an existing subnetwork resource in
-                the network. You can specify the subnetwork as a
-                full or partial URL. For example, the following
-                are all valid URLs:
-                https://www.googleapis.com/compute/v1/projects/{project}/regions/{region}/subnetworks/{subnetwork}
-                projects/{project}/regions/{region}/subnetworks/{subnetwork}
-                regions/{region}/subnetworks/{subnetwork}
+                The URL of an existing subnetwork resource in the network.
+                You can specify the subnetwork as a full or partial URL.
+
+                For example, the following are all valid URLs:
+
+                -  https://www.googleapis.com/compute/v1/projects/{project}/regions/{region}/subnetworks/{subnetwork}
+                -  projects/{project}/regions/{region}/subnetworks/{subnetwork}
+                -  regions/{region}/subnetworks/{subnetwork}
             no_external_ip_address (bool):
                 Default is false (with an external IP
                 address). Required if no external public IP
@@ -762,6 +778,36 @@ class NetworkPolicy(proto.Message):
             message="AllocationPolicy.NetworkInterface",
         )
 
+    class PlacementPolicy(proto.Message):
+        r"""PlacementPolicy describes a group placement policy for the
+        VMs controlled by this AllocationPolicy.
+
+        Attributes:
+            collocation (str):
+                UNSPECIFIED vs. COLLOCATED (default
+                UNSPECIFIED). Use COLLOCATED when you want VMs
+                to be located close to each other for low
+                network latency between the VMs. No placement
+                policy will be generated when collocation is
+                UNSPECIFIED.
+            max_distance (int):
+                When specified, causes the job to fail if more than
+                max_distance logical switches are required between VMs.
+                Batch uses the most compact possible placement of VMs even
+                when max_distance is not specified. An explicit max_distance
+                makes that level of compactness a strict requirement. Not
+                yet implemented
+        """
+
+        collocation: str = proto.Field(
+            proto.STRING,
+            number=1,
+        )
+        max_distance: int = proto.Field(
+            proto.INT64,
+            number=2,
+        )
+
     location: LocationPolicy = proto.Field(
         proto.MESSAGE,
         number=1,
@@ -787,6 +833,11 @@ class NetworkPolicy(proto.Message):
         number=7,
         message=NetworkPolicy,
     )
+    placement: PlacementPolicy = proto.Field(
+        proto.MESSAGE,
+        number=10,
+        message=PlacementPolicy,
+    )
 
 
 class TaskGroup(proto.Message):
@@ -804,10 +855,11 @@ class TaskGroup(proto.Message):
             task spec.
         task_count (int):
             Number of Tasks in the TaskGroup.
-            default is 1
+            Default is 1.
         parallelism (int):
             Max number of tasks that can run in parallel. Default to
-            min(task_count, 1000).
+            min(task_count, 1000). Field parallelism must be 1 if the
+            scheduling_policy is IN_ORDER.
         task_environments (MutableSequence[google.cloud.batch_v1.types.Environment]):
             An array of environment variable mappings, which are passed
             to Tasks with matching indices. If task_environments is used

google/cloud/batch_v1/types/task.py

@@ -138,7 +138,7 @@ class State(proto.Enum):
 
         Values:
             STATE_UNSPECIFIED (0):
-                unknown state
+                Unknown state.
             PENDING (1):
                 The Task is created and waiting for
                 resources.
@@ -150,13 +150,17 @@ class State(proto.Enum):
                 The Task has failed.
             SUCCEEDED (5):
                 The Task has succeeded.
+            UNEXECUTED (6):
+                The Task has not been executed when the Job
+                finishes.
         """
         STATE_UNSPECIFIED = 0
         PENDING = 1
         ASSIGNED = 2
         RUNNING = 3
         FAILED = 4
         SUCCEEDED = 5
+        UNEXECUTED = 6
 
     state: State = proto.Field(
         proto.ENUM,
@@ -222,6 +226,8 @@ class Runnable(proto.Message):
             TaskGroup).
         timeout (google.protobuf.duration_pb2.Duration):
             Timeout for this Runnable.
+        labels (MutableMapping[str, str]):
+            Labels for this Runnable.
     """
 
     class Container(proto.Message):
@@ -247,10 +253,11 @@ class Container(proto.Message):
                 the "docker run" command when running this
                 container, e.g. "--network host".
             block_external_network (bool):
-                If set to true, external network access to
-                and from container will be blocked. The
-                container will use the default internal network
-                'goog-internal'.
+                If set to true, external network access to and from
+                container will be blocked, containers that are with
+                block_external_network as true can still communicate with
+                each other, network cannot be specified in the
+                ``container.options`` field.
             username (str):
                 Optional username for logging in to a docker registry. If
                 username matches ``projects/*/secrets/*/versions/*`` then
@@ -399,6 +406,11 @@ class Barrier(proto.Message):
         number=8,
         message=duration_pb2.Duration,
     )
+    labels: MutableMapping[str, str] = proto.MapField(
+        proto.STRING,
+        proto.STRING,
+        number=9,
+    )
 
 
 class TaskSpec(proto.Message):

tests/unit/gapic/batch_v1/test_batch_service.py

@@ -2549,6 +2549,7 @@ def test_create_job_rest(request_type):
                                 },
                             },
                             "timeout": {"seconds": 751, "nanos": 543},
+                            "labels": {},
                         }
                     ],
                     "compute_resource": {
@@ -2640,6 +2641,7 @@ def test_create_job_rest(request_type):
                     }
                 ]
             },
+            "placement": {"collocation": "collocation_value", "max_distance": 1264},
         },
         "labels": {},
         "status": {
@@ -2896,6 +2898,7 @@ def test_create_job_rest_bad_request(
                                 },
                             },
                             "timeout": {"seconds": 751, "nanos": 543},
+                            "labels": {},
                         }
                     ],
                     "compute_resource": {
@@ -2987,6 +2990,7 @@ def test_create_job_rest_bad_request(
                     }
                 ]
             },
+            "placement": {"collocation": "collocation_value", "max_distance": 1264},
         },
         "labels": {},
         "status": {