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 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.
"""Colab Command Group."""
from googlecloudsdk.calliope import base
@base.DefaultUniverseOnly
@base.ReleaseTracks(base.ReleaseTrack.GA, base.ReleaseTrack.BETA)
class Colab(base.Group):
"""Manage Colab Enterprise resources."""
category = base.AI_AND_MACHINE_LEARNING_CATEGORY
def Filter(self, context, args):
base.RequireProjectID(args)
del context, args

View File

@@ -0,0 +1,28 @@
# -*- 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.
"""The gcloud colab executions command group."""
from googlecloudsdk.calliope import base
@base.DefaultUniverseOnly
@base.ReleaseTracks(base.ReleaseTrack.GA, base.ReleaseTrack.BETA)
class Executions(base.Group):
# pylint: disable=g-doc - this is required for help text.
"""Manage Colab Enterprise notebook execution jobs.
For more information about executions, see
https://cloud.google.com/colab/docs/schedule-notebook-run.
"""

View File

@@ -0,0 +1,79 @@
# -*- 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.
"""Create command for Colab Enterprise Executions."""
from googlecloudsdk.api_lib.ai import operations
from googlecloudsdk.api_lib.colab_enterprise import util
from googlecloudsdk.api_lib.notebook_executor import executions as executions_util
from googlecloudsdk.calliope import base
from googlecloudsdk.command_lib.ai import endpoint_util
from googlecloudsdk.command_lib.notebook_executor import flags
_DETAILED_HELP = {
'DESCRIPTION': """
Create a notebook execution to be used on a Colab Enterprise runtime.
""",
'EXAMPLES': """
To create an execution of a notebook file with Cloud Storage URI `gs://my-bucket/my-notebook.ipynb`, with an execution job display name `my-execution`, compute configured from runtime template `my-runtime-template-id`, running with service account `my-service-account@my-project.iam.gserviceaccount.com`, with results uploaded to Cloud Storage bucket `gs://my-bucket/results`, and in region `us-central1` run:
$ {command} --display-name=my-execution --runtime-template=my-runtime-template-id --gcs-notebook-uri=gs://my-bucket/my-notebook.ipynb --service-account=my-service-account@my-project.iam.gserviceaccount.com --gcs-output-uri=gs://my-bucket/results --region=us-central1
""",
}
@base.DefaultUniverseOnly
@base.ReleaseTracks(base.ReleaseTrack.GA, base.ReleaseTrack.BETA)
class Create(base.CreateCommand):
"""Create an execution."""
@staticmethod
def Args(parser):
"""Register flags for this command."""
flags.AddCreateExecutionFlags(parser)
def Run(self, args):
"""This is what gets called when the user runs this command."""
release_track = self.ReleaseTrack()
messages = util.GetMessages(self.ReleaseTrack())
region_ref = args.CONCEPTS.region.Parse()
region = region_ref.AsDict()['locationsId']
with endpoint_util.AiplatformEndpointOverrides(
version='BETA', region=region
):
api_client = util.GetClient(release_track)
executions_service = (
api_client.projects_locations_notebookExecutionJobs
)
operation = executions_service.Create(
executions_util.CreateExecutionCreateRequest(
args, messages
)
)
return util.WaitForOpMaybe(
operations_client=operations.OperationsClient(client=api_client),
op=operation,
op_ref=executions_util.ParseExecutionOperation(
operation.name
),
asynchronous=util.GetAsyncConfig(args),
kind='notebook execution',
log_method='create',
message='Waiting for execution to be created...',
)
Create.detailed_help = _DETAILED_HELP

View File

@@ -0,0 +1,80 @@
# -*- 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.
"""Delete command for Colab Enterprise Executions."""
from googlecloudsdk.api_lib.ai import operations
from googlecloudsdk.api_lib.colab_enterprise import util
from googlecloudsdk.api_lib.notebook_executor import executions as executions_util
from googlecloudsdk.calliope import base
from googlecloudsdk.command_lib.ai import endpoint_util
from googlecloudsdk.command_lib.notebook_executor import flags
_DETAILED_HELP = {
'DESCRIPTION': """
Delete a Colab Enterprise notebook execution.
""",
'EXAMPLES': """
To delete an execution with id `my-execution`, in region `us-central1`, run:
$ {command} my-execution --region=us-central1
""",
}
@base.DefaultUniverseOnly
@base.ReleaseTracks(base.ReleaseTrack.GA, base.ReleaseTrack.BETA)
class Delete(base.DeleteCommand):
"""Delete an execution."""
@staticmethod
def Args(parser):
"""Register flags for this command."""
flags.AddDeleteExecutionFlags(parser)
def Run(self, args):
"""This is what gets called when the user runs this command."""
release_track = self.ReleaseTrack()
messages = util.GetMessages(self.ReleaseTrack())
execution_ref = args.CONCEPTS.execution.Parse()
region = execution_ref.AsDict()['locationsId']
with endpoint_util.AiplatformEndpointOverrides(
version='BETA', region=region
):
api_client = util.GetClient(release_track)
executions_service = (
api_client.projects_locations_notebookExecutionJobs
)
executions_util.ValidateAndGetColabExecution(args, messages, executions_service)
operation = executions_service.Delete(
executions_util.CreateExecutionDeleteRequest(
args, messages
)
)
return util.WaitForOpMaybe(
operations_client=operations.OperationsClient(client=api_client),
op=operation,
op_ref=executions_util.ParseExecutionOperation(
operation.name
),
asynchronous=util.GetAsyncConfig(args),
kind='notebook execution',
log_method='delete',
message='Waiting for execution to be deleted...',
)
Delete.detailed_help = _DETAILED_HELP

View File

@@ -0,0 +1,66 @@
# -*- 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.
"""Describe command for Colab Enterprise Executions."""
from googlecloudsdk.api_lib.colab_enterprise import util
from googlecloudsdk.api_lib.notebook_executor import executions as executions_util
from googlecloudsdk.calliope import base
from googlecloudsdk.command_lib.ai import endpoint_util
from googlecloudsdk.command_lib.notebook_executor import flags
_DETAILED_HELP = {
'DESCRIPTION': """
Describe a Colab Enterprise notebook execution.
""",
'EXAMPLES': """
To describe a notebook execution with id `my-execution` in region `us-central1`, run:
$ {command} my-execution --region=us-central1
""",
}
@base.DefaultUniverseOnly
@base.ReleaseTracks(base.ReleaseTrack.GA, base.ReleaseTrack.BETA)
class Describe(base.DescribeCommand):
"""Describe an execution."""
@staticmethod
def Args(parser):
"""Register flags for this command."""
flags.AddDescribeExecutionFlags(parser)
def Run(self, args):
"""This is what gets called when the user runs this command."""
release_track = self.ReleaseTrack()
messages = util.GetMessages(self.ReleaseTrack())
execution_ref = args.CONCEPTS.execution.Parse()
region = execution_ref.AsDict()['locationsId']
# Override to regionalize domain as used by the AIPlatform API.
with endpoint_util.AiplatformEndpointOverrides(
version='BETA', region=region
):
api_client = util.GetClient(release_track)
executions_service = (
api_client.projects_locations_notebookExecutionJobs
)
return executions_util.ValidateAndGetColabExecution(
args, messages, executions_service
)
Describe.detailed_help = _DETAILED_HELP

View File

