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

This commit is contained in:
2026-01-01 20:26:35 +01:00
parent 5e23cbece0
commit a19e592eb7
25221 changed files with 8324611 additions and 0 deletions

View File

@@ -0,0 +1,29 @@
# -*- coding: utf-8 -*- #
# Copyright 2020 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 custom job."""
from __future__ import absolute_import
from __future__ import division
from __future__ import unicode_literals
from googlecloudsdk.calliope import base
@base.ReleaseTracks(base.ReleaseTrack.ALPHA, base.ReleaseTrack.BETA,
base.ReleaseTrack.GA)
class CustomJob(base.Group):
"""Manage Vertex AI custom jobs."""
category = base.VERTEX_AI_CATEGORY

View File

@@ -0,0 +1,93 @@
# -*- coding: utf-8 -*- #
# Copyright 2020 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 cancel a custom job in Vertex AI."""
from __future__ import absolute_import
from __future__ import division
from __future__ import unicode_literals
from googlecloudsdk.api_lib.ai.custom_jobs 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.custom_jobs import flags
from googlecloudsdk.command_lib.ai.custom_jobs import validation
from googlecloudsdk.core import log
_CUSTOM_JOB_CANCEL_DISPLAY_MESSAGE = """\
Request to cancel CustomJob [{job_name}] has been sent.
You may view the status of your job with the command
$ {command_prefix} ai custom-jobs describe {job_name}
"""
@base.ReleaseTracks(base.ReleaseTrack.GA)
class CancelGA(base.SilentCommand):
"""Cancel a running custom job.
If the job is already finished,
the command will not perform any operation.
## EXAMPLES
To cancel a job ``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.AddCustomJobResourceArg(parser, 'to cancel')
def _CommandPrefix(self):
return 'gcloud'
def Run(self, args):
custom_job_ref = args.CONCEPTS.custom_job.Parse()
region = custom_job_ref.AsDict()['locationsId']
validation.ValidateRegion(region)
with endpoint_util.AiplatformEndpointOverrides(
version=self._api_version, region=region):
job_name = custom_job_ref.RelativeName()
response = client.CustomJobsClient(
version=self._api_version).Cancel(job_name)
log.status.Print(
_CUSTOM_JOB_CANCEL_DISPLAY_MESSAGE.format(
job_name=job_name, command_prefix=self._CommandPrefix()))
return response
@base.ReleaseTracks(base.ReleaseTrack.BETA, base.ReleaseTrack.ALPHA)
class CancelPreGA(CancelGA):
"""Cancel a running custom job.
If the job is already finished,
the command will not perform any operation.
To cancel a job ``123'' under project ``example'' in region
``us-central1'', run:
$ {command} 123 --project=example --region=us-central1
"""
_api_version = constants.BETA_VERSION
def _CommandPrefix(self):
return 'gcloud ' + self.ReleaseTrack().prefix

View File

@@ -0,0 +1,141 @@
# -*- coding: utf-8 -*- #
# Copyright 2020 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 custom job in Vertex AI."""
from __future__ import absolute_import
from __future__ import division
from __future__ import unicode_literals
from googlecloudsdk.api_lib.ai.custom_jobs 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.custom_jobs import custom_jobs_util
from googlecloudsdk.command_lib.ai.custom_jobs import flags
from googlecloudsdk.command_lib.ai.custom_jobs import validation
from googlecloudsdk.command_lib.util.args import labels_util
from googlecloudsdk.core import log
from googlecloudsdk.core import properties
_JOB_CREATION_DISPLAY_MESSAGE_TEMPLATE = """\
CustomJob [{job_name}] is submitted successfully.
Your job is still active. You may view the status of your job with the command
$ {command_prefix} ai custom-jobs describe {job_name}
or continue streaming the logs with the command
$ {command_prefix} ai custom-jobs stream-logs {job_name}\
"""
@base.ReleaseTracks(base.ReleaseTrack.GA)
class CreateGA(base.CreateCommand):
"""Create a new custom job.
This command will attempt to run the custom job immediately upon creation.
## EXAMPLES
To create a job under project ``example'' in region
``us-central1'', run:
$ {command} --region=us-central1 --project=example
--worker-pool-spec=replica-count=1,machine-type='n1-highmem-2',container-image-uri='gcr.io/ucaip-test/ucaip-training-test'
--display-name=test
"""
_version = constants.GA_VERSION
@staticmethod
def Args(parser):
flags.AddCreateCustomJobFlags(parser)
def _DisplayResult(self, response):
cmd_prefix = 'gcloud'
if self.ReleaseTrack().prefix:
cmd_prefix += ' ' + self.ReleaseTrack().prefix
log.status.Print(
_JOB_CREATION_DISPLAY_MESSAGE_TEMPLATE.format(
job_name=response.name, command_prefix=cmd_prefix))
def _PrepareJobSpec(self, args, api_client, project):
job_config = api_client.ImportResourceMessage(
args.config, 'CustomJobSpec') if args.config else api_client.GetMessage(
'CustomJobSpec')()
validation.ValidateCreateArgs(args, job_config, self._version)
worker_pool_specs = list(
custom_jobs_util.UpdateWorkerPoolSpecsIfLocalPackageRequired(
args.worker_pool_spec or [], args.display_name, project))
job_spec = custom_jobs_util.ConstructCustomJobSpec(
api_client,
base_config=job_config,
worker_pool_specs=worker_pool_specs,
network=args.network,
service_account=args.service_account,
enable_web_access=args.enable_web_access,
enable_dashboard_access=args.enable_dashboard_access,
args=args.args,
command=args.command,
python_package_uri=args.python_package_uris,
persistent_resource_id=args.persistent_resource_id)
return job_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.CustomJobsClient(version=self._version)
job_spec = self._PrepareJobSpec(args, api_client, project)
labels = labels_util.ParseCreateArgs(
args,
api_client.CustomJobMessage().LabelsValue)
response = api_client.Create(
parent=region_ref.RelativeName(),
display_name=args.display_name,
job_spec=job_spec,
kms_key_name=common_validation.GetAndValidateKmsKey(args),
labels=labels)
self._DisplayResult(response)
return response
@base.ReleaseTracks(base.ReleaseTrack.BETA, base.ReleaseTrack.ALPHA)
class CreatePreGA(CreateGA):
"""Create a new custom job.
This command will attempt to run the custom job immediately upon creation.
## EXAMPLES
To create a job under project ``example'' in region
``us-central1'', run:
$ {command} --region=us-central1 --project=example
--worker-pool-spec=replica-count=1,machine-type='n1-highmem-2',container-image-uri='gcr.io/ucaip-test/ucaip-training-test'
--display-name=test
"""
_version = constants.BETA_VERSION

