462 lines
15 KiB
Python
462 lines
15 KiB
Python
# -*- coding: utf-8 -*- #
|
|
# Copyright 2022 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 submit a specified Batch job."""
|
|
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import unicode_literals
|
|
|
|
import datetime
|
|
|
|
from apitools.base.protorpclite.messages import DecodeError
|
|
from apitools.base.py import encoding
|
|
from googlecloudsdk.api_lib.batch import jobs
|
|
from googlecloudsdk.calliope import arg_parsers
|
|
from googlecloudsdk.calliope import base
|
|
from googlecloudsdk.command_lib.batch import resource_args
|
|
from googlecloudsdk.command_lib.util.apis import arg_utils
|
|
from googlecloudsdk.core import exceptions
|
|
from googlecloudsdk.core import log
|
|
from googlecloudsdk.core import yaml
|
|
from googlecloudsdk.core.util import files
|
|
|
|
|
|
def _CommonArgs(parser):
|
|
"""Register flags for this command.
|
|
|
|
Args:
|
|
parser: An argparse.ArgumentParser-like object. It is mocked out in order to
|
|
capture some information, but behaves like an ArgumentParser.
|
|
|
|
Returns:
|
|
network_group flag groups.
|
|
"""
|
|
network_group = parser.add_group()
|
|
network_group.add_argument(
|
|
'--network',
|
|
required=True,
|
|
type=str,
|
|
help="""The URL for the network resource.
|
|
Must specify subnetwork as well if network is specified""",
|
|
)
|
|
network_group.add_argument(
|
|
'--subnetwork',
|
|
required=True,
|
|
type=str,
|
|
help="""The URL for the subnetwork resource.
|
|
Must specify network as well if subnetwork is specified""",
|
|
)
|
|
network_group.add_argument(
|
|
'--no-external-ip-address',
|
|
action='store_true',
|
|
default=False,
|
|
help="""Required if no external public IP address
|
|
is attached to the VM. If no external public IP address,
|
|
additional configuration is required to allow the VM
|
|
to access Google Services.""",
|
|
)
|
|
|
|
task_spec_group = parser.add_group(required=True)
|
|
task_spec_group.add_argument(
|
|
'--config',
|
|
type=arg_parsers.FileContents(),
|
|
help="""The file path of the job config file in either JSON or YAML format.
|
|
It also supports direct input from stdin with '-' or HereDoc
|
|
(in shells with HereDoc support like Bash) with '- <<DELIMITER'. """,
|
|
)
|
|
|
|
runnable_group = task_spec_group.add_group(
|
|
mutex=True,
|
|
help="""Either specify the config file for the job or
|
|
the first runnable in the task spec. Specify either a script file or
|
|
container arguments for the first runnable in the task spec.""",
|
|
)
|
|
|
|
script_group = runnable_group.add_group(
|
|
mutex=True,
|
|
help="""Either specify a path to a script file to run or provide
|
|
inline text to execute directly.""",
|
|
)
|
|
script_group.add_argument(
|
|
'--script-file-path',
|
|
help="""Path to script file to run as first runnable in task spec.
|
|
File path should be a valid path on the instance volume.""",
|
|
)
|
|
script_group.add_argument(
|
|
'--script-text',
|
|
type=str,
|
|
help="""Text to run as first runnable in task spec.""",
|
|
)
|
|
|
|
container_group = runnable_group.add_group(
|
|
help="""Options to specify the container arguments for the first
|
|
runnable in the task spec."""
|
|
)
|
|
container_group.add_argument(
|
|
'--container-image-uri',
|
|
help="""The URI to pull the container image from.""",
|
|
)
|
|
container_group.add_argument(
|
|
'--container-entrypoint',
|
|
help="""Overrides the `ENTRYPOINT` specified in the container.""",
|
|
)
|
|
container_group.add_argument(
|
|
'--container-commands-file',
|
|
help="""Overrides the `CMD` specified in the container. If there is an
|
|
ENTRYPOINT (either in the container image or with the entrypoint field
|
|
below) then commands are appended as arguments to the ENTRYPOINT.""",
|
|
)
|
|
|
|
parser.add_argument(
|
|
'--priority',
|
|
type=arg_parsers.BoundedInt(0, 99),
|
|
help='Job priority [0-99] 0 is the lowest priority.',
|
|
)
|
|
parser.add_argument(
|
|
'--provisioning-model',
|
|
choices={
|
|
'STANDARD': 'The STANDARD VM provisioning model',
|
|
'SPOT': """The SPOT VM provisioning model. Ideal for fault-tolerant
|
|
workloads that can withstand preemption.""",
|
|
},
|
|
type=arg_utils.ChoiceToEnumName,
|
|
help='Specify the allowed provisioning model for the compute instances',
|
|
)
|
|
parser.add_argument(
|
|
'--machine-type',
|
|
type=str,
|
|
help="""Specify the Compute Engine machine type, for
|
|
example, e2-standard-4. Currently only one machine type is supported.""",
|
|
)
|
|
parser.add_argument(
|
|
'--job-prefix',
|
|
type=str,
|
|
help="""Specify the job prefix. A job ID in the format of
|
|
job prefix + %Y%m%d-%H%M%S will be generated. Note that job prefix
|
|
cannot be specified while JOB ID positional argument is
|
|
specified.""",
|
|
)
|
|
|
|
return network_group
|
|
|
|
|
|
def _BuildJobMsg(args, job_msg, batch_msgs):
|
|
"""Build the job API message from the args.
|
|
|
|
Args:
|
|
args: the args from the parser.
|
|
job_msg: the output job message.
|
|
batch_msgs: the related version of the batch message.
|
|
"""
|
|
|
|
if job_msg.taskGroups is None:
|
|
job_msg.taskGroups = []
|
|
if not job_msg.taskGroups:
|
|
job_msg.taskGroups.insert(
|
|
0, batch_msgs.TaskGroup(taskSpec=batch_msgs.TaskSpec(runnables=[]))
|
|
)
|
|
if args.script_file_path:
|
|
job_msg.taskGroups[0].taskSpec.runnables.insert(
|
|
0,
|
|
batch_msgs.Runnable(
|
|
script=batch_msgs.Script(path=args.script_file_path)
|
|
),
|
|
)
|
|
if args.script_text:
|
|
job_msg.taskGroups[0].taskSpec.runnables.insert(
|
|
0, batch_msgs.Runnable(script=batch_msgs.Script(text=args.script_text))
|
|
)
|
|
if (
|
|
args.container_commands_file
|
|
or args.container_image_uri
|
|
or args.container_entrypoint
|
|
):
|
|
container_cmds = []
|
|
if args.container_commands_file:
|
|
container_cmds = files.ReadFileContents(
|
|
args.container_commands_file
|
|
).splitlines()
|
|
job_msg.taskGroups[0].taskSpec.runnables.insert(
|
|
0,
|
|
batch_msgs.Runnable(
|
|
container=batch_msgs.Container(
|
|
entrypoint=args.container_entrypoint,
|
|
imageUri=args.container_image_uri,
|
|
commands=container_cmds,
|
|
)
|
|
),
|
|
)
|
|
|
|
if args.priority:
|
|
job_msg.priority = args.priority
|
|
|
|
# Add default empty allocation policy if there is no allocation policy and an
|
|
# allocation policy is needed by some other argument.
|
|
if job_msg.allocationPolicy is None and (
|
|
args.machine_type
|
|
or (args.network and args.subnetwork)
|
|
or args.provisioning_model
|
|
):
|
|
job_msg.allocationPolicy = batch_msgs.AllocationPolicy()
|
|
|
|
if args.machine_type:
|
|
if job_msg.allocationPolicy.instances is None:
|
|
job_msg.allocationPolicy.instances = []
|
|
if not job_msg.allocationPolicy.instances:
|
|
job_msg.allocationPolicy.instances.insert(
|
|
0, batch_msgs.InstancePolicyOrTemplate()
|
|
)
|
|
if job_msg.allocationPolicy.instances[0].policy is None:
|
|
job_msg.allocationPolicy.instances[0].policy = batch_msgs.InstancePolicy()
|
|
job_msg.allocationPolicy.instances[0].policy.machineType = args.machine_type
|
|
|
|
if args.network and args.subnetwork:
|
|
if job_msg.allocationPolicy.network is None:
|
|
job_msg.allocationPolicy.network = batch_msgs.NetworkPolicy(
|
|
networkInterfaces=[]
|
|
)
|
|
job_msg.allocationPolicy.network.networkInterfaces.insert(
|
|
0,
|
|
batch_msgs.NetworkInterface(
|
|
network=args.network,
|
|
subnetwork=args.subnetwork,
|
|
noExternalIpAddress=args.no_external_ip_address,
|
|
),
|
|
)
|
|
|
|
if args.provisioning_model:
|
|
if job_msg.allocationPolicy.instances is None:
|
|
job_msg.allocationPolicy.instances = []
|
|
if not job_msg.allocationPolicy.instances:
|
|
job_msg.allocationPolicy.instances.insert(
|
|
0, batch_msgs.InstancePolicyOrTemplate()
|
|
)
|
|
if job_msg.allocationPolicy.instances[0].policy is None:
|
|
job_msg.allocationPolicy.instances[0].policy = batch_msgs.InstancePolicy()
|
|
job_msg.allocationPolicy.instances[0].policy.provisioningModel = (
|
|
arg_utils.ChoiceToEnum(
|
|
args.provisioning_model,
|
|
batch_msgs.InstancePolicy.ProvisioningModelValueValuesEnum,
|
|
)
|
|
)
|
|
|
|
|
|
@base.ReleaseTracks(base.ReleaseTrack.GA)
|
|
class Submit(base.Command):
|
|
"""Submit a Batch job.
|
|
|
|
This command creates and submits a Batch job. After you create and
|
|
submit the job, Batch automatically queues, schedules, and executes it.
|
|
|
|
## EXAMPLES
|
|
|
|
To submit a job with a sample JSON configuration file (config.json) and name
|
|
`projects/foo/locations/us-central1/jobs/bar`, run:
|
|
|
|
$ {command} projects/foo/locations/us-central1/jobs/bar --config=config.json
|
|
|
|
To submit a job with a sample YAML configuration file (config.yaml) and
|
|
name projects/foo/locations/us-central1/jobs/bar, run:
|
|
|
|
$ {command} projects/foo/locations/us-central1/jobs/bar --config=config.yaml
|
|
|
|
To submit a job through stdin with a sample job configuration and name
|
|
`projects/foo/locations/us-central1/jobs/bar`, run:
|
|
|
|
$ {command} projects/foo/locations/us-central1/jobs/bar --config=-
|
|
|
|
then input json job config via stdin
|
|
{
|
|
job config
|
|
}
|
|
|
|
To submit a job through HereDoc with a sample job configuration and name
|
|
`projects/foo/locations/us-central1/jobs/bar`, run:
|
|
|
|
$ {command} projects/foo/locations/us-central1/jobs/bar --config=- << EOF
|
|
|
|
{
|
|
job config
|
|
}
|
|
EOF
|
|
|
|
For details about how to define a job's configuration using JSON, see the
|
|
projects.locations.jobs resource in the Batch API Reference.
|
|
If you want to define a job's configuration using YAML, convert the JSON
|
|
syntax to YAML.
|
|
"""
|
|
|
|
@staticmethod
|
|
def Args(parser):
|
|
_CommonArgs(parser)
|
|
resource_args.AddSubmitJobResourceArgs(parser)
|
|
|
|
@classmethod
|
|
def _CreateJobMessage(cls, batch_msgs, config):
|
|
"""Parse into Job message using the config input.
|
|
|
|
Args:
|
|
batch_msgs: Batch defined proto message.
|
|
config: The input content being either YAML or JSON or the HEREDOC
|
|
input.
|
|
|
|
Returns:
|
|
The Parsed job message.
|
|
"""
|
|
try:
|
|
result = encoding.PyValueToMessage(batch_msgs.Job, yaml.load(config))
|
|
except (ValueError, AttributeError, yaml.YAMLParseError):
|
|
try:
|
|
result = encoding.JsonToMessage(batch_msgs.Job, config)
|
|
except (ValueError, DecodeError) as e:
|
|
raise exceptions.Error('Unable to parse config file: {}'.format(e))
|
|
return result
|
|
|
|
def Run(self, args):
|
|
job_ref = args.CONCEPTS.job.Parse()
|
|
location_ref = job_ref.Parent()
|
|
job_id = self._GetJobId(job_ref, args)
|
|
|
|
release_track = self.ReleaseTrack()
|
|
|
|
batch_client = jobs.JobsClient(release_track)
|
|
batch_msgs = batch_client.messages
|
|
job_msg = batch_msgs.Job()
|
|
|
|
if args.config:
|
|
job_msg = self._CreateJobMessage(batch_msgs, args.config)
|
|
|
|
_BuildJobMsg(args, job_msg, batch_msgs)
|
|
|
|
resp = batch_client.Create(job_id, location_ref, job_msg)
|
|
log.status.Print(
|
|
'Job {jobName} was successfully submitted.'.format(jobName=resp.uid)
|
|
)
|
|
return resp
|
|
|
|
def _GetJobId(self, job_ref, args):
|
|
job_id = job_ref.RelativeName().split('/')[-1]
|
|
|
|
if job_id != resource_args.INVALIDID and args.job_prefix:
|
|
raise exceptions.Error(
|
|
'--job-prefix cannot be specified when JOB ID positional '
|
|
'argument is specified'
|
|
)
|
|
# Remove the invalid job_id if no job_id being specified,
|
|
# batch_client would create a valid job_id.
|
|
elif args.job_prefix:
|
|
job_id = (
|
|
args.job_prefix
|
|
+ '-'
|
|
+ datetime.datetime.now().strftime('%Y%m%d-%H%M%S')
|
|
)
|
|
|
|
# The case that both positional JOB ID and prefix are not specified
|
|
elif job_id == resource_args.INVALIDID:
|
|
job_id = None
|
|
|
|
return job_id
|
|
|
|
|
|
@base.ReleaseTracks(base.ReleaseTrack.BETA)
|
|
class SubmitBeta(Submit):
|
|
"""Submit a Batch job.
|
|
|
|
This command creates and submits a Batch job. After you create and
|
|
submit the job, Batch automatically queues, schedules, and executes it.
|
|
|
|
## EXAMPLES
|
|
|
|
To submit a job with a sample JSON configuration file (config.json) and name
|
|
`projects/foo/locations/us-central1/jobs/bar`, run:
|
|
|
|
$ {command} projects/foo/locations/us-central1/jobs/bar --config=config.json
|
|
|
|
To submit a job with a sample YAML configuration file (config.yaml) and
|
|
name projects/foo/locations/us-central1/jobs/bar, run:
|
|
|
|
$ {command} projects/foo/locations/us-central1/jobs/bar --config=config.yaml
|
|
|
|
To submit a job through stdin with a sample job configuration and name
|
|
`projects/foo/locations/us-central1/jobs/bar`, run:
|
|
|
|
$ {command} projects/foo/locations/us-central1/jobs/bar --config=-
|
|
|
|
then input json job config via stdin
|
|
{
|
|
job config
|
|
}
|
|
|
|
To submit a job through HereDoc with a sample job configuration and name
|
|
`projects/foo/locations/us-central1/jobs/bar`, run:
|
|
|
|
$ {command} projects/foo/locations/us-central1/jobs/bar --config=- << EOF
|
|
|
|
{
|
|
job config
|
|
}
|
|
EOF
|
|
|
|
For details about how to define a job's configuration using JSON, see the
|
|
projects.locations.jobs resource in the Batch API Reference.
|
|
If you want to define a job's configuration using YAML, convert the JSON
|
|
syntax to YAML.
|
|
"""
|
|
|
|
|
|
@base.ReleaseTracks(base.ReleaseTrack.ALPHA)
|
|
class SubmitAlpha(SubmitBeta):
|
|
"""Submit a Batch job.
|
|
|
|
This command creates and submits a Batch job. After you create and
|
|
submit the job, Batch automatically queues, schedules, and executes it.
|
|
|
|
## EXAMPLES
|
|
|
|
To submit a job with a sample JSON configuration file (config.json) and name
|
|
`projects/foo/locations/us-central1/jobs/bar`, run:
|
|
|
|
$ {command} projects/foo/locations/us-central1/jobs/bar --config=config.json
|
|
|
|
To submit a job with a sample YAML configuration file (config.yaml) and
|
|
name projects/foo/locations/us-central1/jobs/bar, run:
|
|
|
|
$ {command} projects/foo/locations/us-central1/jobs/bar --config=config.yaml
|
|
|
|
To submit a job through stdin with a sample job configuration and name
|
|
`projects/foo/locations/us-central1/jobs/bar`, run:
|
|
|
|
$ {command} projects/foo/locations/us-central1/jobs/bar --config=-
|
|
|
|
then input json job config via stdin
|
|
{
|
|
job config
|
|
}
|
|
|
|
To submit a job through HereDoc with a sample job configuration and name
|
|
`projects/foo/locations/us-central1/jobs/bar`, run:
|
|
|
|
$ {command} projects/foo/locations/us-central1/jobs/bar --config=- << EOF
|
|
|
|
{
|
|
job config
|
|
}
|
|
EOF
|
|
|
|
For details about how to define a job's configuration using JSON, see the
|
|
projects.locations.jobs resource in the Batch API Reference.
|
|
If you want to define a job's configuration using YAML, convert the JSON
|
|
syntax to YAML.
|
|
"""
|