@@ -0,0 +1,80 @@
# -*- 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.
"""List command for Colab Enterprise Executions."""
from apitools.base.py import list_pager
from googlecloudsdk.api_lib.colab_enterprise import util
from googlecloudsdk.api_lib.notebook_executor import executions as executions_util
from googlecloudsdk.calliope import base
from googlecloudsdk.command_lib.ai import endpoint_util
from googlecloudsdk.command_lib.notebook_executor import flags
_DETAILED_HELP = {
'DESCRIPTION': """
List your project's Colab Enterprise notebook executions in a given region.
""",
'EXAMPLES': """
To list your executions in region `us-central1`, run:
$ {command} --region=us-central1
""",
}
@base.DefaultUniverseOnly
@base.ReleaseTracks(base.ReleaseTrack.GA, base.ReleaseTrack.BETA)
class List(base.ListCommand):
"""List your Colab Enterprise notebook execution jobs."""
@staticmethod
def Args(parser):
"""Register flags for this command."""
flags.AddListExecutionsFlags(parser)
parser.display_info.AddFormat("""
table(name.segment(-1):label=ID,
displayName,
name.segment(-3):label=REGION,
gcsOutputUri,
jobState)
""")
def Run(self, args):
"""This is what gets called when the user runs this command."""
release_track = self.ReleaseTrack()
messages = util.GetMessages(self.ReleaseTrack())
region_ref = args.CONCEPTS.region.Parse()
region = region_ref.AsDict()['locationsId']
# Override to regionalize domain as used by the AIPlatform API.
with endpoint_util.AiplatformEndpointOverrides(
version='BETA', region=region
):
api_client = util.GetClient(release_track)
executions_service = api_client.projects_locations_notebookExecutionJobs
return list_pager.YieldFromList(
service=executions_service,
request=executions_util.CreateExecutionListRequest(args, messages),
field='notebookExecutionJobs',
limit=args.limit,
batch_size_attribute='pageSize',
batch_size=args.page_size,
predicate=lambda execution: not executions_util.IsWorkbenchExecution(
execution
),
)
List.detailed_help = _DETAILED_HELP

View File

@@ -0,0 +1,28 @@
# -*- 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.
"""The gcloud colab runtime-templates command group."""
from googlecloudsdk.calliope import base
@base.DefaultUniverseOnly
@base.ReleaseTracks(base.ReleaseTrack.GA, base.ReleaseTrack.BETA)
class RuntimeTemplates(base.Group):
# pylint: disable=g-doc - this is required for help text.
"""Manage Colab Enterprise runtime templates.
For more information about runtime templates, see
https://cloud.google.com/colab/docs/runtimes#runtime-templates
"""

View File

@@ -0,0 +1,78 @@
# -*- coding: utf-8 -*- #
# Copyright 2024 Google Inc. 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 add an IAM policy binding to a Colab Enterprise runtime template."""
from googlecloudsdk.api_lib.colab_enterprise import runtime_templates as runtime_templates_util
from googlecloudsdk.api_lib.colab_enterprise import util
from googlecloudsdk.calliope import base
from googlecloudsdk.command_lib.ai import endpoint_util
from googlecloudsdk.command_lib.colab_enterprise import flags
from googlecloudsdk.command_lib.iam import iam_util
_DETAILED_HELP = {
'DESCRIPTION': """
Add an IAM policy binding to a Colab Enterprise runtime template.
""",
'EXAMPLES': """
To set `someone@example.com` to have the `roles/aiplatform.notebookRuntimeUser` role for a runtime template with id `my-runtime-template` in region `us-central1`, run:
$ {command} my-runtime-template --region=us-central1 --member=user:someone@example.com --role=roles/aiplatform.notebookRuntimeUser
""",
}
@base.DefaultUniverseOnly
@base.ReleaseTracks(base.ReleaseTrack.GA, base.ReleaseTrack.BETA)
class AddIamPolicyBinding(base.Command):
"""Add an IAM policy binding to a Colab Enterprise runtime template."""
@staticmethod
def Args(parser):
"""Register flags for this command."""
flags.AddFlagsToAddIamPolicyBinding(parser)
def Run(self, args):
"""This is what gets called when the user runs this command."""
release_track = self.ReleaseTrack()
messages = util.GetMessages(self.ReleaseTrack())
runtime_template_ref = args.CONCEPTS.runtime_template.Parse()
region = runtime_template_ref.AsDict()['locationsId']
# Override to regionalize domain as used by the AIPlatform API.
with endpoint_util.AiplatformEndpointOverrides(
version='BETA', region=region
):
api_client = util.GetClient(release_track)
runtime_templates_service = (
api_client.projects_locations_notebookRuntimeTemplates
)
iam_policy = runtime_templates_service.GetIamPolicy(
runtime_templates_util.CreateRuntimeTemplateGetIamPolicyRequest(
args, messages
)
)
iam_util.AddBindingToIamPolicy(
messages.GoogleIamV1Binding,
iam_policy,
args.member,
args.role,
)
return runtime_templates_service.SetIamPolicy(
runtime_templates_util.CreateRuntimeTemplateSetIamPolicyRequest(
iam_policy, args, messages
)
)
AddIamPolicyBinding.detailed_help = _DETAILED_HELP

View File

@@ -0,0 +1,88 @@
# -*- 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.
"""Create command for Colab Enterprise Runtime Templates."""
from googlecloudsdk.api_lib.ai import operations
from googlecloudsdk.api_lib.colab_enterprise import runtime_templates as runtime_templates_util
from googlecloudsdk.api_lib.colab_enterprise import util
from googlecloudsdk.calliope import base
from googlecloudsdk.command_lib.ai import endpoint_util
from googlecloudsdk.command_lib.colab_enterprise import flags
_DETAILED_HELP = {
'DESCRIPTION': """
Create a Colab Enterprise runtime template, a VM configuration for your notebook runtimes.
""",
'EXAMPLES': """
To create a runtime template in region 'us-central1' with the display name 'my-runtime-template', run:
$ {command} --region=us-central1 --display-name=my-runtime-template
To create a runtime template for a VM with machine type n1-standard-2 and one NVIDIA_TESLA_V100 accelerator, run:
$ {command} --machine-type=n1-standard-2 --accelerator-type=NVIDIA_TESLA_V100 --accelerator-count=1 --region=us-central1 --display-name=my-runtime-template
To create a runtime template that disables end user credential access, run:
$ {command} --no-enable-euc --region=us-central1 --display-name=my-runtime-template
""",
}
@base.DefaultUniverseOnly
@base.ReleaseTracks(base.ReleaseTrack.GA, base.ReleaseTrack.BETA)
class Create(base.CreateCommand):
"""Create a runtime template."""
@staticmethod
def Args(parser):
"""Register flags for this command."""
flags.AddCreateRuntimeTemplateFlags(parser)
def Run(self, args):
"""This is what gets called when the user runs this command."""
release_track = self.ReleaseTrack()
messages = util.GetMessages(self.ReleaseTrack())
region_ref = args.CONCEPTS.region.Parse()
region = region_ref.AsDict()['locationsId']
# Override to regionalize domain as used by the AIPlatform API.
with endpoint_util.AiplatformEndpointOverrides(
version='BETA', region=region
):
api_client = util.GetClient(release_track)
runtime_templates_service = (
api_client.projects_locations_notebookRuntimeTemplates
)
operation = runtime_templates_service.Create(
runtime_templates_util.CreateRuntimeTemplateCreateRequest(
args, messages
)
)
return util.WaitForOpMaybe(
operations_client=operations.OperationsClient(client=api_client),
op=operation,
op_ref=runtime_templates_util.ParseRuntimeTemplateOperation(
operation.name
),
asynchronous=util.GetAsyncConfig(args),
kind='runtime template',
log_method='create',
message='Waiting for runtime template to be created...',
)
Create.detailed_help = _DETAILED_HELP

View File

