feat: Add new gcloud commands, API clients, and third-party libraries across various services.

gcloud auth application-default login/google-cloud-sdk/lib/surface/ai/persistent_resources/__init__.py

@@ -0,0 +1,36 @@
+# -*- coding: utf-8 -*- #
+# Copyright 2023 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.
+"""Command group for Vertex AI persistent resources."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import unicode_literals
+
+from googlecloudsdk.calliope import base
+
+
+@base.ReleaseTracks(base.ReleaseTrack.GA)
+class PersistentResourceGA(base.Group):
+  """Create and manage Vertex AI persistent resources."""
+
+  category = base.VERTEX_AI_CATEGORY
+
+
+@base.ReleaseTracks(base.ReleaseTrack.ALPHA, base.ReleaseTrack.BETA)
+class PersistentResourcePreGA(base.Group):
+  """Create and manage Vertex AI rersistent resources."""
+
+  category = base.VERTEX_AI_CATEGORY
+

gcloud auth application-default login/google-cloud-sdk/lib/surface/ai/persistent_resources/create.py

@@ -0,0 +1,160 @@
+# -*- coding: utf-8 -*- #
+# Copyright 2023 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.
+"""Command to create a Persistent Resource in Vertex AI."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import unicode_literals
+
+import re
+
+from googlecloudsdk.api_lib.ai.persistent_resources import client
+from googlecloudsdk.calliope import base
+from googlecloudsdk.command_lib.ai import constants
+from googlecloudsdk.command_lib.ai import endpoint_util
+from googlecloudsdk.command_lib.ai import validation as common_validation
+from googlecloudsdk.command_lib.ai.persistent_resources import flags
+from googlecloudsdk.command_lib.ai.persistent_resources import persistent_resource_util
+from googlecloudsdk.command_lib.ai.persistent_resources import validation
+from googlecloudsdk.command_lib.util.args import labels_util
+from googlecloudsdk.core import log
+from googlecloudsdk.core import properties
+
+_OPERATION_RESOURCE_NAME_TEMPLATE = (
+    'projects/{project_number}/locations/{region}/operations/{operation_id}')
+
+_PERSISTENT_RESOURCE_CREATION_DISPLAY_MESSAGE_TEMPLATE = """\
+Operation to create PersistentResource [{display_name}] is submitted successfully.
+
+You may view the status of your PersistentResource create operation with the command
+
+  $ {command_prefix} ai operations describe {operation_resource_name}
+"""
+
+
+@base.ReleaseTracks(base.ReleaseTrack.GA)
+class CreateGA(base.CreateCommand):
+  """Create a new persistent resource.
+
+  This command will create a persistent resource on the users project to use
+  with Vertex AI custom training jobs. Persistent resources remain active until
+  they are deleted by the user.
+
+  ## EXAMPLES
+
+  To create a PersistentResource under project ``example'' in region
+  ``us-central1'', run:
+
+    $ {command} --region=us-central1 --project=example
+    --resource-pool-spec=replica-count=1,machine-type='n1-standard-4'
+    --display-name=example-resource
+  """
+
+  _version = constants.GA_VERSION
+
+  @staticmethod
+  def Args(parser):
+    flags.AddCreatePersistentResourceFlags(parser)
+
+  def _DisplayResult(self, response, project_number, region):
+    cmd_prefix = 'gcloud'
+    if self.ReleaseTrack().prefix:
+      cmd_prefix += ' ' + self.ReleaseTrack().prefix
+
+    operation_id = re.search(r'operations\/(\d+)', response.name).groups(0)[0]
+    operation_resource_name = _OPERATION_RESOURCE_NAME_TEMPLATE.format(
+        project_number=project_number,
+        region=region,
+        operation_id=operation_id,
+    )
+
+    log.status.Print(
+        _PERSISTENT_RESOURCE_CREATION_DISPLAY_MESSAGE_TEMPLATE.format(
+            display_name=response.name,
+            command_prefix=cmd_prefix,
+            operation_resource_name=operation_resource_name
+        )
+    )
+
+  def _PrepareResourcePools(self, args, api_client):
+    persistent_resource_config = (
+        api_client.ImportResourceMessage(args.config, 'PersistentResource')
+        if args.config
+        else api_client.PersistentResourceMessage()
+    )
+
+    validation.ValidateCreateArgs(
+        args, persistent_resource_config, self._version
+    )
+    resource_pool_specs = list(args.resource_pool_spec or [])
+    persistent_resource_spec = persistent_resource_util.ConstructResourcePools(
+        api_client,
+        persistent_resource_config=persistent_resource_config,
+        resource_pool_specs=resource_pool_specs,
+    )
+    return persistent_resource_spec
+
+  def Run(self, args):
+    project = properties.VALUES.core.project.GetOrFail()
+    region_ref = args.CONCEPTS.region.Parse()
+    region = region_ref.AsDict()['locationsId']
+    validation.ValidateRegion(region)
+
+    with endpoint_util.AiplatformEndpointOverrides(
+        version=self._version, region=region
+    ):
+      api_client = client.PersistentResourcesClient(version=self._version)
+      resource_pools = self._PrepareResourcePools(
+          args, api_client
+      )
+      labels = labels_util.ParseCreateArgs(
+          args, api_client.PersistentResourceMessage().LabelsValue
+      )
+
+      response = api_client.Create(
+          parent=region_ref.RelativeName(),
+          display_name=args.display_name,
+          resource_pools=resource_pools,
+          persistent_resource_id=args.persistent_resource_id,
+          kms_key_name=common_validation.GetAndValidateKmsKey(args),
+          labels=labels,
+          network=args.network,
+          enable_custom_service_account=args.enable_custom_service_account,
+          # TODO(b/262780738): Unimplemented
+          # service_account=args.service_account,
+      )
+      self._DisplayResult(response, project, region)
+      return response
+
+
+@base.ReleaseTracks(base.ReleaseTrack.ALPHA, base.ReleaseTrack.BETA)
+class CreatePreGA(CreateGA):
+  """Create a new persistent resource.
+
+  This command will create a persistent resource on the users project to use
+  with Vertex AI custom training jobs. Persistent resources remain active until
+  they are deleted by the user.
+
+  ## EXAMPLES
+
+  To create a PersistentResource under project ``example'' in region
+  ``us-central1'', run:
+
+    $ {command} --region=us-central1 --project=example
+    --resource-pool-spec=replica-count=1,machine-type='n1-standard-4'
+    --display-name=example-resource
+  """
+
+  _version = constants.BETA_VERSION

