<commit_before>"""
Manage AWS Batch jobs, queues, and compute environments.
"""

from __future__ import absolute_import, division, print_function, unicode_literals

import os, sys, argparse, base64, collections, io, subprocess, json, time, re
from datetime import datetime

from botocore.exceptions import ClientError
import yaml

from . import logger
from .ls import register_parser, register_listing_parser
from .ecr import ecr_image_name_completer
from .util import Timestamp, paginate
from .util.crypto import ensure_ssh_key
from .util.exceptions import AegeaException
from .util.printing import page_output, tabulate, YELLOW, RED, GREEN, BOLD, ENDC
from .util.aws import (ARN, resources, clients, expect_error_codes, ensure_iam_role, ensure_instance_profile,
                       make_waiter, ensure_vpc, ensure_security_group, ensure_s3_bucket, ensure_log_group,
                       IAMPolicyBuilder, resolve_ami)
from .util.aws.spot import SpotFleetBuilder

bash_cmd_preamble = ["/bin/bash", "-c", 'for i in "$@"; do eval "$i"; done', __name__]

ebs_vol_mgr_shellcode = """iid=$(http http://169.254.169.254/latest/dynamic/instance-identity/document)

aws configure set default.region $(echo "$iid" | jq -r .region)

az=$(echo "$iid" | jq -r .availabilityZone)

vid=$(aws ec2 create-volume --availability-zone $az --size %s --volume-type st1 | jq -r .VolumeId)

aws ec2 create-tags --resource $vid --tags Key=aegea_batch_job,Value=$AWS_BATCH_JOB_ID

trap "umount /mnt || umount -l /mnt; aws ec2 detach-volume --volume-id $vid; let delay=2; while ! aws ec2 describe-volumes --volume-ids $vid | jq -re .Volumes[0].Attachments==[]; do if [[ $delay -gt 30 ]]; then aws ec2 detach-volume --force --volume-id $vid; break; fi; sleep $delay; let delay=$delay*2; done; aws ec2 delete-volume --volume-id $vid" EXIT

while [[ $(aws ec2 describe-volumes --volume-ids $vid | jq -r .Volumes[0].State) != available ]]; do sleep 1; done

# let each process start trying from a different /dev/xvd{letter}
let pid=$$

# when N processes compete, for every success there can be N-1 failures; so the appropriate number of retries is O(N^2)
# let us size this for 10 competitors
# NOTE: the constants 9 and 10 in the $ind and $devnode calculation below are based on strlen("/dev/xvda")
let delay=2+$pid%5
for try in {1..100}; do if [[ $try == 100 ]]; then echo "Unable to mount $vid on $devnode"; exit 1; fi; if [[ $try -gt 1 ]]; then sleep $delay; fi; devices=$(echo /dev/xvd* /dev/xvd{a..z} /dev/xvd{a..z} | sed 's= =\n=g' | sort | uniq -c | sort -n | grep ' 2 ' | awk '{print $2}'); let devcnt=${#devices}/10+1; let ind=$pid%devcnt; devnode=${devices:10*$ind:9};  aws ec2 attach-volume --instance-id $(echo "$iid" | jq -r .instanceId) --volume-id $vid --device $devnode || continue; break; done

# attach-volume is not instantaneous, and describe-volume requests are rate-limited
let delay=5+pid%11
sleep $delay

let try=1
let max_tries=32
while [[ $(aws ec2 describe-volumes --volume-ids $vid | jq -r .Volumes[0].State) != in-use ]]; do if [[ $try == $max_tries ]]; break; fi; let foo=1+$try%5; let delay=2**$foo+pid%11; sleep $delay; done

while [[ ! -e $devnode ]]; do sleep 1; done

mkfs.ext4 $devnode

mount $devnode %s""" # noqa

efs_vol_shellcode = """mkdir -p {efs_mountpoint}
MAC=$(curl http://169.254.169.254/latest/meta-data/mac)
export SUBNET_ID=$(curl http://169.254.169.254/latest/meta-data/network/interfaces/macs/$MAC/subnet-id)
NFS_ENDPOINT=$(echo "$AEGEA_EFS_DESC" | jq -r ".[] | select(.SubnetId == env.SUBNET_ID) | .IpAddress")
mount -t nfs -o nfsvers=4.1,rsize=1048576,wsize=1048576,hard,timeo=600,retrans=2 $NFS_ENDPOINT:/ {efs_mountpoint}"""

def batch(args):
    batch_parser.print_help()

batch_parser = register_parser(batch, help="Manage AWS Batch resources", description=__doc__,
                               formatter_class=argparse.RawTextHelpFormatter)

def queues(args):
    table = clients.batch.describe_job_queues()["jobQueues"]
    page_output(tabulate(table, args))

parser = register_listing_parser(queues, parent=batch_parser, help="List Batch queues")

def create_queue(args):
    ces = [dict(computeEnvironment=e, order=i) for i, e in enumerate(args.compute_environments)]
    logger.info("Creating queue %s in %s", args.name, ces)
    queue = clients.batch.create_job_queue(jobQueueName=args.name, priority=args.priority, computeEnvironmentOrder=ces)
    make_waiter(clients.batch.describe_job_queues, "jobQueues[].status", "VALID", "pathAny").wait(jobQueues=[args.name])
    return queue