View File

@@ -0,0 +1,71 @@
# -*- coding: utf-8 -*- #
# Copyright 2020 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 custom job in Vertex AI."""
from __future__ import absolute_import
from __future__ import division
from __future__ import unicode_literals
from googlecloudsdk.api_lib.ai.custom_jobs 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.custom_jobs import flags
from googlecloudsdk.command_lib.ai.custom_jobs import validation
@base.ReleaseTracks(base.ReleaseTrack.GA)
class DescribeGA(base.DescribeCommand):
"""Get detailed information about the custom job by given id.
## EXAMPLES
To get a job ``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.AddCustomJobResourceArg(parser, 'to describe')
def Run(self, args):
custom_job_ref = args.CONCEPTS.custom_job.Parse()
region = custom_job_ref.AsDict()['locationsId']
validation.ValidateRegion(region)
with endpoint_util.AiplatformEndpointOverrides(
version=self._api_version, region=region):
response = client.CustomJobsClient(version=self._api_version).Get(
custom_job_ref.RelativeName())
return response
@base.ReleaseTracks(base.ReleaseTrack.BETA, base.ReleaseTrack.ALPHA)
class DescribePreGA(DescribeGA):
"""Get detailed information about the custom job by given id.
## EXAMPLES
To get a job ``123'' under project ``example'' in region
``us-central1'', run:
$ {command} 123 --project=example --region=us-central1
"""
_api_version = constants.BETA_VERSION

View File

@@ -0,0 +1,93 @@
# -*- coding: utf-8 -*- #
# Copyright 2020 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 custom jobs in Vertex AI."""
from __future__ import absolute_import
from __future__ import division
from __future__ import unicode_literals
from googlecloudsdk.api_lib.ai.custom_jobs 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.custom_jobs import custom_jobs_util
from googlecloudsdk.command_lib.ai.custom_jobs import validation
@base.ReleaseTracks(base.ReleaseTrack.GA)
class ListGA(base.ListCommand):
"""Lists the existing custom jobs.
## EXAMPLES
To list the jobs 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 custom jobs',
prompt_func=region_util.GetPromptForRegionFunc(
constants.SUPPORTED_TRAINING_REGIONS))
flags.AddUriFlags(
parser,
collection=custom_jobs_util.CUSTOM_JOB_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):
return client.CustomJobsClient(version=self._api_version).List(
region=region_ref.RelativeName())
@base.ReleaseTracks(base.ReleaseTrack.BETA, base.ReleaseTrack.ALPHA)
class ListPreGA(ListGA):
"""Lists the existing custom jobs.
## EXAMPLES
To list the jobs of project ``example'' in region
``us-central1'', run:
$ {command} --project=example --region=us-central1
"""
_api_version = constants.BETA_VERSION

View File

@@ -0,0 +1,149 @@
# -*- coding: utf-8 -*- #
# Copyright 2020 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 run a training application locally."""
from __future__ import absolute_import
from __future__ import division
from __future__ import unicode_literals
import os
import textwrap
from googlecloudsdk.calliope import base
from googlecloudsdk.command_lib.ai.custom_jobs import flags
from googlecloudsdk.command_lib.ai.custom_jobs import local_util
from googlecloudsdk.command_lib.ai.custom_jobs import validation
from googlecloudsdk.command_lib.ai.docker import build as docker_builder
from googlecloudsdk.command_lib.ai.docker import run as docker_runner
from googlecloudsdk.core import log
from googlecloudsdk.core.util import files
@base.ReleaseTracks(base.ReleaseTrack.GA, base.ReleaseTrack.BETA,
base.ReleaseTrack.ALPHA)
class Create(base.CreateCommand):
"""Run a custom training locally.
Packages your training code into a Docker image and executes it locally.
"""
detailed_help = {
'DESCRIPTION':
textwrap.dedent("""\
{description}
You should execute this command in the top folder which includes all
the code and resources you want to pack and run, or specify the
'work-dir' flag to point to it. Any other path you specified via flags
should be a relative path to the work-dir and under it; otherwise it
will be unaccessible.
Supposing your directories are like the following structures:
/root
- my_project
- my_training
- task.py
- util.py
- setup.py
- other_modules
- some_module.py
- dataset
- small.dat
- large.dat
- config
- dep
- foo.tar.gz
- bar.whl
- requirements.txt
- another_project
- something
If you set 'my_project' as the package, then you should
execute the task.py by specifying "--script=my_training/task.py" or
"--python-module=my_training.task", the 'requirements.txt' will be
processed. And you will also be able to install extra packages by,
e.g. specifying "--extra-packages=dep/foo.tar.gz,bar.whl" or include
extra directories, e.g. specifying "--extra-dirs=dataset,config".
If you set 'my_training' as the package, then you should
execute the task.py by specifying "--script=task.py" or
"--python-module=task", the 'setup.py' will be processed. However, you
won't be able to access any other files or directories that are not in
'my_training' folder.
See more details in the HELP info of the corresponding flags.
"""),
'EXAMPLES':
"""\
To execute an python module with required dependencies, run:
$ {command} --python-module=my_training.task --executor-image-uri=gcr.io/my/image --requirements=pandas,scipy>=1.3.0
To execute a python script using local GPU, run:
$ {command} --script=my_training/task.py --executor-image-uri=gcr.io/my/image --gpu
To execute an arbitrary script with custom arguments, run:
$ {command} --script=my_run.sh --executor-image-uri=gcr.io/my/image -- --my-arg bar --enable_foo
To run an existing container training without building new image, run:
$ {command} --executor-image-uri=gcr.io/my/custom-training-image
""",
}
@staticmethod
def Args(parser):
flags.AddLocalRunCustomJobFlags(parser)
def Run(self, args):
args = validation.ValidateLocalRunArgs(args)
with files.ChDir(args.local_package_path):
log.status.Print('Package is set to {}.'.format(args.local_package_path))
executable_image = args.executor_image_uri or args.base_image
if args.script:
# TODO(b/176214485): Consider including the image id in build result.
built_image = docker_builder.BuildImage(
base_image=executable_image,
host_workdir=args.local_package_path,
main_script=args.script,
python_module=args.python_module,
requirements=args.requirements,
extra_packages=args.extra_packages,
extra_dirs=args.extra_dirs,
output_image_name=args.output_image_uri)
executable_image = built_image.name
log.status.Print('A training image is built.')
log.status.Print('Starting to run ...')
docker_runner.RunContainer(
image_name=executable_image,
enable_gpu=args.gpu,
service_account_key=args.service_account_key_file,
user_args=args.args)
log.out.Print(
'A local run is finished successfully using custom image: {}.'.format(
executable_image))
# Clean generated cache
cache_dir, _ = os.path.split(
os.path.join(args.local_package_path, args.script or ''))
if local_util.ClearPyCache(cache_dir):
log.status.Print(
'Cleaned Python cache from directory: {}'.format(cache_dir))

