Skip to content

Commit

Permalink
feat: add support for placement policies (#110)
Browse files Browse the repository at this point in the history
* feat: add support for placement policies

feat: per-Runnable labels
PiperOrigin-RevId: 532437427

Source-Link: googleapis/googleapis@288aa7f

Source-Link: googleapis/googleapis-gen@615d60f
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>
  • Loading branch information
gcf-owl-bot[bot] and gcf-owl-bot[bot] authored May 16, 2023
1 parent e279275 commit 09ca683
Show file tree
Hide file tree
Showing 3 changed files with 115 additions and 47 deletions.
136 changes: 94 additions & 42 deletions packages/google-cloud-batch/google/cloud/batch_v1/types/job.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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.
Expand All @@ -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).
Expand All @@ -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):
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand All @@ -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):
Expand All @@ -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
Expand Down
22 changes: 17 additions & 5 deletions packages/google-cloud-batch/google/cloud/batch_v1/types/task.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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,
Expand Down Expand Up @@ -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):
Expand All @@ -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
Expand Down Expand Up @@ -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):
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -2549,6 +2549,7 @@ def test_create_job_rest(request_type):
},
},
"timeout": {"seconds": 751, "nanos": 543},
"labels": {},
}
],
"compute_resource": {
Expand Down Expand Up @@ -2640,6 +2641,7 @@ def test_create_job_rest(request_type):
}
]
},
"placement": {"collocation": "collocation_value", "max_distance": 1264},
},
"labels": {},
"status": {
Expand Down Expand Up @@ -2896,6 +2898,7 @@ def test_create_job_rest_bad_request(
},
},
"timeout": {"seconds": 751, "nanos": 543},
"labels": {},
}
],
"compute_resource": {
Expand Down Expand Up @@ -2987,6 +2990,7 @@ def test_create_job_rest_bad_request(
}
]
},
"placement": {"collocation": "collocation_value", "max_distance": 1264},
},
"labels": {},
"status": {
Expand Down

0 comments on commit 09ca683

Please sign in to comment.