예제 #1
0
def _upload_and_get_command(boto3_factory, args, job_s3_folder, job_name, config, log):
    """
    Get command by parsing args and config.

    The function will also perform an s3 upload, if needed.
    :param boto3_factory: initialized Boto3ClientFactory object
    :param args: input arguments
    :param job_s3_folder: S3 folder for the job files
    :param job_name: job name
    :param config: config object
    :param log: log
    :return: command to submit
    """
    # create S3 folder for the job
    s3_uploader = S3Uploader(boto3_factory, config.s3_bucket, job_s3_folder)

    # upload input files, if there
    if args.input_file:
        for file in args.input_file:
            s3_uploader.put_file(file, os.path.basename(file))

    # upload command, if needed
    if args.command_file or not sys.stdin.isatty() or args.env:
        # define job script name
        job_script = job_name + ".sh"
        log.info("Using command-file option or stdin. Job script name: %s" % job_script)

        env_file = None
        if args.env:
            env_file = job_name + ".env.sh"
            # get environment variables and upload file used to extend the submission environment
            env_blacklist = args.env_blacklist if args.env_blacklist else config.env_blacklist
            _get_env_and_upload(s3_uploader, args.env, env_blacklist, env_file, log)

        # upload job script
        if args.command_file:
            # existing script file
            try:
                s3_uploader.put_file(args.command, job_script)
            except Exception as e:
                fail("Error creating job script. Failed with exception: %s" % e)
        elif not sys.stdin.isatty():
            # stdin
            _get_stdin_and_upload(s3_uploader, job_script)

        # define command to execute
        bash_command = _compose_bash_command(args, config.s3_bucket, config.region, job_s3_folder, job_script, env_file)
        command = ["/bin/bash", "-c", bash_command]
    elif type(args.command) == str:
        log.info("Using command parameter")
        command = [args.command] + args.arguments
    else:
        fail("Unexpected error. Command cannot be empty.")
    log.info("Command: %s" % shell_join(command))
    return command
예제 #2
0
def _compose_bash_command(args, s3_bucket, region, job_s3_folder, job_script,
                          env_file):
    """
    Define bash command to execute.

    :param args: input arguments
    :param s3_bucket: S3 bucket
    :param region: AWS region
    :param job_s3_folder: S3 job folder
    :param job_script: job script file
    :param env_file: environment file
    :return: composed bash command
    """
    command_args = shell_join(args.arguments)
    # download awscli, if required.
    bash_command = []
    if args.awscli:
        bash_command.append(
            "curl -O https://bootstrap.pypa.io/get-pip.py >/dev/null 2>&1 && "
            "python get-pip.py --user >/dev/null 2>&1 && "
            "export PATH=~/.local/bin:$PATH >/dev/null 2>&1 && "
            "pip install awscli --upgrade --user >/dev/null 2>&1")

    # set working directory
    if args.working_dir:
        # create and move to the working dir specified by the user
        bash_command.append('mkdir -p "{JOB_WD}" && cd "{JOB_WD}"'.format(
            JOB_WD=args.working_dir))
    else:
        if args.parent_working_dir:
            # create and move to the parent working dir specified by the user
            bash_command.append(
                'mkdir -p "{JOB_PARENT_WD}" && cd "{JOB_PARENT_WD}"'.format(
                    JOB_PARENT_WD=args.parent_working_dir))

        # create subfolder named job-<$AWS_BATCH_JOB_ID>
        bash_command.append(
            "mkdir -p job-${AWS_BATCH_JOB_ID} && cd job-${AWS_BATCH_JOB_ID}")

    # download all job files to the job folder
    bash_command.append(
        "aws s3 --region {REGION} sync s3://{BUCKET}/{S3_FOLDER} . >/dev/null".
        format(REGION=region, BUCKET=s3_bucket, S3_FOLDER=job_s3_folder))
    if env_file:  # source the environment file
        bash_command.append("source {ENV_FILE}".format(ENV_FILE=env_file))

    # execute the job script + arguments
    bash_command.append("chmod +x {SCRIPT} && ./{SCRIPT} {ARGS}".format(
        SCRIPT=job_script, ARGS=command_args))
    return " && ".join(bash_command)
예제 #3
0
 def _get_command(container):
     command = container.get("command", [])
     if not command:
         return "-"
     return shell_join(command)