@@ -0,0 +1,81 @@
# -*- 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.
"""Delete command for Colab Enterprise Runtime Templates."""
from googlecloudsdk.api_lib.ai import operations
from googlecloudsdk.api_lib.colab_enterprise import runtime_templates as runtime_templates_util
from googlecloudsdk.api_lib.colab_enterprise import util
from googlecloudsdk.calliope import base
from googlecloudsdk.command_lib.ai import endpoint_util
from googlecloudsdk.command_lib.colab_enterprise import flags
_DETAILED_HELP = {
'DESCRIPTION': """
Delete a Colab Enterprise notebook runtime template.
""",
'EXAMPLES': """
To delete a runtime template with id 'my-runtime-template' in region 'us-central1', run:
$ {command} my-runtime-template --region=us-central1
""",
}
@base.DefaultUniverseOnly
@base.ReleaseTracks(base.ReleaseTrack.GA, base.ReleaseTrack.BETA)
class Delete(base.DeleteCommand):
"""Delete a runtime template."""
@staticmethod
def Args(parser):
"""Register flags for this command."""
flags.AddDeleteRuntimeTemplateFlags(parser)
def Run(self, args):
"""This is what gets called when the user runs this command."""
release_track = self.ReleaseTrack()
messages = util.GetMessages(self.ReleaseTrack())
runtime_template_ref = args.CONCEPTS.runtime_template.Parse()
region = runtime_template_ref.AsDict()['locationsId']
# Override to regionalize domain as used by the AIPlatform API.
with endpoint_util.AiplatformEndpointOverrides(
version='BETA', region=region
):
api_client = util.GetClient(release_track)
runtime_templates_service = (
api_client.projects_locations_notebookRuntimeTemplates
)
operation = runtime_templates_service.Delete(
runtime_templates_util.CreateRuntimeTemplateDeleteRequest(
args, messages
)
)
return util.WaitForOpMaybe(
operations_client=operations.OperationsClient(client=api_client),
op=operation,
op_ref=runtime_templates_util.ParseRuntimeTemplateOperation(
operation.name
),
log_method='delete',
kind='runtime template',
asynchronous=util.GetAsyncConfig(args),
message='Waiting for runtime template to be deleted...',
resource=args.CONCEPTS.runtime_template.Parse().RelativeName(),
)
Delete.detailed_help = _DETAILED_HELP

View File

@@ -0,0 +1,68 @@
# -*- 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.
"""Describe command for Colab Enterprise Runtime Templates."""
from googlecloudsdk.api_lib.colab_enterprise import runtime_templates as runtime_templates_util
from googlecloudsdk.api_lib.colab_enterprise import util
from googlecloudsdk.calliope import base
from googlecloudsdk.command_lib.ai import endpoint_util
from googlecloudsdk.command_lib.colab_enterprise import flags
_DETAILED_HELP = {
'DESCRIPTION': """
Describe a Colab Enterprise notebook runtime template.
""",
'EXAMPLES': """
To describe a runtime template with id 'my-runtime-template' in region 'us-central1', run:
$ {command} my-runtime-template --region=us-central1
""",
}
@base.DefaultUniverseOnly
@base.ReleaseTracks(base.ReleaseTrack.GA, base.ReleaseTrack.BETA)
class Describe(base.DescribeCommand):
"""Describe a runtime template."""
@staticmethod
def Args(parser):
"""Register flags for this command."""
flags.AddDescribeRuntimeTemplateFlags(parser)
def Run(self, args):
"""This is what gets called when the user runs this command."""
release_track = self.ReleaseTrack()
messages = util.GetMessages(self.ReleaseTrack())
runtime_template_ref = args.CONCEPTS.runtime_template.Parse()
region = runtime_template_ref.AsDict()['locationsId']
# Override to regionalize domain as used by the AIPlatform API.
with endpoint_util.AiplatformEndpointOverrides(
version='BETA', region=region
):
api_client = util.GetClient(release_track)
runtime_templates_service = (
api_client.projects_locations_notebookRuntimeTemplates
)
return runtime_templates_service.Get(
runtime_templates_util.CreateRuntimeTemplateGetRequest(
args, messages
)
)
Describe.detailed_help = _DETAILED_HELP

View File

@@ -0,0 +1,71 @@
# -*- coding: utf-8 -*- #
# Copyright 2024 Google Inc. 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 the IAM policy for a Colab Enterprise runtime template."""
from googlecloudsdk.api_lib.colab_enterprise import runtime_templates as runtime_templates_util
from googlecloudsdk.api_lib.colab_enterprise import util
from googlecloudsdk.calliope import base
from googlecloudsdk.command_lib.ai import constants
from googlecloudsdk.command_lib.ai import endpoint_util
from googlecloudsdk.command_lib.colab_enterprise import flags
_DETAILED_HELP = {
'DESCRIPTION': """
Get the IAM policy for a Colab Enterprise runtime template.
""",
'EXAMPLES': """
To get the IAM policy for a runtime template with id `my-runtime-template` in region `us-central1`, run:
$ {command} my-runtime-template --location=us-central1 \
--member=user:someone@example.com --role=roles/aiplatform.notebookRuntimeUser
""",
}
@base.DefaultUniverseOnly
@base.ReleaseTracks(base.ReleaseTrack.GA, base.ReleaseTrack.BETA)
class GetIamPolicy(base.ListCommand):
"""Get IAM policy for a Colab Enterprise runtime template."""
@staticmethod
def Args(parser):
"""Register flags for this command."""
flags.AddGetIamPolicyFlags(parser)
# Inherited URI flag from list command does not apply for listing IAM
# policies.
base.URI_FLAG.RemoveFromParser(parser)
def Run(self, args):
"""This is what gets called when the user runs this command."""
release_track = self.ReleaseTrack()
messages = util.GetMessages(self.ReleaseTrack())
runtime_template_ref = args.CONCEPTS.runtime_template.Parse()
region = runtime_template_ref.AsDict()['locationsId']
# Override to regionalize domain as used by the AIPlatform API.
with endpoint_util.AiplatformEndpointOverrides(
version=constants.BETA_VERSION, region=region
):
api_client = util.GetClient(release_track)
runtime_templates_service = (
api_client.projects_locations_notebookRuntimeTemplates
)
return runtime_templates_service.GetIamPolicy(
runtime_templates_util.CreateRuntimeTemplateGetIamPolicyRequest(
args, messages
)
)
GetIamPolicy.detailed_help = _DETAILED_HELP

View File

@@ -0,0 +1,79 @@
# -*- 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.
"""List command for Colab Enterprise Runtime Templates."""
from apitools.base.py import list_pager
from googlecloudsdk.api_lib.colab_enterprise import runtime_templates as runtime_templates_util
from googlecloudsdk.api_lib.colab_enterprise import util
from googlecloudsdk.calliope import base
from googlecloudsdk.command_lib.ai import endpoint_util
from googlecloudsdk.command_lib.colab_enterprise import flags
_DETAILED_HELP = {
'DESCRIPTION': """
List your project's Colab Enterprise notebook runtime templates in a given region.
""",
'EXAMPLES': """
To list your runtime templates in region 'us-central1', run:
$ {command} --region=us-central1
""",
}
@base.DefaultUniverseOnly
@base.ReleaseTracks(base.ReleaseTrack.GA, base.ReleaseTrack.BETA)
class List(base.ListCommand):
"""List your runtime templates."""
@staticmethod
def Args(parser):
"""Register flags for this command."""
flags.AddListRuntimeTemplatesFlags(parser)
parser.display_info.AddFormat("""
table(name.segment(-1):label=ID,
displayName.segment(-1),
name.segment(-3):label=REGION,
machineSpec.machineType.segment(-1))
""")
def Run(self, args):
"""This is what gets called when the user runs this command."""
release_track = self.ReleaseTrack()
messages = util.GetMessages(self.ReleaseTrack())
region_ref = args.CONCEPTS.region.Parse()
region = region_ref.AsDict()['locationsId']
# Override to regionalize domain as used by the AIPlatform API.
with endpoint_util.AiplatformEndpointOverrides(
version='BETA', region=region
):
api_client = util.GetClient(release_track)
runtime_templates_service = (
api_client.projects_locations_notebookRuntimeTemplates
)
return list_pager.YieldFromList(
service=runtime_templates_service,
request=runtime_templates_util.CreateRuntimeTemplateListRequest(
args, messages
),
field='notebookRuntimeTemplates',
limit=args.limit,
batch_size_attribute='pageSize',
batch_size=args.page_size,
)
List.detailed_help = _DETAILED_HELP