gcloud auth application-default login/google-cloud-sdk/lib/surface/ai/persistent_resources/delete.py

@@ -0,0 +1,94 @@
+# -*- coding: utf-8 -*- #
+# Copyright 2023 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.
+"""Command to delete a Persistent Resource in Vertex AI."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import unicode_literals
+
+from googlecloudsdk.api_lib.ai.persistent_resources import client
+from googlecloudsdk.calliope import base
+from googlecloudsdk.command_lib.ai import constants
+from googlecloudsdk.command_lib.ai import endpoint_util
+from googlecloudsdk.command_lib.ai.persistent_resources import flags
+from googlecloudsdk.command_lib.ai.persistent_resources import validation
+from googlecloudsdk.core import log
+
+_PERSISTENT_RESOURCE_DELETE_DISPLAY_MESSAGE = """\
+Request to delete the PersistentResource [{name}] has been sent.
+
+You may view the status of your persistent resource with the command
+
+  $ {command_prefix} ai persistent-resources describe {name}
+"""
+
+
+@base.ReleaseTracks(base.ReleaseTrack.GA)
+class DeleteGA(base.SilentCommand):
+  """Delete an active Persistent Resource.
+
+  If the Persistent Resource is not in the active state,
+  the command will not perform any operation.
+
+  ## EXAMPLES
+
+  To delete a persistent resource ``123'' under project ``example'' in region
+  ``us-central1'', run:
+
+    $ {command} 123 --project=example --region=us-central1
+  """
+  _api_version = constants.GA_VERSION
+
+  @staticmethod
+  def Args(parser):
+    flags.AddPersistentResourceResourceArg(parser, 'to delete')
+
+  def _CommandPrefix(self):
+    cmd_prefix = 'gcloud'
+    if self.ReleaseTrack().prefix:
+      cmd_prefix += ' ' + self.ReleaseTrack().prefix
+    return cmd_prefix
+
+  def Run(self, args):
+    persistent_resource_ref = args.CONCEPTS.persistent_resource.Parse()
+    region = persistent_resource_ref.AsDict()['locationsId']
+    validation.ValidateRegion(region)
+
+    with endpoint_util.AiplatformEndpointOverrides(
+        version=self._api_version, region=region):
+      resource_name = persistent_resource_ref.RelativeName()
+      response = client.PersistentResourcesClient(
+          version=self._api_version).Delete(resource_name)
+      log.status.Print(
+          _PERSISTENT_RESOURCE_DELETE_DISPLAY_MESSAGE.format(
+              name=resource_name, command_prefix=self._CommandPrefix()))
+      return response
+
+
+@base.ReleaseTracks(base.ReleaseTrack.ALPHA, base.ReleaseTrack.BETA)
+class DeletePreGA(DeleteGA):
+  """Delete an active Persistent Resource.
+
+  If the Persistent Resource is not in the active state,
+  the command will not perform any operation.
+
+  ## EXAMPLES
+
+  To delete a persistent resource ``123'' under project ``example'' in region
+  ``us-central1'', run:
+
+    $ {command} 123 --project=example --region=us-central1
+  """
+  _api_version = constants.BETA_VERSION

gcloud auth application-default login/google-cloud-sdk/lib/surface/ai/persistent_resources/describe.py

@@ -0,0 +1,73 @@
+# -*- coding: utf-8 -*- #
+# Copyright 2023 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.
+"""Command to get a Persistent Resource in Vertex AI."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import unicode_literals
+
+from googlecloudsdk.api_lib.ai.persistent_resources import client
+from googlecloudsdk.calliope import base
+from googlecloudsdk.command_lib.ai import constants
+from googlecloudsdk.command_lib.ai import endpoint_util
+from googlecloudsdk.command_lib.ai.persistent_resources import flags
+from googlecloudsdk.command_lib.ai.persistent_resources import validation
+
+
+@base.ReleaseTracks(base.ReleaseTrack.GA)
+class DescribeGA(base.DescribeCommand):
+  """Get detailed information about a PersistentResource with a given id.
+
+  ## EXAMPLES
+
+  To get the persistent resource with the PersistentResource id ``123'' under
+  project ``example'' in region ``us-central1'', run:
+
+    $ {command} 123 --project=example --region=us-central1
+  """
+
+  _api_version = constants.GA_VERSION
+
+  @staticmethod
+  def Args(parser):
+    flags.AddPersistentResourceResourceArg(parser, 'to describe')
+
+  def Run(self, args):
+    persistent_resource_ref = args.CONCEPTS.persistent_resource.Parse()
+    region = persistent_resource_ref.AsDict()['locationsId']
+    validation.ValidateRegion(region)
+
+    with endpoint_util.AiplatformEndpointOverrides(
+        version=self._api_version, region=region
+    ):
+      response = client.PersistentResourcesClient(
+          version=self._api_version
+      ).Get(persistent_resource_ref.RelativeName())
+      return response
+
+
+@base.ReleaseTracks(base.ReleaseTrack.ALPHA, base.ReleaseTrack.BETA)
+class DescribePreGA(DescribeGA):
+  """Get detailed information about a PersistentResource with a given id.
+
+  ## EXAMPLES
+
+  To get the persistent resource with the PersistentResource id ``123'' under
+  project ``example'' in region ``us-central1'', run:
+
+    $ {command} 123 --project=example --region=us-central1
+  """
+
+  _api_version = constants.BETA_VERSION

