# Copyright 2021 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.

accelerator_type:
  api_field: node.acceleratorType
  arg_name: accelerator-type
  required: false
  help_text: |
    TPU accelerator type for the TPU.
    If not specified, this defaults to `v2-8`.

    For a list of available accelerator types run:

    `{parent_command} accelerator-types list`
  type: googlecloudsdk.command_lib.util.hooks.types:LowerCaseType
  default: 'v2-8'

description:
  api_field: node.description
  arg_name: description
  required: false
  help_text: |
    Text description of the TPU.

network:
  api_field: node.networkConfig.network
  arg_name: network
  required: false
  help_text: |
    Network that this TPU will be a part of.
  default: 'default'

subnetwork:
  api_field: node.networkConfig.subnetwork
  arg_name: subnetwork
  required: false
  help_text: |
    Subnetwork that this TPU will be a part of.

enable_external_ips:
  api_field: node.networkConfig.enableExternalIps
  arg_name: internal-ips
  required: false
  default: false
  type: bool
  processor: googlecloudsdk.command_lib.compute.tpus.tpu_vm.util:InvertBoolean
  help_text: |
    Indicate that the IP addresses for the node should be internal. The default is that external IP
    addresses will be associated with the TPU workers.

queue_count:
  api_field: node.networkConfig.queueCount
  arg_name: queue-count
  type: int
  required: false
  help_text: |
    Specifies the networking queue count for TPU VM instances. Both Rx and Tx queues will be set to
    this number. If it's not specified, a default queue count will be assigned. For Virtio-net,
    each interface will get min(floor(#vCPU / #vNIC), 32) queues. For gVNIC, each interface will
    get min(floor(#vCPU / #vNIC / 2), 16) queues.

version:
  api_field: node.runtimeVersion
  arg_name: version
  required: true
  help_text: |
    Runtime version for the TPU, such as `2.3`.

    For a list of available versions run:

    `{parent_command} versions list`

range: &range
  api_field: node.cidrBlock
  arg_name: range
  help_text: |
    CIDR Range for the TPU.

    The IP range that the TPU will select an IP address from.
    Must be in CIDR notation and a `/29` range, for example
    `192.168.0.0/29`. Errors will occur if the CIDR range has already been
    used for a currently existing TPU, the CIDR range conflicts with any
    networks in the user's provided network, or the provided network is
    peered with another network that is using that CIDR range.

preemptible:
  api_field: node.schedulingConfig.preemptible
  arg_name: preemptible
  type: bool
  default: false
  required: false
  help_text: |
    If provided, the TPU will be preemptible and time-limited. It may be
    preempted to free up resources for standard TPUs, and will only be able
    to run for a limited amount of time.

    Preemptible TPUs cannot be restarted.

spot:
  api_field: node.schedulingConfig.spot
  arg_name: spot
  type: bool
  default: false
  required: false
  help_text: |
    If specified, create this VM as a spot VM.
    Spot VMs make unused capacity available at highly discounted rates.
    Spot VMs may be preempted at any time if the capacity is needed, but unless preempted there is
    no limit on runtime duration. Spot VM TPUs cannot be restarted, and must be recreated again.

reserved:
  api_field: node.schedulingConfig.reserved
  arg_name: reserved
  type: bool
  default: false
  required: false
  help_text: |
    When specified, will attempt to create the TPU node under reservations made in the current
    project. The reservations can be made separately but used in aggregated form. i.e., the user can
    make a reservation of 128 V2 TPUs and later on make another reservation of 128 V2 TPUs then
    creates a v2-256 TPU instance. If there exists no reservation or not sufficient amount of
    reserved cores under the project, the request will fail due to lack of capacity.

service_account:
  api_field: node.serviceAccount.email
  arg_name: service-account
  required: false
  help_text: |
    Email address of the service account. If empty, default Google Compute Engine service
    account will be used.

service_account_scopes:
  api_field: node.serviceAccount.scope
  arg_name: scopes
  type: "googlecloudsdk.calliope.arg_parsers:ArgList:"
  required: false
  help_text: |
    List of comma-separated scopes to be made available for the service account.

labels:
  api_field: node.labels
  arg_name: labels
  metavar: KEY=VALUE
  type: "googlecloudsdk.calliope.arg_parsers:ArgDict:"
  required: false
  help_text: |
    Resource labels to represent user-provided metadata. See
    https://cloud.google.com/compute/docs/labeling-resources for details.

update_labels:
  arg_name: update-labels
  type: "googlecloudsdk.calliope.arg_parsers:ArgDict:"
  required: false
  metavar: KEY=VALUE
  help_text: |
    Resource labels to update that represent user-provided metadata. If a label exists, its value is
    modified. Otherwise, a new label is created.
    See https://cloud.google.com/compute/docs/labeling-resources for details.

clear_labels:
  arg_name: clear-labels
  type: bool
  default: false
  required: false
  help_text: |
    Remove all labels. If `--update-labels` is also specified then `--clear-labels` is applied
    first.

    For example, to remove all labels:

      $ {command} --clear-labels

    To remove all existing labels and create two new labels, 'foo' and 'baz':

      $ {command} --clear-labels --update-labels foo=bar,baz=qux

remove_labels:
  arg_name: remove-labels
  type: "googlecloudsdk.calliope.arg_parsers:ArgList:"
  required: false
  metavar: KEY
  help_text: |
    List of label keys to remove. If a label does not exist it is silently ignored. If
    `--update-labels` is also specified then `--update-labels` is applied first.

tags:
  api_field: node.tags
  arg_name: tags
  type: "googlecloudsdk.calliope.arg_parsers:ArgList:"
  required: false
  help_text: |
    Tags to apply to the TPU Node. Tags are used to identify valid sources or
    targets for network firewalls. See https://cloud.google.com/vpc/docs/add-remove-network-tags for
    more details.

add_tags:
  arg_name: add-tags
  type: "googlecloudsdk.calliope.arg_parsers:ArgList:"
  required: false
  metavar: TAGS
  help_text: |
    Tags to add to the TPU Node. Tags are used to identify valid sources or
    targets for network firewalls. See https://cloud.google.com/vpc/docs/add-remove-network-tags for
    more details.

clear_tags:
  arg_name: clear-tags
  type: bool
  default: false
  required: false
  help_text: |
    Remove all tags. If `--add-tags` is also specified then `--clear-tags` is applied
    first.

    For example, to remove all tags:

      $ {command} --clear-tags

    To remove all existing tags and create two new tags, 'foo' and 'bar':

      $ {command} --clear-tags --add-tags foo,bar

remove_tags:
  arg_name: remove-tags
  type: "googlecloudsdk.calliope.arg_parsers:ArgList:"
  required: false
  metavar: TAG
  help_text: |
    List of tags to remove. If a tag does not exist it is silently ignored. If
    `--add-tags` is also specified then `--add-tags` is applied first.

data_disks:
  api_field: node.dataDisks
  arg_name: data-disk
  type:
    arg_dict:
      flatten: false
      spec:
      - api_field: sourceDisk
        arg_name: source
        type: str
        required: true
      - api_field: mode
        arg_name: mode
        type: str
        required: false
        choices:
        - arg_value: read-write
          enum_value: READ_WRITE
        - arg_value: read-only
          enum_value: READ_ONLY
  required: false
  help_text: |
    Additional data disks for the TPU VM.

    This flag must be repeated to provide multiple data disks. For example:

      $ {command} --data-disk source=projects/my-project/zones/us-central1-c/disks/my-disk,mode=read-only

    The following keys are allowed:

    *source*::: Specifies the full path to an existing disk. Required. The disk must be in the same zone.

    *mode*::: Specifies the mode in which to attach this disk. Valid options are 'read-write',
    'read-only'. If not specified, the default is 'read-write'.

attach_disk:
  arg_name: attach-disk
  type: "googlecloudsdk.calliope.arg_parsers:ArgDict:"
  required: false
  metavar: SOURCE=DATA_DISK
  help_text: |
    Attach a data disk to the TPU VM. For example,

      $ {command} example-tpu --attach-disk source=projects/my-project/zones/us-central1-c/disks/my-disk,mode=read-only --zone=us-central1-c

    attaches the disk named 'projects/my-project/zones/us-central1-c/disks/my-disk' to a TPU VM named
    'example-tpu' in read-only mode in zone `us-central1-c`.

    The following keys are allowed:

    *source*::: Specify the full path to an existing disk. Required. The disk must be in the same zone.

    *mode*::: Specify the mode in which to attach this disk. Valid options are 'read-write',
    'read-only'. If not specified, the default is 'read-write'.

detach_disk:
  arg_name: detach-disk
  type: str
  required: false
  metavar: DATA_DISK
  help_text: |
    Detach a data disk from the TPU VM. For example,

    This flag must be repeated to provide multiple data disks. For example:

      $ {command} example-tpu --detach-disk=projects/my-project/zones/us-central1-c/disks/my-disk --zone=us-central1-c

    detaches the disk named 'projects/my-project/zones/us-central1-c/disks/my-disk' from a TPU VM
    named 'example-tpu' in zone `us-central1-c`.

enable_secure_boot:
  api_field: node.shieldedInstanceConfig.enableSecureBoot
  arg_name: shielded-secure-boot
  type: bool
  default: false
  required: false
  help_text: |
    Specifies that the TPU instances are created with secure boot enabled. This implicitly makes
    them Shielded VM instances.

boot_disk:
  arg_name: boot-disk
  required: false
  metavar: KEY=VALUE
  type: "googlecloudsdk.calliope.arg_parsers:ArgDict:"
  help_text: |
    Specifies the boot disk configuration.

      $ {command} \
            --boot-disk kms-key=<full_kms_key_name_here>

    The following keys are allowed:

    *kms-key*::: Specifies the fully qualified Cloud KMS cryptokey name
    which will be used to protect the disk. KMS cryptokey name format:
      projects/<kms-project>/locations/<kms-location>/keyRings/<kms-keyring>/cryptoKeys/<key-name>
  ALPHA:
    help_text: |
      Specifies the boot disk configuration.

        $ {command} \
              --boot-disk confidential-compute=True,kms-key=<full_kms_key_name_here>

      The following keys are allowed:

      *confidential-compute*::: Create the boot disk in confidential compute mode.
       CMEK layer is required.
      *kms-key*::: Specifies the fully qualified Cloud KMS cryptokey name
      which will be used to protect the disk. KMS cryptokey name format:
        projects/<kms-project>/locations/<kms-location>/keyRings/<kms-keyring>/cryptoKeys/<key-name>

autocheckpoint_enabled:
  api_field: node.autocheckpointEnabled
  arg_name: autocheckpoint-enabled
  type: bool
  default: false
  required: false
  help_text: |
    If specified, the TPU node is created with the Autocheckpoint feature enabled.

reservation:
  api_field: node.schedulingConfig.reservationName
  arg_name: reservation
  type: str
  required: false
  help_text: |
    The name of the reservation. This can either be the full name or just the name.
    However, a full name is required if the reservation is not in the default project.

    Full reservation name format is:

    projects/<project>/reservations/<reservation-name>

provisioning_model:
  arg_name: provisioning-model
  required: false
  help_text: |
    The provisioning model of the TPU node. One of [standard, spot, reservation-bound]. Setting
    `--provisioning-model=spot` is functionally equivalent to setting `--spot`.

worker:
  arg_name: worker
  type: "googlecloudsdk.calliope.arg_parsers:ArgList:"
  required: false
  help_text: |
    List of worker IDs to apply attach/detach disk to.
    If not specified, the update is applied to all workers.
    Only numeric values are allowed.
    The only exception is 'all'.

    `--attach-disk` and `--worker` can be used to specify which workers to attach the disk to. For example:

      $ {command} example-tpu --attach-disk source=projects/my-project/zones/us-central1-c/disks/my-disk,mode=read-only --zone=us-central1-c
            --worker=0,1

    attaches the disk named 'projects/my-project/zones/us-central1-c/disks/my-disk' to worker 0 and 1 of a TPU VM
    named 'example-tpu' in read-only mode in zone `us-central1-c`.
    if `--worker` is not specified or `--worker=all` is specified, the disk is attached to all workers.

    `--detach-disk` and `--worker` can be used to specify which workers to detach the disk from. For example:

      $ {command} example-tpu --detach-disk=projects/my-project/zones/us-central1-c/disks/my-disk --zone=us-central1-c
            --worker=0,1

    detaches the disk named 'projects/my-project/zones/us-central1-c/disks/my-disk' from worker 0 and 1 of a TPU VM
    named 'example-tpu' in zone `us-central1-c`.
    if `--worker` is not specified or `--worker=all` is specified, the disk is attached to all workers.