View File

@@ -0,0 +1,77 @@
# -*- coding: utf-8 -*- #
# Copyright 2024 Google Inc. 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 remove an IAM policy binding from a Colab Enterprise runtime template."""
from googlecloudsdk.api_lib.colab_enterprise import runtime_templates as runtime_templates_util
from googlecloudsdk.api_lib.colab_enterprise import util
from googlecloudsdk.calliope import base
from googlecloudsdk.command_lib.ai import constants
from googlecloudsdk.command_lib.ai import endpoint_util
from googlecloudsdk.command_lib.colab_enterprise import flags
from googlecloudsdk.command_lib.iam import iam_util
_DETAILED_HELP = {
'DESCRIPTION': """
Remove an IAM policy binding from a Colab Enterprise runtime template.
""",
'EXAMPLES': """
To remove an IAM policy binding of `roles/aiplatform.notebookRuntimeUser` for `someone@example.com`, from the runtime template with id `my-runtime-template` in region `us-central1`, run:
$ {command} my-runtime-template --region=us-central1 --member=user:someone@example.com --role=roles/aiplatform.notebookRuntimeUser
""",
}
@base.DefaultUniverseOnly
@base.ReleaseTracks(base.ReleaseTrack.GA, base.ReleaseTrack.BETA)
class RemoveIamPolicyBinding(base.Command):
"""Remove an IAM policy binding from a Colab Enterprise runtime template."""
@staticmethod
def Args(parser):
"""Register flags for this command."""
flags.AddRemoveIamPolicyBindingFlags(parser)
def Run(self, args):
"""This is what gets called when the user runs this command."""
release_track = self.ReleaseTrack()
messages = util.GetMessages(self.ReleaseTrack())
runtime_template_ref = args.CONCEPTS.runtime_template.Parse()
region = runtime_template_ref.AsDict()['locationsId']
# Override to regionalize domain as used by the AIPlatform API.
with endpoint_util.AiplatformEndpointOverrides(
version=constants.BETA_VERSION, region=region
):
api_client = util.GetClient(release_track)
runtime_templates_service = (
api_client.projects_locations_notebookRuntimeTemplates
)
iam_policy = runtime_templates_service.GetIamPolicy(
runtime_templates_util.CreateRuntimeTemplateGetIamPolicyRequest(
args, messages
)
)
iam_util.RemoveBindingFromIamPolicy(
iam_policy,
args.member,
args.role,
)
return runtime_templates_service.SetIamPolicy(
runtime_templates_util.CreateRuntimeTemplateSetIamPolicyRequest(
iam_policy, args, messages
)
)
RemoveIamPolicyBinding.detailed_help = _DETAILED_HELP

View File

@@ -0,0 +1,69 @@
# -*- coding: utf-8 -*- #
# Copyright 2024 Google Inc. 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 set an IAM policy binding on a Colab Enterprise runtime template."""
from googlecloudsdk.api_lib.colab_enterprise import runtime_templates as runtime_templates_util
from googlecloudsdk.api_lib.colab_enterprise import util
from googlecloudsdk.calliope import base
from googlecloudsdk.command_lib.ai import constants
from googlecloudsdk.command_lib.ai import endpoint_util
from googlecloudsdk.command_lib.colab_enterprise import flags
_DETAILED_HELP = {
'DESCRIPTION': """
Set the IAM policy for a Colab Enterprise runtime template as defined in a JSON or YAML file.
See https://cloud.google.com/iam/docs/managing-policies for details of
the policy file format and contents.
""",
'EXAMPLES': """
To set the IAM policy for a runtime template with id `my-runtime-template` in region `us-central1` to the policy defined in `policy.json`, run:
$ {command} my-runtime-template policy.json --region=us-central1
""",
}
@base.DefaultUniverseOnly
@base.ReleaseTracks(base.ReleaseTrack.GA, base.ReleaseTrack.BETA)
class SetIamPolicy(base.Command):
"""Set IAM policy for a Colab Enterprise runtime template as defined in a JSON or YAML file."""
@staticmethod
def Args(parser):
"""Register flags for this command."""
flags.AddSetIamPolicyBindingFlags(parser)
def Run(self, args):
"""This is what gets called when the user runs this command."""
release_track = self.ReleaseTrack()
messages = util.GetMessages(self.ReleaseTrack())
runtime_template_ref = args.CONCEPTS.runtime_template.Parse()
region = runtime_template_ref.AsDict()['locationsId']
# Override to regionalize domain as used by the AIPlatform API.
with endpoint_util.AiplatformEndpointOverrides(
version=constants.BETA_VERSION, region=region
):
api_client = util.GetClient(release_track)
runtime_templates_service = (
api_client.projects_locations_notebookRuntimeTemplates
)
return runtime_templates_service.SetIamPolicy(
runtime_templates_util.CreateRuntimeTemplateSetIamPolicyRequestFromFile(
args, messages
)
)
SetIamPolicy.detailed_help = _DETAILED_HELP

View File

@@ -0,0 +1,28 @@
# -*- 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.
"""The gcloud colab runtimes command group."""
from googlecloudsdk.calliope import base
@base.DefaultUniverseOnly
@base.ReleaseTracks(base.ReleaseTrack.GA, base.ReleaseTrack.BETA)
class Runtimes(base.Group):
# pylint: disable=g-doc - this is required for help text.
"""Manage Colab Enterprise runtimes.
For more information about runtimes, see
https://cloud.google.com/colab/docs/runtimes#runtimes
"""

View File

@@ -0,0 +1,92 @@
# -*- 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.
"""Assign command for Colab Enterprise Runtimes."""
from googlecloudsdk.api_lib.ai import operations
from googlecloudsdk.api_lib.colab_enterprise import runtimes as runtimes_util
from googlecloudsdk.api_lib.colab_enterprise import util
from googlecloudsdk.calliope import base
from googlecloudsdk.command_lib.ai import endpoint_util
from googlecloudsdk.command_lib.colab_enterprise import flags
_DETAILED_HELP = {
'DESCRIPTION': """
Assign a notebook runtime to run code from your notebook (IPYNB file).
""",
'EXAMPLES': """
To create a runtime in region 'us-central1' with the display name 'my-runtime' and template with id 'my-template', run:
$ {command} --region=us-central1 --display-name=my-runtime --runtime-template=my-template
To create a runtime for user 'USER@DOMAIN.COM', run:
$ {command} --runtime-user=USER@DOMAIN.COM --region=us-central1 --display-name=my-runtime --runtime-template=my-template
""",
}
@base.Deprecate(
is_removed=True,
warning=(
'This command is deprecated. '
'Please use `gcloud beta colab runtimes create` instead.'
),
error=(
'This command has been removed. '
'Please use `gcloud beta colab runtimes create` instead.'
),
)
@base.DefaultUniverseOnly
@base.ReleaseTracks(base.ReleaseTrack.BETA)
class Create(base.CreateCommand):
"""Assign a notebook runtime."""
@staticmethod
def Args(parser):
"""Register flags for this command."""
flags.AddCreateRuntimeFlags(parser)
def Run(self, args):
"""This is what gets called when the user runs this command."""
release_track = self.ReleaseTrack()
messages = util.GetMessages(self.ReleaseTrack())
region_ref = args.CONCEPTS.region.Parse()
region = region_ref.AsDict()['locationsId']
# Override to regionalize domain as used by the AIPlatform API.
with endpoint_util.AiplatformEndpointOverrides(
version='BETA', region=region
):
api_client = util.GetClient(release_track)
runtimes_service = api_client.projects_locations_notebookRuntimes
operation = runtimes_service.Assign(
runtimes_util.CreateRuntimeAssignRequestMessage(args, messages)
)
return util.WaitForOpMaybe(
operations_client=operations.OperationsClient(client=api_client),
op=operation,
op_ref=runtimes_util.ParseRuntimeOperation(operation.name),
asynchronous=util.GetAsyncConfig(args),
kind='runtime',
log_method='create',
message=(
'Assigning a runtime to'
f' {runtimes_util.GetRuntimeUserFromArgsOrProperties(args)}...'
),
)
Create.detailed_help = _DETAILED_HELP