예제 #4
0
    def __add_jobs(self, jobs, details):
        """
        Get job info from AWS Batch and add to the output
        :param jobs: list of jobs items (output of the list_jobs function)
        :param details: ask for job details
        """
        try:
            if jobs:
                self.log.debug("Adding jobs to the output (%s)" % jobs)
                if details:
                    self.log.info("Asking for jobs details")
                    jobs_to_show = []
                    for index in range(0, len(jobs), 100):
                        jobs_chunk = jobs[index:index + 100]
                        job_ids = []
                        for job in jobs_chunk:
                            job_ids.append(job['jobId'])
                        jobs_to_show.extend(
                            self.batch_client.describe_jobs(
                                jobs=job_ids)['jobs'])
                else:
                    jobs_to_show = jobs

                for job in jobs_to_show:
                    nodes = 1
                    if 'nodeProperties' in job:
                        # MNP job
                        container = job['nodeProperties'][
                            'nodeRangeProperties'][0]['container']
                        nodes = job['nodeProperties']['numNodes']
                    elif 'container' in job:
                        container = job['container']
                    else:
                        container = {}

                    if is_job_array(job):
                        # parent job array
                        job_id = '{0}[{1}]'.format(
                            job['jobId'], job['arrayProperties']['size'])
                        log_stream = '-'
                        log_stream_url = '-'
                    else:
                        job_id = job['jobId']
                        if 'logStreamName' in container:
                            log_stream = container.get('logStreamName')
                            log_stream_url = _compose_log_stream_url(
                                self.boto3_factory.region, log_stream)
                        else:
                            log_stream = '-'
                            log_stream_url = '-'

                    command = container.get('command', [])
                    self.log.debug("Adding job to the output (%s)", job)
                    job = Job(job_id=job_id,
                              name=job['jobName'],
                              creation_time=convert_to_date(job['createdAt']),
                              start_time=convert_to_date(job['startedAt'])
                              if 'startedAt' in job else '-',
                              stop_time=convert_to_date(job['stoppedAt'])
                              if 'stoppedAt' in job else '-',
                              status=job.get('status', 'UNKNOWN'),
                              status_reason=job.get('statusReason', '-'),
                              job_definition=get_job_definition_name_by_arn(
                                  job['jobDefinition'], version=True)
                              if 'jobQueue' in job else '-',
                              queue=job['jobQueue'].split('/')[1]
                              if 'jobQueue' in job else '-',
                              command=shell_join(command) if command else '-',
                              reason=container.get('reason', '-'),
                              exit_code=container.get('exitCode', '-'),
                              vcpus=container.get('vcpus', '-'),
                              memory=container.get('memory', '-'),
                              nodes=nodes,
                              log_stream=log_stream,
                              log_stream_url=log_stream_url)
                    self.output.add(job)
        except KeyError as e:
            fail("Error building Job item. Key (%s) not found." % e)
        except Exception as e:
            fail("Error adding jobs to the output. Failed with exception: %s" %
                 e)
예제 #5
0
def _upload_and_get_command(boto3_factory, args, job_name, region, s3_bucket,
                            log):
    """
    Get command by parsing args and config.
    The function will also perform an s3 upload, if needed
    :param boto3_factory: initialized Boto3ClientFactory object
    :param args: input arguments
    :param job_name: job name
    :param region: region
    :param s3_bucket: S3 bucket to use
    :param log: log
    :return: command to submit
    """
    if args.command_file or not sys.stdin.isatty():
        # define job script name
        job_script = 'job-{0}-{1}.sh'.format(job_name, int(time.time() * 1000))
        log.info("Using command-file option or stdin. Job script name: %s" %
                 job_script)

        # upload job script
        if args.command_file:
            # existing script file
            try:
                _upload_to_s3(boto3_factory, s3_bucket, args.command,
                              job_script, args.timeout)
            except Exception as e:
                fail("Error creating job script. Failed with exception: %s" %
                     e)
        else:
            try:
                # copy stdin to temporary file
                with os.fdopen(sys.stdin.fileno(), 'rb') as src:
                    with open(job_script, 'wb') as dst:
                        shutil.copyfileobj(src, dst)

                _upload_to_s3(boto3_factory, s3_bucket, job_script, job_script,
                              args.timeout)
            except (OSError, Exception) as e:
                fail("Error creating job script. Failed with exception: %s" %
                     e)
            finally:
                # remove temporary file
                if os.path.exists(job_script):
                    os.remove(job_script)

        # define command to execute
        command_args = shell_join(args.arguments)
        # download awscli, if required. TODO: remove
        s3_command = "curl -O https://bootstrap.pypa.io/get-pip.py >/dev/null 2>&1; " \
                     "python get-pip.py --user >/dev/null 2>&1; " \
                     "export PATH=~/.local/bin:$PATH >/dev/null 2>&1; " \
                     "pip install awscli --upgrade --user >/dev/null 2>&1; " if args.awscli else ""
        s3_command += "aws s3 --region {REGION} cp s3://{BUCKET}/{SCRIPT} /tmp/{SCRIPT}; " \
                      "bash /tmp/{SCRIPT} {ARGS}".format(REGION=region, BUCKET=s3_bucket,
                                                         SCRIPT=_prepend_s3_folder(job_script), ARGS=command_args)
        command = ['/bin/bash', '-c', s3_command]
    elif type(args.command) == str:
        log.info("Using command parameter")
        command = [args.command] + args.arguments
    else:
        fail("Unexpected error. Command cannot be empty.")
    log.info("Command: %s" % shell_join(command))
    return command