View File

@@ -0,0 +1,79 @@
# -*- coding: utf-8 -*- #
# Copyright 2020 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 check stream logs of a custom job in Vertex AI."""
from __future__ import absolute_import
from __future__ import division
from __future__ import unicode_literals
from googlecloudsdk.api_lib.ai.custom_jobs 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 as common_flags
from googlecloudsdk.command_lib.ai import log_util
from googlecloudsdk.command_lib.ai.custom_jobs import flags as custom_job_flags
from googlecloudsdk.command_lib.ai.custom_jobs import validation
@base.ReleaseTracks(base.ReleaseTrack.GA)
class StreamLogsGA(base.Command):
"""Show stream logs from a running custom job.
## EXAMPLES
To stream logs of custom job ``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):
custom_job_flags.AddCustomJobResourceArg(parser, 'to fetch stream log')
common_flags.AddStreamLogsFlags(parser)
parser.display_info.AddFormat(log_util.LOG_FORMAT)
def Run(self, args):
custom_job_ref = args.CONCEPTS.custom_job.Parse()
region = custom_job_ref.AsDict()['locationsId']
validation.ValidateRegion(region)
with endpoint_util.AiplatformEndpointOverrides(
version=self._api_version, region=region):
return log_util.StreamLogs(
custom_job_ref.AsDict()['customJobsId'],
continue_function=client.CustomJobsClient(
version=self._api_version).CheckJobComplete(
custom_job_ref.RelativeName()),
polling_interval=args.polling_interval,
task_name=args.task_name,
allow_multiline=args.allow_multiline_logs)
@base.ReleaseTracks(base.ReleaseTrack.BETA, base.ReleaseTrack.ALPHA)
class StreamLogsPreGA(StreamLogsGA):
"""Show stream logs from a running custom job.
## EXAMPLES
To stream logs of custom job ``123'' under project ``example'' in region
``us-central1'', run:
$ {command} 123 --project=example --region=us-central1
"""
_api_version = constants.BETA_VERSION