View File

@@ -0,0 +1,81 @@
# -*- 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.
"""Create command for Colab Enterprise Runtimes."""
from googlecloudsdk.api_lib.ai import operations
from googlecloudsdk.api_lib.colab_enterprise import runtimes as runtimes_util
from googlecloudsdk.api_lib.colab_enterprise import util
from googlecloudsdk.calliope import base
from googlecloudsdk.command_lib.ai import endpoint_util
from googlecloudsdk.command_lib.colab_enterprise import flags
_DETAILED_HELP = {
'DESCRIPTION': """
Create a notebook runtime that can be used to run code from your notebook (IPYNB file).
""",
'EXAMPLES': """
To create a runtime in region 'us-central1' with the display name 'my-runtime' and template with id 'my-template', run:
$ {command} --region=us-central1 --display-name=my-runtime --runtime-template=my-template
To create a runtime for user 'USER@DOMAIN.COM', run:
$ {command} --runtime-user=USER@DOMAIN.COM --region=us-central1 --display-name=my-runtime --runtime-template=my-template
""",
}
@base.DefaultUniverseOnly
@base.ReleaseTracks(base.ReleaseTrack.GA, base.ReleaseTrack.BETA)
class Create(base.CreateCommand):
"""Create a notebook runtime."""
@staticmethod
def Args(parser):
"""Register flags for this command."""
flags.AddCreateRuntimeFlags(parser)
def Run(self, args):
"""This is what gets called when the user runs this command."""
release_track = self.ReleaseTrack()
messages = util.GetMessages(self.ReleaseTrack())
region_ref = args.CONCEPTS.region.Parse()
region = region_ref.AsDict()['locationsId']
# Override to regionalize domain as used by the AIPlatform API.
with endpoint_util.AiplatformEndpointOverrides(
version='BETA', region=region
):
api_client = util.GetClient(release_track)
runtimes_service = api_client.projects_locations_notebookRuntimes
operation = runtimes_service.Assign(
runtimes_util.CreateRuntimeAssignRequestMessage(args, messages)
)
return util.WaitForOpMaybe(
operations_client=operations.OperationsClient(client=api_client),
op=operation,
op_ref=runtimes_util.ParseRuntimeOperation(operation.name),
asynchronous=util.GetAsyncConfig(args),
kind='runtime',
log_method='create',
message=(
'Creating a runtime for'
f' {runtimes_util.GetRuntimeUserFromArgsOrProperties(args)}...'
),
)
Create.detailed_help = _DETAILED_HELP

View File

@@ -0,0 +1,79 @@
# -*- 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.
"""Delete command for Colab Enterprise Runtimes."""
from googlecloudsdk.api_lib.ai import operations
from googlecloudsdk.api_lib.colab_enterprise import runtimes as runtimes_util
from googlecloudsdk.api_lib.colab_enterprise import util
from googlecloudsdk.calliope import base
from googlecloudsdk.command_lib.ai import endpoint_util
from googlecloudsdk.command_lib.colab_enterprise import flags
_DETAILED_HELP = {
'DESCRIPTION': """
Delete a Colab Enterprise notebook runtime.
""",
'EXAMPLES': """
To delete a runtime with id 'my-runtime' in region 'us-central1', run:
$ {command} my-runtime --region=us-central1
""",
}
@base.DefaultUniverseOnly
@base.ReleaseTracks(base.ReleaseTrack.GA, base.ReleaseTrack.BETA)
class Delete(base.DeleteCommand):
"""Delete a runtime."""
@staticmethod
def Args(parser):
"""Register flags for this command."""
flags.AddDeleteRuntimeFlags(parser)
def Run(self, args):
"""This is what gets called when the user runs this command."""
release_track = self.ReleaseTrack()
messages = util.GetMessages(self.ReleaseTrack())
runtime_ref = args.CONCEPTS.runtime.Parse()
region = runtime_ref.AsDict()['locationsId']
# Override to regionalize domain as used by the AIPlatform API.
with endpoint_util.AiplatformEndpointOverrides(
version='BETA', region=region
):
api_client = util.GetClient(release_track)
runtime_service = (
api_client.projects_locations_notebookRuntimes
)
operation = runtime_service.Delete(
runtimes_util.CreateRuntimeDeleteRequestMessage(args, messages)
)
return util.WaitForOpMaybe(
operations_client=operations.OperationsClient(client=api_client),
op=operation,
op_ref=runtimes_util.ParseRuntimeOperation(
operation.name
),
log_method='delete',
kind='runtime',
asynchronous=util.GetAsyncConfig(args),
message='Waiting for runtime to be deleted...',
resource=args.CONCEPTS.runtime.Parse().RelativeName(),
)
Delete.detailed_help = _DETAILED_HELP

View File

@@ -0,0 +1,67 @@
# -*- 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.
"""Describe command for Colab Enterprise Runtimes."""
from googlecloudsdk.api_lib.colab_enterprise import runtimes as runtimes_util
from googlecloudsdk.api_lib.colab_enterprise import util
from googlecloudsdk.calliope import base
from googlecloudsdk.command_lib.ai import endpoint_util
from googlecloudsdk.command_lib.colab_enterprise import flags
_DETAILED_HELP = {
'DESCRIPTION': """
Describe a Colab Enterprise notebook runtime.
""",
'EXAMPLES': """
To describe a runtime with id 'my-runtime' in region 'us-central1', run:
$ {command} my-runtime --region=us-central1
""",
}
@base.DefaultUniverseOnly
@base.ReleaseTracks(base.ReleaseTrack.GA, base.ReleaseTrack.BETA)
class Describe(base.DescribeCommand):
"""Describe a runtime."""
@staticmethod
def Args(parser):
"""Register flags for this command."""
flags.AddDescribeRuntimeFlags(parser)
def Run(self, args):
"""This is what gets called when the user runs this command."""
release_track = self.ReleaseTrack()
messages = util.GetMessages(self.ReleaseTrack())
runtime_ref = args.CONCEPTS.runtime.Parse()
region = runtime_ref.AsDict()['locationsId']
# Override to regionalize domain as used by the AIPlatform API.
with endpoint_util.AiplatformEndpointOverrides(
version='BETA', region=region
):
api_client = util.GetClient(release_track)
runtimes_service = (
api_client.projects_locations_notebookRuntimes
)
return runtimes_service.Get(
runtimes_util.CreateRuntimeGetRequestMessage(
args, messages
)
)
Describe.detailed_help = _DETAILED_HELP

View File