gcloud auth application-default login/google-cloud-sdk/lib/surface/ai/persistent_resources/list.py

@@ -0,0 +1,100 @@
+# -*- coding: utf-8 -*- #
+# Copyright 2023 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.
+"""Command to list persistent resources in Vertex AI."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import unicode_literals
+
+from googlecloudsdk.api_lib.ai.persistent_resources import client
+from googlecloudsdk.calliope import base
+from googlecloudsdk.command_lib.ai import constants
+from googlecloudsdk.command_lib.ai import endpoint_util
+from googlecloudsdk.command_lib.ai import flags
+from googlecloudsdk.command_lib.ai import region_util
+from googlecloudsdk.command_lib.ai.persistent_resources import persistent_resource_util
+from googlecloudsdk.command_lib.ai.persistent_resources import validation
+
+
+@base.ReleaseTracks(base.ReleaseTrack.GA)
+class ListGA(base.ListCommand):
+  """Lists the active persistent resources.
+
+  ## EXAMPLES
+
+  To list the persistent resources of project ``example'' in region
+  ``us-central1'', run:
+
+    $ {command} --project=example --region=us-central1
+  """
+  _api_version = constants.GA_VERSION
+
+  @classmethod
+  def Args(cls, parser):
+    """Method called by Calliope to set up arguments for this command.
+
+    Args:
+      parser: A argparse.Parser to register accepted arguments in command input.
+    """
+    flags.AddRegionResourceArg(
+        parser,
+        'to list persistent resources',
+        prompt_func=region_util.GetPromptForRegionFunc(
+            constants.SUPPORTED_TRAINING_REGIONS))
+    flags.AddUriFlags(
+        parser,
+        collection=persistent_resource_util.PERSISTENT_RESOURCE_COLLECTION,
+        api_version=constants.AI_PLATFORM_API_VERSION[cls._api_version])
+
+  def Run(self, args):
+    """Executes the list command.
+
+    Args:
+      args: an argparse.Namespace, it contains all arguments that this command
+        was invoked with.
+
+    Returns:
+      The list of resources
+    """
+    region_ref = args.CONCEPTS.region.Parse()
+    region = region_ref.AsDict()['locationsId']
+    validation.ValidateRegion(region)
+
+    with endpoint_util.AiplatformEndpointOverrides(
+        version=self._api_version, region=region):
+      response = client.PersistentResourcesClient(
+          version=self._api_version).List(
+              region=region_ref.RelativeName())
+      # Filter out resources with a raySpec present to hide ray clusters from
+      # the user. In both the error and success case, the response will be a
+      # generator object `YieldFromList`
+      response = [resource for resource in response if (
+          (resource.resourceRuntimeSpec is None) or
+          (resource.resourceRuntimeSpec.raySpec is None))]
+      return response
+
+
+@base.ReleaseTracks(base.ReleaseTrack.ALPHA, base.ReleaseTrack.BETA)
+class ListPreGA(ListGA):
+  """Lists the active persistent resources.
+
+  ## EXAMPLES
+
+  To list the persistent resources of project ``example'' in region
+  ``us-central1'', run:
+
+    $ {command} --project=example --region=us-central1
+  """
+  _api_version = constants.BETA_VERSION

gcloud auth application-default login/google-cloud-sdk/lib/surface/ai/persistent_resources/reboot.py

@@ -0,0 +1,106 @@
+# -*- coding: utf-8 -*- #
+# Copyright 2024 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.
+"""Command to reboot a Persistent Resource in Vertex AI."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import unicode_literals
+
+import re
+
+from googlecloudsdk.api_lib.ai.persistent_resources import client
+from googlecloudsdk.calliope import base
+from googlecloudsdk.command_lib.ai import constants
+from googlecloudsdk.command_lib.ai import endpoint_util
+from googlecloudsdk.command_lib.ai.persistent_resources import flags
+from googlecloudsdk.command_lib.ai.persistent_resources import validation
+from googlecloudsdk.core import log
+from googlecloudsdk.core import properties
+
+_OPERATION_RESOURCE_NAME_TEMPLATE = (
+    'projects/{project}/locations/{region}/operations/{operation_id}')
+
+_PERSISTENT_RESOURCE_REBOOT_DISPLAY_MESSAGE = """\
+Request to reboot the PersistentResource [{name}] has been sent.
+
+You may view the status of your PersistentResource reboot operation with the command
+
+  $ {command_prefix} ai operations describe {operation_resource_name}
+
+"""
+
+
+@base.ReleaseTracks(base.ReleaseTrack.GA)
+class RebootGA(base.SilentCommand):
+  """Reboot a Persistent Resource.
+
+  ## EXAMPLES
+
+  To reboot a persistent resource ``123'' under project ``example'' in region
+  ``us-central1'', run:
+
+    $ {command} 123 --project=example --region=us-central1
+  """
+  _api_version = constants.GA_VERSION
+
+  @staticmethod
+  def Args(parser):
+    flags.AddPersistentResourceResourceArg(parser, 'to reboot')
+
+  def _CommandPrefix(self):
+    cmd_prefix = 'gcloud'
+    if self.ReleaseTrack().prefix:
+      cmd_prefix += ' ' + self.ReleaseTrack().prefix
+    return cmd_prefix
+
+  def Run(self, args):
+    persistent_resource_ref = args.CONCEPTS.persistent_resource.Parse()
+    region = persistent_resource_ref.AsDict()['locationsId']
+    validation.ValidateRegion(region)
+
+    with endpoint_util.AiplatformEndpointOverrides(
+        version=self._api_version, region=region):
+      project = properties.VALUES.core.project.GetOrFail()
+      resource_name = persistent_resource_ref.RelativeName()
+      response = client.PersistentResourcesClient(
+          version=self._api_version).Reboot(resource_name)
+
+      operation_id = re.search(r'operations\/(\d+)', response.name).groups(0)[0]
+      operation_resource_name = _OPERATION_RESOURCE_NAME_TEMPLATE.format(
+          project=project,
+          region=region,
+          operation_id=operation_id,
+      )
+
+      log.status.Print(
+          _PERSISTENT_RESOURCE_REBOOT_DISPLAY_MESSAGE.format(
+              name=resource_name,
+              command_prefix=self._CommandPrefix(),
+              operation_resource_name=operation_resource_name))
+      return response
+
+
+@base.ReleaseTracks(base.ReleaseTrack.ALPHA, base.ReleaseTrack.BETA)
+class RebootPreGA(RebootGA):
+  """Reboot an active Persistent Resource.
+
+  ## EXAMPLES
+
+  To reboot a persistent resource ``123'' under project ``example'' in region
+  ``us-central1'', run:
+
+    $ {command} 123 --project=example --region=us-central1
+  """
+  _api_version = constants.BETA_VERSION