parser = register_parser(create_queue, parent=batch_parser, help="Create a Batch queue")
parser.add_argument("name")
parser.add_argument("--priority", type=int, default=5)
parser.add_argument("--compute-environments", nargs="+", required=True)<commit_msg>fix batch: newlines and percent characters have special meaning (#32)<commit_after>"""
Manage AWS Batch jobs, queues, and compute environments.
"""

from __future__ import absolute_import, division, print_function, unicode_literals

import os, sys, argparse, base64, collections, io, subprocess, json, time, re
from datetime import datetime

from botocore.exceptions import ClientError
import yaml

from . import logger
from .ls import register_parser, register_listing_parser
from .ecr import ecr_image_name_completer
from .util import Timestamp, paginate
from .util.crypto import ensure_ssh_key
from .util.exceptions import AegeaException
from .util.printing import page_output, tabulate, YELLOW, RED, GREEN, BOLD, ENDC
from .util.aws import (ARN, resources, clients, expect_error_codes, ensure_iam_role, ensure_instance_profile,
                       make_waiter, ensure_vpc, ensure_security_group, ensure_s3_bucket, ensure_log_group,
                       IAMPolicyBuilder, resolve_ami)
from .util.aws.spot import SpotFleetBuilder

bash_cmd_preamble = ["/bin/bash", "-c", 'for i in "$@"; do eval "$i"; done', __name__]

ebs_vol_mgr_shellcode = """iid=$(http http://169.254.169.254/latest/dynamic/instance-identity/document)
aws configure set default.region $(echo "$iid" | jq -r .region)
az=$(echo "$iid" | jq -r .availabilityZone)
vid=$(aws ec2 create-volume --availability-zone $az --size %s --volume-type st1 | jq -r .VolumeId)
aws ec2 create-tags --resource $vid --tags Key=aegea_batch_job,Value=$AWS_BATCH_JOB_ID
trap "umount /mnt || umount -l /mnt; aws ec2 detach-volume --volume-id $vid; let delay=2; while ! aws ec2 describe-volumes --volume-ids $vid | jq -re .Volumes[0].Attachments==[]; do if [[ $delay -gt 30 ]]; then aws ec2 detach-volume --force --volume-id $vid; break; fi; sleep $delay; let delay=$delay*2; done; aws ec2 delete-volume --volume-id $vid" EXIT
while [[ $(aws ec2 describe-volumes --volume-ids $vid | jq -r .Volumes[0].State) != available ]]; do sleep 1; done
# let each process start trying from a different /dev/xvd{letter}
let pid=$$
# when N processes compete, for every success there can be N-1 failures; so the appropriate number of retries is O(N^2)
# let us size this for 10 competitors
# NOTE: the constants 9 and 10 in the $ind and $devnode calculation below are based on strlen("/dev/xvda")
let delay=2+$pid%%5
for try in {1..100}; do if [[ $try == 100 ]]; then echo "Unable to mount $vid on $devnode"; exit 1; fi; if [[ $try -gt 1 ]]; then sleep $delay; fi; devices=$(echo /dev/xvd* /dev/xvd{a..z} /dev/xvd{a..z} | sed 's= =\n=g' | sort | uniq -c | sort -n | grep ' 2 ' | awk '{print $2}'); let devcnt=${#devices}/10+1; let ind=$pid%%devcnt; devnode=${devices:10*$ind:9};  aws ec2 attach-volume --instance-id $(echo "$iid" | jq -r .instanceId) --volume-id $vid --device $devnode || continue; break; done
# attach-volume is not instantaneous, and describe-volume requests are rate-limited
let delay=5+pid%%11
sleep $delay
let try=1
let max_tries=32
while [[ $(aws ec2 describe-volumes --volume-ids $vid | jq -r .Volumes[0].State) != in-use ]]; do if [[ $try == $max_tries ]]; break; fi; let foo=1+$try%%5; let delay=2**$foo+pid%%11; sleep $delay; done
while [[ ! -e $devnode ]]; do sleep 1; done
mkfs.ext4 $devnode
mount $devnode %s""" # noqa

efs_vol_shellcode = """mkdir -p {efs_mountpoint}
MAC=$(curl http://169.254.169.254/latest/meta-data/mac)
export SUBNET_ID=$(curl http://169.254.169.254/latest/meta-data/network/interfaces/macs/$MAC/subnet-id)
NFS_ENDPOINT=$(echo "$AEGEA_EFS_DESC" | jq -r ".[] | select(.SubnetId == env.SUBNET_ID) | .IpAddress")
mount -t nfs -o nfsvers=4.1,rsize=1048576,wsize=1048576,hard,timeo=600,retrans=2 $NFS_ENDPOINT:/ {efs_mountpoint}"""

def batch(args):
    batch_parser.print_help()

batch_parser = register_parser(batch, help="Manage AWS Batch resources", description=__doc__,
                               formatter_class=argparse.RawTextHelpFormatter)

def queues(args):
    table = clients.batch.describe_job_queues()["jobQueues"]
    page_output(tabulate(table, args))

parser = register_listing_parser(queues, parent=batch_parser, help="List Batch queues")

def create_queue(args):
    ces = [dict(computeEnvironment=e, order=i) for i, e in enumerate(args.compute_environments)]
    logger.info("Creating queue %s in %s", args.name, ces)
    queue = clients.batch.create_job_queue(jobQueueName=args.name, priority=args.priority, computeEnvironmentOrder=ces)
    make_waiter(clients.batch.describe_job_queues, "jobQueues[].status", "VALID", "pathAny").wait(jobQueues=[args.name])
    return queue

parser = register_parser(create_queue, parent=batch_parser, help="Create a Batch queue")
parser.add_argument("name")
parser.add_argument("--priority", type=int, default=5)
parser.add_argument("--compute-environments", nargs="+", required=True)
{'fraction_non_alphanumeric': 0.14388264669163545, 'fraction_numerical': 0.02455264253017062, 'mean_word_length': 4.71182412358883, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 6, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2303021', 'n_tokens_mistral': 3405, 'n_tokens_neox': 3226, 'n_words': 932}