@@ -0,0 +1,80 @@
# -*- 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.
"""List command for Colab Enterprise Runtimes."""
from apitools.base.py import list_pager
from googlecloudsdk.api_lib.colab_enterprise import runtimes as runtimes_util
from googlecloudsdk.api_lib.colab_enterprise import util
from googlecloudsdk.calliope import base
from googlecloudsdk.command_lib.ai import endpoint_util
from googlecloudsdk.command_lib.colab_enterprise import flags
_DETAILED_HELP = {
'DESCRIPTION': """
List your project's Colab Enterprise notebook runtimes in a given region.
""",
'EXAMPLES': """
To list your runtimes in region 'us-central1', run:
$ {command} --region=us-central1
""",
}
@base.DefaultUniverseOnly
@base.ReleaseTracks(base.ReleaseTrack.GA, base.ReleaseTrack.BETA)
class List(base.ListCommand):
"""List your project's runtimes."""
@staticmethod
def Args(parser):
"""Register flags for this command."""
flags.AddListRuntimeFlags(parser)
parser.display_info.AddFormat("""
table(name.segment(-1):label=ID,
displayName,
runtimeUser,
notebookRuntimeTemplateRef.notebookRuntimeTemplate.segment(-1):label=RUNTIME_TEMPLATE,
name.segment(-3):label=REGION)
""")
def Run(self, args):
"""This is what gets called when the user runs this command."""
release_track = self.ReleaseTrack()
messages = util.GetMessages(self.ReleaseTrack())
region_ref = args.CONCEPTS.region.Parse()
region = region_ref.AsDict()['locationsId']
# Override to regionalize domain as used by the AIPlatform API.
with endpoint_util.AiplatformEndpointOverrides(
version='BETA', region=region
):
api_client = util.GetClient(release_track)
runtimes_service = (
api_client.projects_locations_notebookRuntimes
)
return list_pager.YieldFromList(
service=runtimes_service,
request=runtimes_util.CreateRuntimeListRequestMessage(
args, messages
),
field='notebookRuntimes',
limit=args.limit,
batch_size_attribute='pageSize',
batch_size=args.page_size,
)
List.detailed_help = _DETAILED_HELP

View File

@@ -0,0 +1,79 @@
# -*- 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.
"""Start command for Colab Enterprise Runtimes."""
from googlecloudsdk.api_lib.ai import operations
from googlecloudsdk.api_lib.colab_enterprise import runtimes as runtimes_util
from googlecloudsdk.api_lib.colab_enterprise import util
from googlecloudsdk.calliope import base
from googlecloudsdk.command_lib.ai import endpoint_util
from googlecloudsdk.command_lib.colab_enterprise import flags
_DETAILED_HELP = {
'DESCRIPTION': """
Start a stopped Colab Enterprise notebook runtime.
""",
'EXAMPLES': """
To start a runtime with id 'my-runtime' in region 'us-central1', run:
$ {command} my-runtime --region=us-central1
""",
}
@base.DefaultUniverseOnly
@base.ReleaseTracks(base.ReleaseTrack.GA, base.ReleaseTrack.BETA)
class Start(base.UpdateCommand):
"""Start a stopped runtime."""
@staticmethod
def Args(parser):
"""Register flags for this command."""
flags.AddStartRuntimeFlags(parser)
def Run(self, args):
"""This is what gets called when the user runs this command."""
release_track = self.ReleaseTrack()
messages = util.GetMessages(self.ReleaseTrack())
runtime_ref = args.CONCEPTS.runtime.Parse()
region = runtime_ref.AsDict()['locationsId']
# Override to regionalize domain as used by the AIPlatform API.
with endpoint_util.AiplatformEndpointOverrides(
version='BETA', region=region
):
api_client = util.GetClient(release_track)
runtime_service = (
api_client.projects_locations_notebookRuntimes
)
operation = runtime_service.Start(
runtimes_util.CreateRuntimeStartRequestMessage(args, messages)
)
return util.WaitForOpMaybe(
operations_client=operations.OperationsClient(client=api_client),
op=operation,
op_ref=runtimes_util.ParseRuntimeOperation(
operation.name
),
log_method='update',
kind='runtime',
asynchronous=util.GetAsyncConfig(args),
message='Starting runtime...',
resource=args.CONCEPTS.runtime.Parse().RelativeName(),
)
Start.detailed_help = _DETAILED_HELP

View File

@@ -0,0 +1,79 @@
# -*- 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.
"""Stop command for Colab Enterprise Runtimes."""
from googlecloudsdk.api_lib.ai import operations
from googlecloudsdk.api_lib.colab_enterprise import runtimes as runtimes_util
from googlecloudsdk.api_lib.colab_enterprise import util
from googlecloudsdk.calliope import base
from googlecloudsdk.command_lib.ai import endpoint_util
from googlecloudsdk.command_lib.colab_enterprise import flags
_DETAILED_HELP = {
'DESCRIPTION': """
Stop a Colab Enterprise notebook runtime.
""",
'EXAMPLES': """
To stop a runtime with id `my-runtime` in region `us-central1`, run:
$ {command} my-runtime --region=us-central1
""",
}
@base.DefaultUniverseOnly
@base.ReleaseTracks(base.ReleaseTrack.GA, base.ReleaseTrack.BETA)
class Stop(base.UpdateCommand):
"""Stop a runtime."""
@staticmethod
def Args(parser):
"""Register flags for this command."""
flags.AddStopRuntimeFlags(parser)
def Run(self, args):
"""This is what gets called when the user runs this command."""
release_track = self.ReleaseTrack()
messages = util.GetMessages(self.ReleaseTrack())
runtime_ref = args.CONCEPTS.runtime.Parse()
region = runtime_ref.AsDict()['locationsId']
# Override to regionalize domain as used by the AIPlatform API.
with endpoint_util.AiplatformEndpointOverrides(
version='BETA', region=region
):
api_client = util.GetClient(release_track)
runtime_service = (
api_client.projects_locations_notebookRuntimes
)
operation = runtime_service.Stop(
runtimes_util.CreateRuntimeStopRequestMessage(args, messages)
)
return util.WaitForOpMaybe(
operations_client=operations.OperationsClient(client=api_client),
op=operation,
op_ref=runtimes_util.ParseRuntimeOperation(
operation.name
),
log_method='update',
kind='runtime',
asynchronous=util.GetAsyncConfig(args),
message='Stopping runtime...',
resource=args.CONCEPTS.runtime.Parse().RelativeName(),
)
Stop.detailed_help = _DETAILED_HELP

View File

@@ -0,0 +1,79 @@
# -*- 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.
"""Upgrade command for Colab Enterprise Runtimes."""
from googlecloudsdk.api_lib.ai import operations
from googlecloudsdk.api_lib.colab_enterprise import runtimes as runtimes_util
from googlecloudsdk.api_lib.colab_enterprise import util
from googlecloudsdk.calliope import base
from googlecloudsdk.command_lib.ai import endpoint_util
from googlecloudsdk.command_lib.colab_enterprise import flags
_DETAILED_HELP = {
'DESCRIPTION': """
Upgrade a Colab Enterprise notebook runtime.
""",
'EXAMPLES': """
To upgrade a runtime with id 'my-runtime' in region 'us-central1', run:
$ {command} my-runtime --region=us-central1
""",
}
@base.DefaultUniverseOnly
@base.ReleaseTracks(base.ReleaseTrack.GA, base.ReleaseTrack.BETA)
class Upgrade(base.UpdateCommand):
"""Upgrade a runtime."""
@staticmethod
def Args(parser):
"""Register flags for this command."""
flags.AddUpgradeRuntimeFlags(parser)
def Run(self, args):
"""This is what gets called when the user runs this command."""
release_track = self.ReleaseTrack()
messages = util.GetMessages(self.ReleaseTrack())
runtime_ref = args.CONCEPTS.runtime.Parse()
region = runtime_ref.AsDict()['locationsId']
# Override to regionalize domain as used by the AIPlatform API.
with endpoint_util.AiplatformEndpointOverrides(
version='BETA', region=region
):
api_client = util.GetClient(release_track)
runtime_service = (
api_client.projects_locations_notebookRuntimes
)
operation = runtime_service.Upgrade(
runtimes_util.CreateRuntimeUpgradeRequestMessage(args, messages)
)
return util.WaitForOpMaybe(
operations_client=operations.OperationsClient(client=api_client),
op=operation,
op_ref=runtimes_util.ParseRuntimeOperation(
operation.name
),
log_method='update',
kind='runtime',
asynchronous=util.GetAsyncConfig(args),
message='Upgrading runtime...',
resource=args.CONCEPTS.runtime.Parse().RelativeName(),
)
Upgrade.detailed_help = _DETAILED_HELP

