325 lines
11 KiB
Python
325 lines
11 KiB
Python
# -*- coding: utf-8 -*- #
|
|
# Copyright 2025 Google LLC. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
"""Advice group command flags."""
|
|
|
|
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import unicode_literals
|
|
|
|
from googlecloudsdk.api_lib.compute import utils
|
|
from googlecloudsdk.calliope import arg_parsers
|
|
from googlecloudsdk.calliope import exceptions
|
|
from googlecloudsdk.command_lib.compute import flags as compute_flags
|
|
from googlecloudsdk.command_lib.compute.reservations import flags as reservation_flags
|
|
from googlecloudsdk.command_lib.util.apis import arg_utils
|
|
|
|
|
|
def AddRegionFlag(parser):
|
|
"""Add the --region flag."""
|
|
compute_flags.AddRegionFlag(
|
|
parser=parser,
|
|
resource_type="resources",
|
|
operation_type="get advice on")
|
|
|
|
|
|
def AddLocationPolicyFlag(parser):
|
|
"""Add the --location_policy flag."""
|
|
parser.add_argument(
|
|
"--location-policy",
|
|
metavar="ZONE=POLICY",
|
|
type=arg_parsers.ArgDict(),
|
|
help="""
|
|
Policy for which zones to include or exclude when looking for the optimal
|
|
time window and zone for Future Reservations within a region. Policy is
|
|
defined as a list of key-value pairs, with the key being the zone name,
|
|
and value being the applied policy. Available policies are `allow` and
|
|
`deny`. Default for zones if left unspecified is `allow`.
|
|
|
|
Example:
|
|
|
|
gcloud compute advice calendar-mode --location-policy=us-central1-a=allow,us-central1-b=deny
|
|
""",
|
|
)
|
|
|
|
|
|
def AddStartTimeRangeFlag(parser):
|
|
"""Add the --start-time-range flag."""
|
|
parser.add_argument(
|
|
"--start-time-range",
|
|
type=arg_parsers.ArgDict(
|
|
spec={
|
|
"from": arg_parsers.Datetime.Parse,
|
|
"to": arg_parsers.Datetime.Parse,
|
|
},
|
|
max_length=2,
|
|
),
|
|
help="""
|
|
A time range for the start time of the Future Reservation. Defined as
|
|
a list of key-value pairs.
|
|
|
|
The key is either "from" or "to", and the value is a datetime.
|
|
See $ gcloud topic datetimes for information on time formats.
|
|
|
|
*from*:: The earliest possible start time for the reservation.
|
|
*to*:: The latest possible start time for the reservation.
|
|
|
|
Example:
|
|
|
|
gcloud compute advice calendar-mode --start-time-range=from=2024-08-01T00:00:00Z,to=2024-08-02T00:00:00Z
|
|
""",
|
|
)
|
|
|
|
|
|
def AddEndTimeRangeFlag(parser):
|
|
"""Add the --end-time-range flag."""
|
|
parser.add_argument(
|
|
"--end-time-range",
|
|
type=arg_parsers.ArgDict(
|
|
spec={
|
|
"from": arg_parsers.Datetime.Parse,
|
|
"to": arg_parsers.Datetime.Parse,
|
|
},
|
|
max_length=2,
|
|
),
|
|
help="""
|
|
A time range for the end time of the Future Reservation. Defined as
|
|
a list of key-value pairs.
|
|
|
|
The key is either "from" or "to", and the value is a datetime.
|
|
See $ gcloud topic datetimes for information on time formats.
|
|
|
|
*from*:: The earliest possible end time for the reservation.
|
|
*to*:: The latest possible end time for the reservation.
|
|
|
|
Example:
|
|
|
|
gcloud compute advice calendar-mode --end-time-range=from=2024-08-01T00:00:00Z,to=2024-08-02T00:00:00Z
|
|
""",
|
|
)
|
|
|
|
|
|
def AddDurationRangeFlag(parser):
|
|
"""Add the --duration-range flag."""
|
|
parser.add_argument(
|
|
"--duration-range",
|
|
type=arg_parsers.ArgDict(
|
|
spec={
|
|
"min": arg_parsers.Duration(),
|
|
"max": arg_parsers.Duration(),
|
|
},
|
|
max_length=2,
|
|
),
|
|
help="""
|
|
A duration range for the duration of the Future Reservation. Defined as
|
|
a list of key-value pairs.
|
|
|
|
The key is either "min" or "max", and the value is a duration in seconds.
|
|
For example, specify `30m` for a duration of 30 minutes or specify
|
|
`1d2h3m4s` for a duration of 1 day, 2 hours, 3 minutes, and 4 seconds.
|
|
|
|
See $ gcloud topic datetimes for information on duration formats.
|
|
|
|
*min*:: The minimum duration of the Future Reservation.
|
|
*max*:: The maximum duration of the Future Reservation.
|
|
|
|
Example:
|
|
|
|
gcloud compute advice calendar-mode --duration-range=min=24h,max=72h
|
|
""",
|
|
)
|
|
|
|
|
|
def AddSkuPropertiesFlags(accelerator_properties_group):
|
|
"""Add the SKU properties flags."""
|
|
|
|
instance_properties_group = accelerator_properties_group.add_group(help="""
|
|
Define individual instance properties for the specific SKU reservation.
|
|
""")
|
|
|
|
# add --machine-type
|
|
reservation_flags.GetMachineType(True).AddToParser(instance_properties_group)
|
|
|
|
# add --vm-count
|
|
vm_count_flag = reservation_flags.GetVmCountFlag(True)
|
|
vm_count_flag.kwargs["help"] = (
|
|
"""The number of instances to check for availability."""
|
|
)
|
|
vm_count_flag.AddToParser(instance_properties_group)
|
|
|
|
# add --local-ssd
|
|
local_ssd_flag = reservation_flags.GetLocalSsdFlag()
|
|
local_ssd_flag.kwargs["help"] = """\
|
|
Manage the size and the interface of local SSD to use. See
|
|
https://cloud.google.com/compute/docs/disks/local-ssd for more information.
|
|
*interface*::: The kind of disk interface exposed to the VM for this SSD.
|
|
The only valid value is `nvme`.
|
|
*size*::: The size of the local SSD in base-2 GB."""
|
|
local_ssd_flag.AddToParser(instance_properties_group)
|
|
|
|
|
|
def AddAggregatePropertiesFlags(accelerator_properties_group):
|
|
"""Add the aggregate properties flags."""
|
|
|
|
aggregate_properties_group = accelerator_properties_group.add_group(help="""
|
|
You must define the version and number of TPUs to reserve.
|
|
""")
|
|
|
|
# add --tpu-version
|
|
reservation_flags.GetTpuVersion(True).AddToParser(aggregate_properties_group)
|
|
|
|
# add --chip-count
|
|
chip_count_flag = reservation_flags.GetChipCount(True)
|
|
chip_count_flag.kwargs["help"] = (
|
|
"""The number of chips to check for availability."""
|
|
)
|
|
chip_count_flag.AddToParser(aggregate_properties_group)
|
|
|
|
# add --workload-type
|
|
workload_type_flag = reservation_flags.GetWorkloadType(False)
|
|
workload_type_flag.kwargs["help"] = (
|
|
"""Type of the workload that will be running on the reserved TPUs."""
|
|
)
|
|
workload_type_flag.kwargs["choices"]["SERVING"] = (
|
|
"Reserved resources will be optimized for SERVING workloads that"
|
|
" handle concurrent requests and require minimal network latency,"
|
|
" such as ML inference."
|
|
)
|
|
workload_type_flag.kwargs["choices"]["BATCH"] = (
|
|
"Reserved resources will be optimized for BATCH workloads that"
|
|
" handle large amounts of data in single or multiple operations,"
|
|
" such as ML training workloads."
|
|
)
|
|
workload_type_flag.AddToParser(aggregate_properties_group)
|
|
|
|
|
|
def AddAcceleratorPropertiesFlags(parser):
|
|
"""Add the accelerator properties flags."""
|
|
accelerator_properties_group = parser.add_group(
|
|
help="""
|
|
Specify the properties of the resources that you want to view the availability of.
|
|
""",
|
|
mutex=True,
|
|
required=True,
|
|
)
|
|
|
|
AddSkuPropertiesFlags(accelerator_properties_group)
|
|
AddAggregatePropertiesFlags(accelerator_properties_group)
|
|
|
|
|
|
def AddDeploymentTypeFlag(parser):
|
|
"""Add the --deployment-type flag."""
|
|
|
|
# Calendar Mode currently only supports Dense deployment as a MVP; this will
|
|
# be the default. Once the Flexible deployment is supported we will use the
|
|
# deployment-type flag from the Future Reservations.
|
|
|
|
help_text = """\
|
|
The deployment type for the reserved capacity.
|
|
"""
|
|
parser.add_argument(
|
|
"--deployment-type",
|
|
choices={
|
|
"DENSE": "DENSE mode is for densely deployed reservation blocks.",
|
|
},
|
|
default="DENSE",
|
|
help=help_text,
|
|
hidden=True,
|
|
)
|
|
|
|
|
|
def AddProvisioningModelFlag(parser):
|
|
"""Add the --provisioning-model flag."""
|
|
provisioning_model_choices = {
|
|
"SPOT": (
|
|
"Compute Engine may preempt a Spot VM whenever it needs capacity. "
|
|
"Because Spot VMs don't have a guaranteed runtime, they come at a "
|
|
"discounted price."
|
|
),
|
|
}
|
|
parser.add_argument(
|
|
"--provisioning-model",
|
|
choices=provisioning_model_choices,
|
|
type=arg_utils.ChoiceToEnumName,
|
|
required=True,
|
|
help="""
|
|
Specifies the provisioning model.
|
|
""",
|
|
)
|
|
|
|
|
|
def AddTargetDistributionShapeFlag(parser):
|
|
"""Add the --target-distribution-shape flag."""
|
|
help_text = """\
|
|
Defines the distribution requirement for the requested VMs.
|
|
"""
|
|
choices = {
|
|
"ANY": (
|
|
"When you specify ANY for VM instance creation across multiple"
|
|
" zones, you specify that you want to create your VM instances in one"
|
|
" or more zones based on resource availability, and to use any"
|
|
" unused, matching reservation in your project. Use ANY for batch"
|
|
" workloads that don't require high availability."
|
|
),
|
|
"ANY_SINGLE_ZONE": (
|
|
"When you specify ANY_SINGLE_ZONE for VM instance creation, you"
|
|
" specify that you want to create all VM instances in a single"
|
|
" available zone. Compute Engine selects this zone based on resource"
|
|
" availability and on any unused, matching reservation in your"
|
|
" project. However, zonal resource availability constraints might"
|
|
" prevent Compute Engine from creating all your requested VM"
|
|
" instances. Use ANY_SINGLE_ZONE when the VM instances in your"
|
|
" workloads need to frequently communicate among each other."
|
|
),
|
|
}
|
|
parser.add_argument(
|
|
"--target-distribution-shape",
|
|
choices=choices,
|
|
type=arg_utils.ChoiceToEnumName,
|
|
required=True,
|
|
help=help_text)
|
|
|
|
|
|
def ValidateZonesAndRegionFlags(args, resources):
|
|
"""Validate --zones and --region flags."""
|
|
if not args.zones:
|
|
return
|
|
|
|
ignored_required_params = {"project": "fake"}
|
|
zone_names = []
|
|
for zone in args.zones:
|
|
zone_ref = resources.Parse(
|
|
zone, collection="compute.zones", params=ignored_required_params
|
|
)
|
|
zone_names.append(zone_ref.Name())
|
|
|
|
zone_regions = set([utils.ZoneNameToRegionName(z) for z in zone_names])
|
|
if len(zone_regions) > 1:
|
|
raise exceptions.InvalidArgumentException(
|
|
"--zones", "All zones must be in the same region."
|
|
)
|
|
elif len(zone_regions) == 1 and args.region:
|
|
zone_region = zone_regions.pop()
|
|
region_ref = resources.Parse(
|
|
args.region,
|
|
collection="compute.regions",
|
|
params=ignored_required_params,
|
|
)
|
|
region = region_ref.Name()
|
|
if zone_region != region:
|
|
raise exceptions.InvalidArgumentException(
|
|
"--zones",
|
|
"Specified zones are not in specified region [{0}].".format(region),
|
|
)
|