View File

@@ -0,0 +1,28 @@
# -*- 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.
"""The gcloud colab schedules command group."""
from googlecloudsdk.calliope import base
@base.DefaultUniverseOnly
@base.ReleaseTracks(base.ReleaseTrack.GA, base.ReleaseTrack.BETA)
class Schedules(base.Group):
# pylint: disable=g-doc - this is required for help text.
"""Manage Colab Enterprise notebook execution schedules.
For more information about schedules, see
https://cloud.google.com/colab/docs/schedule-notebook-run.
"""

View File

@@ -0,0 +1,89 @@
# -*- 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.
"""Create command for Colab Enterprise Schedules."""
from googlecloudsdk.api_lib.colab_enterprise import util
from googlecloudsdk.api_lib.notebook_executor import schedules as schedules_util
from googlecloudsdk.calliope import base
from googlecloudsdk.command_lib.ai import endpoint_util
from googlecloudsdk.command_lib.notebook_executor import flags
from googlecloudsdk.core import log
_DETAILED_HELP = {
'DESCRIPTION': """
Create a schedule to run a Colab Enterprise notebook execution job.
""",
'EXAMPLES': """
To create a schedule in region `us-central1` with the following schedule properties:
- display name: `my-schedule`,
- cron schedule: `TZ=America/Los_Angeles * * * * *`,
- maximum concurrent runs allowed: 1,
- start time: 2025-01-01T00:00:00-06:00,
for a notebook execution job:
- with display name `my-execution`,
- running notebook file from Cloud Storage URI `gs://my-bucket/my-notebook.ipynb`,
- compute configured from runtime template `my-runtime-template-id`,
- running with service account `my-service-account@my-project.iam.gserviceaccount.com`,
- with results uploaded to Cloud Storage bucket `gs://my-bucket/results`
Run the following command:
$ {command} --region=us-central1 --display-name=my-schedule --cron-schedule='TZ=America/Los_Angeles * * * * *' --max-concurrent-runs=1 --start-time=2025-01-01T00:00:00-06:00 --execution-display-name=my-execution --notebook-runtime-template=my-runtime-template-id --gcs-notebook-uri=gs://my-bucket/my-notebook.ipynb --service-account=my-service-account@my-project.iam.gserviceaccount.com --gcs-output-uri=gs://my-bucket/results
""",
}
@base.DefaultUniverseOnly
@base.ReleaseTracks(base.ReleaseTrack.GA, base.ReleaseTrack.BETA)
class Create(base.CreateCommand):
"""Create a schedule."""
@staticmethod
def Args(parser):
"""Register flags for this command.
Args:
parser: argparse parser for the command
"""
flags.AddCreateOrUpdateScheduleFlags(parser, is_update=False)
def Run(self, args):
"""This is what gets called when the user runs this command.
Args:
args: argparse.Namespace, The arguments that this command was invoked
with.
Returns:
The API response for creating the schedule.
"""
release_track = self.ReleaseTrack()
messages = util.GetMessages(self.ReleaseTrack())
region_ref = args.CONCEPTS.region.Parse()
region = region_ref.AsDict()['locationsId']
with endpoint_util.AiplatformEndpointOverrides(
version='BETA', region=region
):
api_client = util.GetClient(release_track)
schedules_service = api_client.projects_locations_schedules
schedule = schedules_service.Create(
schedules_util.CreateScheduleCreateRequest(args, messages)
)
log.CreatedResource(resource=schedule.name, kind='schedule')
return schedule
Create.detailed_help = _DETAILED_HELP

View File

@@ -0,0 +1,82 @@
# -*- 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.
"""Delete command for Colab Enterprise Schedules."""
from googlecloudsdk.api_lib.ai import operations
from googlecloudsdk.api_lib.colab_enterprise import util
from googlecloudsdk.api_lib.notebook_executor import schedules as schedules_util
from googlecloudsdk.calliope import base
from googlecloudsdk.command_lib.ai import endpoint_util
from googlecloudsdk.command_lib.notebook_executor import flags
_DETAILED_HELP = {
'DESCRIPTION': """
Delete a Colab Enterprise notebook execution schedule.
""",
'EXAMPLES': """
To delete a schedule with id `my-schedule`, in region `us-central1`, run:
$ {command} my-schedule --region=us-central1
""",
}
@base.DefaultUniverseOnly
@base.ReleaseTracks(base.ReleaseTrack.GA, base.ReleaseTrack.BETA)
class Delete(base.DeleteCommand):
"""Delete a schedule."""
@staticmethod
def Args(parser):
"""Register flags for this command."""
flags.AddDeleteScheduleFlags(parser)
def Run(self, args):
"""This is what gets called when the user runs this command."""
release_track = self.ReleaseTrack()
messages = util.GetMessages(self.ReleaseTrack())
schedule_ref = args.CONCEPTS.schedule.Parse()
region = schedule_ref.AsDict()['locationsId']
with endpoint_util.AiplatformEndpointOverrides(
version='BETA', region=region
):
api_client = util.GetClient(release_track)
schedules_service = (
api_client.projects_locations_schedules
)
schedules_util.ValidateAndGetColabSchedule(
args, messages, schedules_service
)
operation = schedules_service.Delete(
schedules_util.CreateScheduleDeleteRequest(
args, messages
)
)
return util.WaitForOpMaybe(
operations_client=operations.OperationsClient(client=api_client),
op=operation,
op_ref=schedules_util.ParseScheduleOperation(
operation.name
),
asynchronous=util.GetAsyncConfig(args),
kind='schedule',
log_method='delete',
message='Waiting for schedule to be deleted...',
)
Delete.detailed_help = _DETAILED_HELP

View File

@@ -0,0 +1,76 @@
# -*- 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.
"""Describe command for Colab Enterprise Schedules."""
from googlecloudsdk.api_lib.colab_enterprise import util
from googlecloudsdk.api_lib.notebook_executor import schedules as schedules_util
from googlecloudsdk.calliope import base
from googlecloudsdk.command_lib.ai import endpoint_util
from googlecloudsdk.command_lib.notebook_executor import flags
_DETAILED_HELP = {
'DESCRIPTION': """
Describe a Colab Enterprise notebook execution schedule.
""",
'EXAMPLES': """
To describe a schedule with id `my-schedule` in region `us-central1`, run:
$ {command} my-schedule --region=us-central1
""",
}
@base.DefaultUniverseOnly
@base.ReleaseTracks(base.ReleaseTrack.GA, base.ReleaseTrack.BETA)
class Describe(base.DescribeCommand):
"""Describe a schedule."""
@staticmethod
def Args(parser):
"""Register flags for this command.
Args:
parser: argparse parser for the command
"""
flags.AddDescribeScheduleFlags(parser)
def Run(self, args):
"""This is what gets called when the user runs this command.
Args:
args: argparse.Namespace, The arguments that this command was invoked
with.
Returns:
The API response for describing the schedule.
"""
release_track = self.ReleaseTrack()
messages = util.GetMessages(self.ReleaseTrack())
schedule_ref = args.CONCEPTS.schedule.Parse()
region = schedule_ref.AsDict()['locationsId']
# Override to regionalize domain as used by the AIPlatform API.
with endpoint_util.AiplatformEndpointOverrides(
version='BETA', region=region
):
api_client = util.GetClient(release_track)
schedules_service = api_client.projects_locations_schedules
return schedules_util.ValidateAndGetColabSchedule(
args, messages, schedules_service
)
Describe.detailed_help = _DETAILED_HELP

View File

@@ -0,0 +1,83 @@
# -*- 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.
"""List command for Colab Enterprise Schedules."""
from apitools.base.py import list_pager
from googlecloudsdk.api_lib.colab_enterprise import util
from googlecloudsdk.api_lib.notebook_executor import schedules as schedules_util
from googlecloudsdk.calliope import base
from googlecloudsdk.command_lib.ai import endpoint_util
from googlecloudsdk.command_lib.notebook_executor import flags
_DETAILED_HELP = {
'DESCRIPTION': """
List your project's Colab Enterprise notebook execution schedules in a given region.
""",
'EXAMPLES': """
To list your schedules in region `us-central1`, run:
$ {command} --region=us-central1
""",
}
@base.DefaultUniverseOnly
@base.ReleaseTracks(base.ReleaseTrack.GA, base.ReleaseTrack.BETA)
class List(base.ListCommand):
"""List your Colab Enterprise notebook execution schedules."""
@staticmethod
def Args(parser):
"""Register flags for this command."""
flags.AddListSchedulesFlags(parser)
parser.display_info.AddFormat("""
table(name.segment(-1):label=ID,
displayName,
state,
cron,
nextRunTime)
""")
def Run(self, args):
"""This is what gets called when the user runs this command."""
release_track = self.ReleaseTrack()
messages = util.GetMessages(self.ReleaseTrack())
region_ref = args.CONCEPTS.region.Parse()
region = region_ref.AsDict()['locationsId']
# Override to regionalize domain as used by the AIPlatform API.
with endpoint_util.AiplatformEndpointOverrides(
version='BETA', region=region
):
api_client = util.GetClient(release_track)
schedules_service = (
api_client.projects_locations_schedules
)
return list_pager.YieldFromList(
service=schedules_service,
request=schedules_util.CreateScheduleListRequest(
args, messages
),
field='schedules',
limit=args.limit,
batch_size_attribute='pageSize',
batch_size=args.page_size,
predicate=lambda schedule: not schedules_util.FilterWorkbenchSchedule(
schedule
),
)
List.detailed_help = _DETAILED_HELP

View File

@@ -0,0 +1,72 @@
# -*- 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.
"""Pause command for Colab Enterprise Schedules."""
from googlecloudsdk.api_lib.colab_enterprise import util
from googlecloudsdk.api_lib.notebook_executor import schedules as schedules_util
from googlecloudsdk.calliope import base
from googlecloudsdk.command_lib.ai import endpoint_util
from googlecloudsdk.command_lib.notebook_executor import flags
from googlecloudsdk.core import log
_DETAILED_HELP = {
'DESCRIPTION': """
Pause a Colab Enterprise notebook execution schedule.
""",
'EXAMPLES': """
To pause a schedule with id `my-schedule`, in region `us-central1`, run:
$ {command} my-schedule --region=us-central1
""",
}
@base.DefaultUniverseOnly
@base.ReleaseTracks(base.ReleaseTrack.GA, base.ReleaseTrack.BETA)
class Pause(base.UpdateCommand):
"""Pause a schedule."""
@staticmethod
def Args(parser):
"""Register flags for this command."""
flags.AddPauseScheduleFlags(parser)
def Run(self, args):
"""This is what gets called when the user runs this command."""
release_track = self.ReleaseTrack()
messages = util.GetMessages(self.ReleaseTrack())
schedule_ref = args.CONCEPTS.schedule.Parse()
region = schedule_ref.AsDict()['locationsId']
with endpoint_util.AiplatformEndpointOverrides(
version='BETA', region=region
):
api_client = util.GetClient(release_track)
schedules_service = (
api_client.projects_locations_schedules
)
schedules_util.ValidateAndGetColabSchedule(
args, messages, schedules_service
)
api_response = schedules_service.Pause(
schedules_util.CreateSchedulePauseRequest(
args, messages
)
)
log.status.Print('Paused schedule.')
return api_response
Pause.detailed_help = _DETAILED_HELP

View File

@@ -0,0 +1,72 @@
# -*- 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.
"""Resume command for Colab Enterprise Schedules."""
from googlecloudsdk.api_lib.colab_enterprise import util
from googlecloudsdk.api_lib.notebook_executor import schedules as schedules_util
from googlecloudsdk.calliope import base
from googlecloudsdk.command_lib.ai import endpoint_util
from googlecloudsdk.command_lib.notebook_executor import flags
from googlecloudsdk.core import log
_DETAILED_HELP = {
'DESCRIPTION': """
Resume a Colab Enterprise notebook execution schedule.
""",
'EXAMPLES': """
To resume a paused schedule with id `my-schedule`, in region `us-central1`, run:
$ {command} my-schedule --region=us-central1
""",
}
@base.DefaultUniverseOnly
@base.ReleaseTracks(base.ReleaseTrack.GA, base.ReleaseTrack.BETA)
class Resume(base.UpdateCommand):
"""Resume a schedule."""
@staticmethod
def Args(parser):
"""Register flags for this command."""
flags.AddResumeScheduleFlags(parser)
def Run(self, args):
"""This is what gets called when the user runs this command."""
release_track = self.ReleaseTrack()
messages = util.GetMessages(self.ReleaseTrack())
schedule_ref = args.CONCEPTS.schedule.Parse()
region = schedule_ref.AsDict()['locationsId']
with endpoint_util.AiplatformEndpointOverrides(
version='BETA', region=region
):
api_client = util.GetClient(release_track)
schedules_service = (
api_client.projects_locations_schedules
)
schedules_util.ValidateAndGetColabSchedule(
args, messages, schedules_service
)
api_response = schedules_service.Resume(
schedules_util.CreateScheduleResumeRequest(
args, messages
)
)
log.status.Print('Resumed schedule.')
return api_response
Resume.detailed_help = _DETAILED_HELP

View File

@@ -0,0 +1,83 @@
# -*- 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.
"""Update command for Colab Enterprise Schedules."""
from googlecloudsdk.api_lib.colab_enterprise import util
from googlecloudsdk.api_lib.notebook_executor import schedules as schedules_util
from googlecloudsdk.calliope import base
from googlecloudsdk.command_lib.ai import endpoint_util
from googlecloudsdk.command_lib.notebook_executor import flags
from googlecloudsdk.core import log
_DETAILED_HELP = {
'DESCRIPTION': """
Update a Colab Enterprise notebook execution schedule.
""",
'EXAMPLES': """
To update the cron schedule and max runs of a schedule with id `my-schedule`, in region `us-central1`, run:
$ {command} my-schedule --region=us-central1 --cron-schedule='TZ=America/Los_Angeles 0 0 * * 0' --max-runs=2
""",
}
@base.DefaultUniverseOnly
@base.ReleaseTracks(base.ReleaseTrack.GA, base.ReleaseTrack.BETA)
class Update(base.UpdateCommand):
"""Update a schedule."""
@staticmethod
def Args(parser):
"""Register flags for this command.
Args:
parser: argparse parser for the command
"""
flags.AddCreateOrUpdateScheduleFlags(parser, is_update=True)
def Run(self, args):
"""This is what gets called when the user runs this command.
Args:
args: argparse.Namespace, The arguments that this command was invoked
with.
Returns:
The API response for updating the schedule.
"""
release_track = self.ReleaseTrack()
messages = util.GetMessages(self.ReleaseTrack())
schedule_ref = args.CONCEPTS.schedule.Parse()
region = schedule_ref.AsDict()['locationsId']
with endpoint_util.AiplatformEndpointOverrides(
version='BETA', region=region
):
api_client = util.GetClient(release_track)
schedules_service = (
api_client.projects_locations_schedules
)
# Although API will error if schedule is not of notebook execution type,
# validate client-side for better ux / consistency with other commands.
schedules_util.ValidateAndGetColabSchedule(
args, messages, schedules_service
)
schedule = schedules_service.Patch(
schedules_util.CreateSchedulePatchRequest(args, messages)
)
log.UpdatedResource(resource=schedule.name, kind='schedule')
return schedule
Update.detailed_help = _DETAILED_HELP