<commit_before>"""
Manage AWS Elastic Container Service (ECS) resources, including Fargate tasks.
"""

from __future__ import absolute_import, division, print_function, unicode_literals

import argparse, time, json

from botocore.exceptions import ClientError

from . import logger
from .batch import add_command_args, add_job_defn_args, set_ulimits, set_volumes, get_ecr_image_uri
from .ls import register_parser, register_listing_parser
from .util import Timestamp, paginate
from .util.compat import USING_PYTHON2
from .util.printing import page_output, tabulate, YELLOW, RED, GREEN, BOLD, ENDC
from .util.aws import (ARN, clients, ensure_security_group, ensure_vpc, ensure_iam_role, ensure_log_group,
                       expect_error_codes)
from .util.aws.logs import CloudwatchLogReader
from .util.aws.batch import get_command_and_env

def ecs(args):
    ecs_parser.print_help()

ecs_parser = register_parser(ecs, help="Manage Elastic Container Service resources", description=__doc__)

def clusters(args):
    if not args.clusters:
        args.clusters = list(paginate(clients.ecs.get_paginator("list_clusters")))
    cluster_desc = clients.ecs.describe_clusters(clusters=args.clusters)["clusters"]
    page_output(tabulate(cluster_desc, args))

parser = register_listing_parser(clusters, parent=ecs_parser, help="List ECS clusters")
parser.add_argument("clusters", nargs="*")

def tasks(args):
    list_tasks_args = {}
    if args.cluster:
        list_tasks_args["cluster"] = args.cluster
    if args.launch_type:
        list_tasks_args["launchType"] = args.launch_type
    if args.desired_status:
        list_tasks_args["desiredStatus"] = args.desired_status
    if not args.tasks:
        list_tasks = clients.ecs.get_paginator("list_tasks")
        args.tasks = list(paginate(list_tasks, **list_tasks_args))
        if not args.desired_status:
            args.tasks += list(paginate(list_tasks, desiredStatus="STOPPED", **list_tasks_args))
    task_desc = clients.ecs.describe_tasks(cluster=args.cluster, tasks=args.tasks)["tasks"] if args.tasks else []
    page_output(tabulate(task_desc, args))

parser = register_listing_parser(tasks, parent=ecs_parser, help="List ECS tasks")
parser.add_argument("tasks", nargs="*")
parser.add_argument("--cluster")
parser.add_argument("--desired-status", choices={"RUNNING", "STOPPED"})
parser.add_argument("--launch-type", choices={"EC2", "FARGATE"})

def run(args):
    args.storage = args.efs_storage = None
    command, environment = get_command_and_env(args)
    vpc = ensure_vpc()
    clients.ecs.create_cluster(clusterName=args.cluster)
    log_config = {
        "logDriver": "awslogs",
        "options": {
            "awslogs-region": clients.ecs.meta.region_name,
            "awslogs-group": args.task_name,
            "awslogs-stream-prefix": args.task_name
        }
    }
    ensure_log_group(log_config["options"]["awslogs-group"])

    if args.ecs_image:
        args.image = get_ecr_image_uri(args.ecs_image)

    container_defn = dict(name=args.task_name,
                          image=args.image,
                          memory=args.memory,
                          command=command,
                          environment=environment,
                          logConfiguration=log_config)
    set_volumes(args, container_defn)
    set_ulimits(args, container_defn)
    exec_role = ensure_iam_role(args.execution_role, trust=["ecs-tasks"], policies=["service-role/AWSBatchServiceRole"])
    task_role = ensure_iam_role(args.task_role, trust=["ecs-tasks"])
    clients.ecs.register_task_definition(family=args.task_name,
                                         containerDefinitions=[container_defn],
                                         requiresCompatibilities=["FARGATE"],
                                         executionRoleArn=exec_role.arn,
                                         taskRoleArn=task_role.arn,
                                         networkMode="awsvpc",
                                         cpu=args.fargate_cpu,
                                         memory=args.fargate_memory)
    network_config = {
        'awsvpcConfiguration': {
            'subnets': [
                subnet.id for subnet in vpc.subnets.all()
            ],
            'securityGroups': [ensure_security_group(args.security_group, vpc).id],
            'assignPublicIp': 'ENABLED'
        }
    }
    res = clients.ecs.run_task(cluster=args.cluster,
                               taskDefinition=args.task_name,
                               launchType="FARGATE",
                               networkConfiguration=network_config)
    task_arn = res["tasks"][0]["taskArn"]
    if args.watch:
        watch(watch_parser.parse_args([args.task_name, task_arn, "--cluster", args.cluster]))
    elif args.wait:
        raise NotImplementedError()
    if args.watch or args.wait:
        res = clients.ecs.describe_tasks(cluster=args.cluster, tasks=[task_arn])
        print(json.dumps(res["tasks"][0], indent=2, default=lambda x: str(x)))
        return SystemExit(res["tasks"][0]["containers"][0]["exitCode"])
    else:
        return res["tasks"][0]

register_parser_args = dict(parent=ecs_parser, help="Run a Fargate task")
if not USING_PYTHON2:
    register_parser_args["aliases"] = ["launch"]

parser = register_parser(run, **register_parser_args)
add_command_args(parser)
add_job_defn_args(parser)
parser.add_argument("--execution-role", metavar="IAM_ROLE", default=__name__)
parser.add_argument("--task-role", metavar="IAM_ROLE", default=__name__)
parser.add_argument("--security-group", default=__name__)
parser.add_argument("--cluster", default=__name__.replace(".", "_"))
parser.add_argument("--task-name", default=__name__.replace(".", "_"))
parser.add_argument("--fargate-cpu", help="Execution vCPU count")
parser.add_argument("--fargate-memory")

task_status_colors = dict(PROVISIONING=YELLOW(), PENDING=BOLD() + YELLOW(), ACTIVATING=BOLD() + YELLOW(),
                          RUNNING=GREEN(),
                          DEACTIVATING=BOLD() + GREEN(), STOPPING=BOLD() + GREEN(), DEPROVISIONING=BOLD() + GREEN(),
                          STOPPED=BOLD() + GREEN())

def format_task_status(status):
    return task_status_colors[status] + status + ENDC()

def watch(args):
    _, _, task_id = ARN(args.task_arn).resource.split("/")
    logger.info("Watching task %s (%s)", task_id, args.cluster)
    last_status, events_received = None, 0
    while last_status != "STOPPED":
        res = clients.ecs.describe_tasks(cluster=args.cluster, tasks=[args.task_arn])
        if len(res["tasks"]) == 1:
            task_desc = res["tasks"][0]
            if task_desc["lastStatus"] != last_status:
                logger.info("Task %s %s", args.task_arn, format_task_status(task_desc["lastStatus"]))
                last_status = task_desc["lastStatus"]
        try:
            for event in CloudwatchLogReader("/".join([args.task_name, args.task_name, task_id]),
                                             log_group_name=args.task_name):
                print(str(Timestamp(event["timestamp"])), event["message"])
                events_received += 1
        except ClientError as e:
            expect_error_codes(e, "ResourceNotFoundException")
        if last_status is None and events_received > 0:
            break  # Logs retrieved successfully but task record is no longer in ECS
        time.sleep(1)

watch_parser = register_parser(watch, parent=ecs_parser, help="Monitor a running ECS Fargate task and stream its logs")
watch_parser.add_argument("task_name")
watch_parser.add_argument("task_arn")
watch_parser.add_argument("--cluster", default=__name__.replace(".", "_"))
lines_group = watch_parser.add_mutually_exclusive_group()
lines_group.add_argument("--head", type=int, nargs="?", const=10,
                         help="Retrieve this number of lines from the beginning of the log (default 10)")
lines_group.add_argument("--tail", type=int, nargs="?", const=10,
                         help="Retrieve this number of lines from the end of the log (default 10)")
<commit_msg>aegea ecs run: utilize 4G scratch space<commit_after>"""
Manage AWS Elastic Container Service (ECS) resources, including Fargate tasks.
"""

from __future__ import absolute_import, division, print_function, unicode_literals

import argparse, time, json

from botocore.exceptions import ClientError

from . import logger
from .batch import add_command_args, add_job_defn_args, set_ulimits, set_volumes, get_ecr_image_uri
from .ls import register_parser, register_listing_parser
from .util import Timestamp, paginate
from .util.compat import USING_PYTHON2
from .util.printing import page_output, tabulate, YELLOW, RED, GREEN, BOLD, ENDC
from .util.aws import (ARN, clients, ensure_security_group, ensure_vpc, ensure_iam_role, ensure_log_group,
                       expect_error_codes)
from .util.aws.logs import CloudwatchLogReader
from .util.aws.batch import get_command_and_env

def ecs(args):
    ecs_parser.print_help()

ecs_parser = register_parser(ecs, help="Manage Elastic Container Service resources", description=__doc__)

def clusters(args):
    if not args.clusters:
        args.clusters = list(paginate(clients.ecs.get_paginator("list_clusters")))
    cluster_desc = clients.ecs.describe_clusters(clusters=args.clusters)["clusters"]
    page_output(tabulate(cluster_desc, args))

parser = register_listing_parser(clusters, parent=ecs_parser, help="List ECS clusters")
parser.add_argument("clusters", nargs="*")

def tasks(args):
    list_tasks_args = {}
    if args.cluster:
        list_tasks_args["cluster"] = args.cluster
    if args.launch_type:
        list_tasks_args["launchType"] = args.launch_type
    if args.desired_status:
        list_tasks_args["desiredStatus"] = args.desired_status
    if not args.tasks:
        list_tasks = clients.ecs.get_paginator("list_tasks")
        args.tasks = list(paginate(list_tasks, **list_tasks_args))
        if not args.desired_status:
            args.tasks += list(paginate(list_tasks, desiredStatus="STOPPED", **list_tasks_args))
    task_desc = clients.ecs.describe_tasks(cluster=args.cluster, tasks=args.tasks)["tasks"] if args.tasks else []
    page_output(tabulate(task_desc, args))

parser = register_listing_parser(tasks, parent=ecs_parser, help="List ECS tasks")
parser.add_argument("tasks", nargs="*")
parser.add_argument("--cluster")
parser.add_argument("--desired-status", choices={"RUNNING", "STOPPED"})
parser.add_argument("--launch-type", choices={"EC2", "FARGATE"})

def run(args):
    args.storage = args.efs_storage = None
    command, environment = get_command_and_env(args)
    vpc = ensure_vpc()
    clients.ecs.create_cluster(clusterName=args.cluster)
    log_config = {
        "logDriver": "awslogs",
        "options": {
            "awslogs-region": clients.ecs.meta.region_name,
            "awslogs-group": args.task_name,
            "awslogs-stream-prefix": args.task_name
        }
    }
    ensure_log_group(log_config["options"]["awslogs-group"])

    if args.ecs_image:
        args.image = get_ecr_image_uri(args.ecs_image)

    container_defn = dict(name=args.task_name,
                          image=args.image,
                          memory=args.memory,
                          command=command,
                          environment=environment,
                          logConfiguration=log_config,
                          mountPoints=[dict(sourceVolume="scratch", containerPath="/mnt")])
    set_volumes(args, container_defn)
    set_ulimits(args, container_defn)
    exec_role = ensure_iam_role(args.execution_role, trust=["ecs-tasks"], policies=["service-role/AWSBatchServiceRole"])
    task_role = ensure_iam_role(args.task_role, trust=["ecs-tasks"])
    clients.ecs.register_task_definition(family=args.task_name,
                                         containerDefinitions=[container_defn],
                                         requiresCompatibilities=["FARGATE"],
                                         executionRoleArn=exec_role.arn,
                                         taskRoleArn=task_role.arn,
                                         networkMode="awsvpc",
                                         cpu=args.fargate_cpu,
                                         memory=args.fargate_memory,
                                         volumes=[dict(name="scratch", host={})])
    network_config = {
        'awsvpcConfiguration': {
            'subnets': [
                subnet.id for subnet in vpc.subnets.all()
            ],
            'securityGroups': [ensure_security_group(args.security_group, vpc).id],
            'assignPublicIp': 'ENABLED'
        }
    }
    res = clients.ecs.run_task(cluster=args.cluster,
                               taskDefinition=args.task_name,
                               launchType="FARGATE",
                               networkConfiguration=network_config)
    task_arn = res["tasks"][0]["taskArn"]
    if args.watch:
        watch(watch_parser.parse_args([args.task_name, task_arn, "--cluster", args.cluster]))
    elif args.wait:
        raise NotImplementedError()
    if args.watch or args.wait:
        res = clients.ecs.describe_tasks(cluster=args.cluster, tasks=[task_arn])
        print(json.dumps(res["tasks"][0], indent=2, default=lambda x: str(x)))
        return SystemExit(res["tasks"][0]["containers"][0]["exitCode"])
    else:
        return res["tasks"][0]

register_parser_args = dict(parent=ecs_parser, help="Run a Fargate task")
if not USING_PYTHON2:
    register_parser_args["aliases"] = ["launch"]

parser = register_parser(run, **register_parser_args)
add_command_args(parser)
add_job_defn_args(parser)
parser.add_argument("--execution-role", metavar="IAM_ROLE", default=__name__)
parser.add_argument("--task-role", metavar="IAM_ROLE", default=__name__)
parser.add_argument("--security-group", default=__name__)
parser.add_argument("--cluster", default=__name__.replace(".", "_"))
parser.add_argument("--task-name", default=__name__.replace(".", "_"))
parser.add_argument("--fargate-cpu", help="Execution vCPU count")
parser.add_argument("--fargate-memory")

task_status_colors = dict(PROVISIONING=YELLOW(), PENDING=BOLD() + YELLOW(), ACTIVATING=BOLD() + YELLOW(),
                          RUNNING=GREEN(),
                          DEACTIVATING=BOLD() + GREEN(), STOPPING=BOLD() + GREEN(), DEPROVISIONING=BOLD() + GREEN(),
                          STOPPED=BOLD() + GREEN())

def format_task_status(status):
    return task_status_colors[status] + status + ENDC()

def watch(args):
    _, _, task_id = ARN(args.task_arn).resource.split("/")
    logger.info("Watching task %s (%s)", task_id, args.cluster)
    last_status, events_received = None, 0
    while last_status != "STOPPED":
        res = clients.ecs.describe_tasks(cluster=args.cluster, tasks=[args.task_arn])
        if len(res["tasks"]) == 1:
            task_desc = res["tasks"][0]
            if task_desc["lastStatus"] != last_status:
                logger.info("Task %s %s", args.task_arn, format_task_status(task_desc["lastStatus"]))
                last_status = task_desc["lastStatus"]
        try:
            for event in CloudwatchLogReader("/".join([args.task_name, args.task_name, task_id]),
                                             log_group_name=args.task_name):
                print(str(Timestamp(event["timestamp"])), event["message"])
                events_received += 1
        except ClientError as e:
            expect_error_codes(e, "ResourceNotFoundException")
        if last_status is None and events_received > 0:
            break  # Logs retrieved successfully but task record is no longer in ECS
        time.sleep(1)

watch_parser = register_parser(watch, parent=ecs_parser, help="Monitor a running ECS Fargate task and stream its logs")
watch_parser.add_argument("task_name")
watch_parser.add_argument("task_arn")
watch_parser.add_argument("--cluster", default=__name__.replace(".", "_"))
lines_group = watch_parser.add_mutually_exclusive_group()
lines_group.add_argument("--head", type=int, nargs="?", const=10,
                         help="Retrieve this number of lines from the beginning of the log (default 10)")
lines_group.add_argument("--tail", type=int, nargs="?", const=10,
                         help="Retrieve this number of lines from the end of the log (default 10)")
<|endoftext|>


<commit_before>"""
Manage AWS Elastic Container Service (ECS) resources, including Fargate tasks.
"""

from __future__ import absolute_import, division, print_function, unicode_literals

import argparse, time, json, hashlib
from itertools import product
from functools import partial

from botocore.exceptions import ClientError

from . import logger
from .batch import add_command_args, add_job_defn_args
from .ls import register_parser, register_listing_parser
from .util import Timestamp, paginate, ThreadPoolExecutor
from .util.compat import USING_PYTHON2
from .util.printing import page_output, tabulate, YELLOW, RED, GREEN, BOLD, ENDC
from .util.aws import (ARN, clients, ensure_security_group, ensure_vpc, ensure_iam_role, ensure_log_group,
                       ensure_ecs_cluster, expect_error_codes)
from .util.aws.logs import CloudwatchLogReader
from .util.aws.batch import get_command_and_env, set_ulimits, set_volumes, get_ecr_image_uri

def complete_cluster_name(**kwargs):
    return [ARN(c).resource.partition("/")[2] for c in paginate(clients.ecs.get_paginator("list_clusters"))]

def ecs(args):
    ecs_parser.print_help()

ecs_parser = register_parser(ecs, help="Manage Elastic Container Service resources", description=__doc__)

def clusters(args):
    if not args.clusters:
        args.clusters = list(paginate(clients.ecs.get_paginator("list_clusters")))
    cluster_desc = clients.ecs.describe_clusters(clusters=args.clusters)["clusters"]
    page_output(tabulate(cluster_desc, args))

parser = register_listing_parser(clusters, parent=ecs_parser, help="List ECS clusters")
parser.add_argument("clusters", nargs="*").completer = complete_cluster_name

def tasks(args):
    list_clusters = clients.ecs.get_paginator("list_clusters")
    list_tasks = clients.ecs.get_paginator("list_tasks")

    def list_tasks_worker(worker_args):
        cluster, status = worker_args
        return cluster, status, list(paginate(list_tasks, cluster=cluster, desiredStatus=status))

    def describe_tasks_worker(t, cluster=None):
        return clients.ecs.describe_tasks(cluster=cluster, tasks=t)["tasks"] if t else []

    task_descs = []
    if args.clusters is None:
        args.clusters = [__name__.replace(".", "_")] if args.tasks else list(paginate(list_clusters))
    if args.tasks:
        task_descs = describe_tasks_worker(args.tasks, cluster=args.clusters[0])
    else:
        with ThreadPoolExecutor() as executor:
            for cluster, status, tasks in executor.map(list_tasks_worker, product(args.clusters, args.desired_status)):
                worker = partial(describe_tasks_worker, cluster=cluster)
                descs = executor.map(worker, (tasks[pos:pos + 100] for pos in range(0, len(tasks), 100)))
                task_descs += sum(descs, [])
    page_output(tabulate(task_descs, args))

parser = register_listing_parser(tasks, parent=ecs_parser, help="List ECS tasks")
parser.add_argument("tasks", nargs="*")
parser.add_argument("--clusters", nargs="*").completer = complete_cluster_name
parser.add_argument("--desired-status", nargs=1, choices={"RUNNING", "STOPPED"}, default=["RUNNING", "STOPPED"])
parser.add_argument("--launch-type", nargs=1, choices={"EC2", "FARGATE"}, default=["EC2", "FARGATE"])

def run(args):
    args.storage = args.efs_storage = args.mount_instance_storage = None
    command, environment = get_command_and_env(args)
    vpc = ensure_vpc()
    ensure_ecs_cluster(args.cluster)
    log_config = {
        "logDriver": "awslogs",
        "options": {
            "awslogs-region": clients.ecs.meta.region_name,
            "awslogs-group": args.task_name,
            "awslogs-stream-prefix": args.task_name
        }
    }
    ensure_log_group(log_config["options"]["awslogs-group"])

    if args.ecs_image:
        args.image = get_ecr_image_uri(args.ecs_image)

    container_defn = dict(name=args.task_name,
                          image=args.image,
                          cpu=0,
                          memory=args.memory,
                          command=[],
                          environment=[],
                          portMappings=[],
                          essential=True,
                          logConfiguration=log_config,
                          mountPoints=[dict(sourceVolume="scratch", containerPath="/mnt")],
                          volumesFrom=[])
    set_volumes(args, container_defn)
    set_ulimits(args, container_defn)
    exec_role = ensure_iam_role(args.execution_role, trust=["ecs-tasks"],
                                policies=["service-role/AmazonEC2ContainerServiceforEC2Role",
                                          "service-role/AWSBatchServiceRole"])
    task_role = ensure_iam_role(args.task_role, trust=["ecs-tasks"])

    expect_task_defn = dict(containerDefinitions=[container_defn],
                            requiresCompatibilities=["FARGATE"],
                            taskRoleArn=task_role.arn,
                            executionRoleArn=exec_role.arn,
                            networkMode="awsvpc",
                            cpu=args.fargate_cpu,
                            memory=args.fargate_memory,
                            volumes=[dict(name="scratch", host={})])

    task_hash = hashlib.sha256(json.dumps(expect_task_defn, sort_keys=True).encode()).hexdigest()[:8]
    task_defn_name = __name__.replace(".", "_") + "_" + task_hash

    try:
        task_defn = clients.ecs.describe_task_definition(taskDefinition=task_defn_name)["taskDefinition"]
        assert task_defn["status"] == "ACTIVE"
        assert "FARGATE" in task_defn["compatibilities"]
        desc_keys = ["family", "revision", "taskDefinitionArn", "status", "compatibilities", "placementConstraints",
                     "requiresAttributes"]
        task_desc = {key: task_defn.pop(key) for key in desc_keys}
        if expect_task_defn["cpu"].endswith(" vCPU"):
            expect_task_defn["cpu"] = str(int(expect_task_defn["cpu"][:-len(" vCPU")]) * 1024)
        if expect_task_defn["memory"].endswith(" GB"):
            expect_task_defn["memory"] = str(int(expect_task_defn["memory"][:-len(" GB")]) * 1024)
        assert task_defn == expect_task_defn
        logger.debug("Reusing task definition %s", task_desc["taskDefinitionArn"])
    except (ClientError, AssertionError):
        logger.debug("Registering new ECS task definition %s", task_defn_name)
        task_desc = clients.ecs.register_task_definition(family=task_defn_name, **expect_task_defn)["taskDefinition"]

    network_config = {
        'awsvpcConfiguration': {
            'subnets': [
                subnet.id for subnet in vpc.subnets.all()
            ],
            'securityGroups': [ensure_security_group(args.security_group, vpc).id],
            'assignPublicIp': 'ENABLED'
        }
    }
    container_overrides = [dict(name=args.task_name, command=command, environment=environment)]
    res = clients.ecs.run_task(cluster=args.cluster,
                               taskDefinition=task_desc["taskDefinitionArn"],
                               launchType="FARGATE",
                               networkConfiguration=network_config,
                               overrides=dict(containerOverrides=container_overrides))
    task_arn = res["tasks"][0]["taskArn"]
    if args.watch:
        watch(watch_parser.parse_args([task_arn, "--task-name", args.task_name]))
    elif args.wait:
        raise NotImplementedError()
    if args.watch or args.wait:
        res = clients.ecs.describe_tasks(cluster=args.cluster, tasks=[task_arn])
        print(json.dumps(res["tasks"][0], indent=2, default=lambda x: str(x)))
        return SystemExit(res["tasks"][0]["containers"][0]["exitCode"])
    else:
        return res["tasks"][0]

register_parser_args = dict(parent=ecs_parser, help="Run a Fargate task")
if not USING_PYTHON2:
    register_parser_args["aliases"] = ["launch"]

parser = register_parser(run, **register_parser_args)
add_command_args(parser)
add_job_defn_args(parser)
parser.add_argument("--execution-role", metavar="IAM_ROLE", default=__name__)
parser.add_argument("--task-role", metavar="IAM_ROLE", default=__name__)
parser.add_argument("--security-group", default=__name__)
parser.add_argument("--cluster", default=__name__.replace(".", "_"))
parser.add_argument("--task-name", default=__name__.replace(".", "_"))
parser.add_argument("--fargate-cpu", help="Execution vCPU count")
parser.add_argument("--fargate-memory")

task_status_colors = dict(PROVISIONING=YELLOW(), PENDING=BOLD() + YELLOW(), ACTIVATING=BOLD() + YELLOW(),
                          RUNNING=GREEN(),
                          DEACTIVATING=BOLD() + GREEN(), STOPPING=BOLD() + GREEN(), DEPROVISIONING=BOLD() + GREEN(),
                          STOPPED=BOLD() + GREEN())

def format_task_status(status):
    return task_status_colors[status] + status + ENDC()

def watch(args):
    _, cluster, task_id = ARN(args.task_arn).resource.split("/")
    logger.info("Watching task %s (%s)", task_id, cluster)
    last_status, events_received = None, 0
    log_reader = CloudwatchLogReader("/".join([args.task_name, args.task_name, task_id]),
                                     log_group_name=args.task_name)
    while last_status != "STOPPED":
        res = clients.ecs.describe_tasks(cluster=cluster, tasks=[args.task_arn])
        if len(res["tasks"]) == 1:
            task_desc = res["tasks"][0]
            if task_desc["lastStatus"] != last_status:
                logger.info("Task %s %s", args.task_arn, format_task_status(task_desc["lastStatus"]))
                last_status = task_desc["lastStatus"]
        try:
            for event in log_reader:
                print(str(Timestamp(event["timestamp"])), event["message"])
                events_received += 1
        except ClientError as e:
            expect_error_codes(e, "ResourceNotFoundException")
        if last_status is None and events_received > 0:
            break  # Logs retrieved successfully but task record is no longer in ECS
        time.sleep(1)

watch_parser = register_parser(watch, parent=ecs_parser, help="Monitor a running ECS Fargate task and stream its logs")
watch_parser.add_argument("task_arn")
watch_parser.add_argument("--task-name", default=__name__.replace(".", "_"))
lines_group = watch_parser.add_mutually_exclusive_group()
lines_group.add_argument("--head", type=int, nargs="?", const=10,
                         help="Retrieve this number of lines from the beginning of the log (default 10)")
lines_group.add_argument("--tail", type=int, nargs="?", const=10,
                         help="Retrieve this number of lines from the end of the log (default 10)")
<commit_msg>ecs watch: make compatible with new ECS task IDs<commit_after>"""
Manage AWS Elastic Container Service (ECS) resources, including Fargate tasks.
"""

from __future__ import absolute_import, division, print_function, unicode_literals

import os, argparse, time, json, hashlib
from itertools import product
from functools import partial

from botocore.exceptions import ClientError

from . import logger
from .batch import add_command_args, add_job_defn_args, print_event
from .ls import register_parser, register_listing_parser
from .util import Timestamp, paginate, ThreadPoolExecutor
from .util.compat import USING_PYTHON2
from .util.printing import page_output, tabulate, YELLOW, RED, GREEN, BOLD, ENDC
from .util.aws import (ARN, clients, ensure_security_group, ensure_vpc, ensure_iam_role, ensure_log_group,
                       ensure_ecs_cluster, expect_error_codes)
from .util.aws.logs import CloudwatchLogReader
from .util.aws.batch import get_command_and_env, set_ulimits, set_volumes, get_ecr_image_uri

def complete_cluster_name(**kwargs):
    return [ARN(c).resource.partition("/")[2] for c in paginate(clients.ecs.get_paginator("list_clusters"))]

def ecs(args):
    ecs_parser.print_help()

ecs_parser = register_parser(ecs, help="Manage Elastic Container Service resources", description=__doc__)

def clusters(args):
    if not args.clusters:
        args.clusters = list(paginate(clients.ecs.get_paginator("list_clusters")))
    cluster_desc = clients.ecs.describe_clusters(clusters=args.clusters)["clusters"]
    page_output(tabulate(cluster_desc, args))

parser = register_listing_parser(clusters, parent=ecs_parser, help="List ECS clusters")
parser.add_argument("clusters", nargs="*").completer = complete_cluster_name

def tasks(args):
    list_clusters = clients.ecs.get_paginator("list_clusters")
    list_tasks = clients.ecs.get_paginator("list_tasks")

    def list_tasks_worker(worker_args):
        cluster, status = worker_args
        return cluster, status, list(paginate(list_tasks, cluster=cluster, desiredStatus=status))

    def describe_tasks_worker(t, cluster=None):
        return clients.ecs.describe_tasks(cluster=cluster, tasks=t)["tasks"] if t else []

    task_descs = []
    if args.clusters is None:
        args.clusters = [__name__.replace(".", "_")] if args.tasks else list(paginate(list_clusters))
    if args.tasks:
        task_descs = describe_tasks_worker(args.tasks, cluster=args.clusters[0])
    else:
        with ThreadPoolExecutor() as executor:
            for cluster, status, tasks in executor.map(list_tasks_worker, product(args.clusters, args.desired_status)):
                worker = partial(describe_tasks_worker, cluster=cluster)
                descs = executor.map(worker, (tasks[pos:pos + 100] for pos in range(0, len(tasks), 100)))
                task_descs += sum(descs, [])
    page_output(tabulate(task_descs, args))

parser = register_listing_parser(tasks, parent=ecs_parser, help="List ECS tasks")
parser.add_argument("tasks", nargs="*")
parser.add_argument("--clusters", nargs="*").completer = complete_cluster_name
parser.add_argument("--desired-status", nargs=1, choices={"RUNNING", "STOPPED"}, default=["RUNNING", "STOPPED"])
parser.add_argument("--launch-type", nargs=1, choices={"EC2", "FARGATE"}, default=["EC2", "FARGATE"])

def run(args):
    args.storage = args.efs_storage = args.mount_instance_storage = None
    command, environment = get_command_and_env(args)
    vpc = ensure_vpc()
    ensure_ecs_cluster(args.cluster)
    log_config = {
        "logDriver": "awslogs",
        "options": {
            "awslogs-region": clients.ecs.meta.region_name,
            "awslogs-group": args.task_name,
            "awslogs-stream-prefix": args.task_name
        }
    }
    ensure_log_group(log_config["options"]["awslogs-group"])

    if args.ecs_image:
        args.image = get_ecr_image_uri(args.ecs_image)

    container_defn = dict(name=args.task_name,
                          image=args.image,
                          cpu=0,
                          memory=args.memory,
                          command=[],
                          environment=[],
                          portMappings=[],
                          essential=True,
                          logConfiguration=log_config,
                          mountPoints=[dict(sourceVolume="scratch", containerPath="/mnt")],
                          volumesFrom=[])
    set_volumes(args, container_defn)
    set_ulimits(args, container_defn)
    exec_role = ensure_iam_role(args.execution_role, trust=["ecs-tasks"],
                                policies=["service-role/AmazonEC2ContainerServiceforEC2Role",
                                          "service-role/AWSBatchServiceRole"])
    task_role = ensure_iam_role(args.task_role, trust=["ecs-tasks"])

    expect_task_defn = dict(containerDefinitions=[container_defn],
                            requiresCompatibilities=["FARGATE"],
                            taskRoleArn=task_role.arn,
                            executionRoleArn=exec_role.arn,
                            networkMode="awsvpc",
                            cpu=args.fargate_cpu,
                            memory=args.fargate_memory,
                            volumes=[dict(name="scratch", host={})])

    task_hash = hashlib.sha256(json.dumps(expect_task_defn, sort_keys=True).encode()).hexdigest()[:8]
    task_defn_name = __name__.replace(".", "_") + "_" + task_hash

    try:
        task_defn = clients.ecs.describe_task_definition(taskDefinition=task_defn_name)["taskDefinition"]
        assert task_defn["status"] == "ACTIVE"
        assert "FARGATE" in task_defn["compatibilities"]
        desc_keys = ["family", "revision", "taskDefinitionArn", "status", "compatibilities", "placementConstraints",
                     "requiresAttributes"]
        task_desc = {key: task_defn.pop(key) for key in desc_keys}
        if expect_task_defn["cpu"].endswith(" vCPU"):
            expect_task_defn["cpu"] = str(int(expect_task_defn["cpu"][:-len(" vCPU")]) * 1024)
        if expect_task_defn["memory"].endswith(" GB"):
            expect_task_defn["memory"] = str(int(expect_task_defn["memory"][:-len(" GB")]) * 1024)
        assert task_defn == expect_task_defn
        logger.debug("Reusing task definition %s", task_desc["taskDefinitionArn"])
    except (ClientError, AssertionError):
        logger.debug("Registering new ECS task definition %s", task_defn_name)
        task_desc = clients.ecs.register_task_definition(family=task_defn_name, **expect_task_defn)["taskDefinition"]

    network_config = {
        'awsvpcConfiguration': {
            'subnets': [
                subnet.id for subnet in vpc.subnets.all()
            ],
            'securityGroups': [ensure_security_group(args.security_group, vpc).id],
            'assignPublicIp': 'ENABLED'
        }
    }
    container_overrides = [dict(name=args.task_name, command=command, environment=environment)]
    res = clients.ecs.run_task(cluster=args.cluster,
                               taskDefinition=task_desc["taskDefinitionArn"],
                               launchType="FARGATE",
                               networkConfiguration=network_config,
                               overrides=dict(containerOverrides=container_overrides))
    task_arn = res["tasks"][0]["taskArn"]
    if args.watch:
        watch(watch_parser.parse_args([task_arn, "--cluster", args.cluster, "--task-name", args.task_name]))
    elif args.wait:
        raise NotImplementedError()
    if args.watch or args.wait:
        res = clients.ecs.describe_tasks(cluster=args.cluster, tasks=[task_arn])
        print(json.dumps(res["tasks"][0], indent=2, default=lambda x: str(x)))
        return SystemExit(res["tasks"][0]["containers"][0]["exitCode"])
    else:
        return res["tasks"][0]

register_parser_args = dict(parent=ecs_parser, help="Run a Fargate task")
if not USING_PYTHON2:
    register_parser_args["aliases"] = ["launch"]

parser = register_parser(run, **register_parser_args)
add_command_args(parser)
add_job_defn_args(parser)
parser.add_argument("--execution-role", metavar="IAM_ROLE", default=__name__)
parser.add_argument("--task-role", metavar="IAM_ROLE", default=__name__)
parser.add_argument("--security-group", default=__name__)
parser.add_argument("--cluster", default=__name__.replace(".", "_"))
parser.add_argument("--task-name", default=__name__.replace(".", "_"))
parser.add_argument("--fargate-cpu", help="Execution vCPU count")
parser.add_argument("--fargate-memory")

task_status_colors = dict(PROVISIONING=YELLOW(), PENDING=BOLD() + YELLOW(), ACTIVATING=BOLD() + YELLOW(),
                          RUNNING=GREEN(),
                          DEACTIVATING=BOLD() + GREEN(), STOPPING=BOLD() + GREEN(), DEPROVISIONING=BOLD() + GREEN(),
                          STOPPED=BOLD() + GREEN())

def format_task_status(status):
    return task_status_colors[status] + status + ENDC()

def watch(args):
    logger.info("Watching task %s (%s)", args.task_id, args.cluster)
    last_status, events_received = None, 0
    log_reader = CloudwatchLogReader("/".join([args.task_name, args.task_name, os.path.basename(args.task_id)]),
                                     log_group_name=args.task_name)
    while last_status != "STOPPED":
        res = clients.ecs.describe_tasks(cluster=args.cluster, tasks=[args.task_id])
        if len(res["tasks"]) == 1:
            task_desc = res["tasks"][0]
            if task_desc["lastStatus"] != last_status:
                logger.info("Task %s %s", args.task_id, format_task_status(task_desc["lastStatus"]))
                last_status = task_desc["lastStatus"]
        try:
            for event in log_reader:
                print_event(event)
                events_received += 1
        except ClientError as e:
            expect_error_codes(e, "ResourceNotFoundException")
        if last_status is None and events_received > 0:
            break  # Logs retrieved successfully but task record is no longer in ECS
        time.sleep(1)

watch_parser = register_parser(watch, parent=ecs_parser, help="Monitor a running ECS Fargate task and stream its logs")
watch_parser.add_argument("task_id")
watch_parser.add_argument("--cluster", default=__name__.replace(".", "_"))
watch_parser.add_argument("--task-name", default=__name__.replace(".", "_"))
lines_group = watch_parser.add_mutually_exclusive_group()
lines_group.add_argument("--head", type=int, nargs="?", const=10,
                         help="Retrieve this number of lines from the beginning of the log (default 10)")
lines_group.add_argument("--tail", type=int, nargs="?", const=10,
                         help="Retrieve this number of lines from the end of the log (default 10)")
<|endoftext|>


<commit_before>def print_matrix(m):
    print ' '
    for line in m:
        spTupel = ()
        breite = len(line)
        for column in line:
            spTupel = spTupel + (column, )
        print "%3i"*breite % spTupel

#calculate the affine gap distance between 2 strings default weights
#taken from page 28 of Bilenko's Ph. D dissertation: Learnable
#Similarity Functions and their Application to Record Linkage and
#Clustering

def affineGapDistance(string1, string2,
                      matchWeight = -5,
                      mismatchWeight = 5,
                      gapWeight = 5,
                      spaceWeight = 1):

  length1 = len(string1)
  length2 = len(string2)

  #set up recurrence matrices

  #v_matrix = minimum distance matrix
  v_matrix = [[None for _ in xrange(length1+1)] for _ in xrange(length2+1)]

  # define base case of recurrences
  # Matrix(0,0) = 0
  # Matrix(0,j) = gapWeight + spaceWeight * j
  v_matrix[0] = f_matrix = [(j * spaceWeight + gapWeight) for j in xrange(length1 + 1)]
  v_matrix[0][0] = 0

  #e_matrix: minimum distance matrix when string1 prefix is left aligned to
  #string2
  #e_matrix = [row[:] for row in v_matrix]

  #f_matrix: minimum distance matrix when string1 prefix is right
  #aligned to string2
  #f_matrix = [row[:] for row in v_matrix]

  for i in xrange(1,length2 + 1) :
    # Matrix(i,0) = gapWeight + spaceWeight * i
    v_matrix[i][0] = e = i * spaceWeight + gapWeight

    for j in xrange(1,length1 + 1) :
      if string2[i-1] == string1[j-1] :
        g = v_matrix[i-1][j-1] + matchWeight
      else :
        g = v_matrix[i-1][j-1] + mismatchWeight

      e = (e + spaceWeight
           if e > v_matrix[i][j-1] + gapWeight
           else v_matrix[i][j-1] + gapWeight + spaceWeight)

      #print "g: ", g_matrix[row][col]
      #e_matrix[i][j] = (min(e_matrix[i][j-1],
      #                      v_matrix[i][j-1] + gapWeight)
      #                  + spaceWeight)
      #print "e: ", e_matrix[row][col]

      f_matrix[j] = (f_matrix[j] + spaceWeight
                     if f_matrix[j] > v_matrix[i-1][j] + gapWeight
                     else v_matrix[i-1][j] + gapWeight + spaceWeight)

      #print "f: ", f_matrix[row][col]
      v_matrix[i][j] = min(e,
                           f_matrix[j],
                           g)
      #print "v: ", v_matrix[row][col]

  #print "Affine Gap v_matrix:"
  #print_matrix(g_matrix)
  #print_matrix(e_matrix)
  #print_matrix(f_matrix)
  #print_matrix(v_matrix)

  return v_matrix[length2][length1]/float(length1 + length2)

if __name__ == "__main__" :
    import cProfile
    def test() :
        for i in xrange(100000) :
            string1 = 'asdf'
            string2 = 'fdsa'
            distance = affineGapDistance(string1, string2)

    cProfile.run('test()')
<commit_msg>more bugfixes and moved last min<commit_after>def print_matrix(m):
    print ' '
    for line in m:
        spTupel = ()
        breite = len(line)
        for column in line:
            spTupel = spTupel + (column, )
        print "%3i"*breite % spTupel

#calculate the affine gap distance between 2 strings default weights
#taken from page 28 of Bilenko's Ph. D dissertation: Learnable
#Similarity Functions and their Application to Record Linkage and
#Clustering

def affineGapDistance(string1, string2,
                      matchWeight = -5,
                      mismatchWeight = 5,
                      gapWeight = 5,
                      spaceWeight = 1):

  length1 = len(string1)
  length2 = len(string2)

  #set up recurrence matrices

  #V_matrix = minimum distance matrix
  v_matrix = [[None for _ in xrange(length1+1)] for _ in xrange(length2+1)]

  # define base case of recurrences
  # V(0,0) = F(0,0) = 0
  # V(0,j) = F(0,j) = gapWeight + spaceWeight * j

  v_matrix[0] = f = [(j * spaceWeight + gapWeight)
                     for j in xrange(length1 + 1)]
  v_matrix[0][0] = 0

  for i in xrange(1,length2 + 1) :
    # Base conditions
    # V(i,0) = E(i,0) = gapWeight + spaceWeight * i
    v_matrix[i][0] = e = i * spaceWeight + gapWeight

    for j in xrange(1,length1 + 1) :

      # E: minimum distance matrix when string1 prefix is left aligned
      # to string2
      #
      # E(i,j) = min(E(i,j-1), V(i,j-1) + gapWeight) + spaceWeight
      e = (e + spaceWeight
           if e < v_matrix[i][j-1] + gapWeight
           else v_matrix[i][j-1] + gapWeight + spaceWeight)

      # F: minimum distance matrix when string1 prefix is right
      # aligned to string2
      #
      # F(i,j) = min(F(i-1,j), V(i-1,j) + gapWeight) + spaceWeight
      f[j] = (f[j] + spaceWeight
              if f[j] < v_matrix[i-1][j] + gapWeight
              else v_matrix[i-1][j] + gapWeight + spaceWeight)

      # G: minimum distance matrix when string1 prefix is aligned to
      # string2
      #
      # G(i,j) = V(i-1,j-1) + (matchWeight | misMatchWeight)
      if string2[i-1] == string1[j-1] :
        g = v_matrix[i-1][j-1] + matchWeight
      else :
        g = v_matrix[i-1][j-1] + mismatchWeight

      # V(i,j) = min(E(i,j), F(i,j), G(i,j))
      v_matrix[i][j] = (e
                        if e < f[j]
                        else f[j]
                        if f[j] < g
                        else g)

  #print "Affine Gap v_matrix:"
  #print_matrix(g_matrix)
  #print_matrix(e_matrix)
  #print_matrix(f_matrix)
  #print_matrix(v_matrix)

  return v_matrix[length2][length1]/float(length1 + length2)

if __name__ == "__main__" :
    import cProfile
    def test() :
        for i in xrange(100000) :
            string1 = 'asdf'
            string2 = 'fdsa'
            distance = affineGapDistance(string1, string2)

    cProfile.run('test()')
<|endoftext|>


<commit_before>"""
The ``agar.auth`` module contains classes, functions, and decorators to help secure a `webapp2.Requesthandler`_.
"""
from functools import wraps

from agar.config import Config

class AuthConfig(Config):
    """
    :py:class:`~agar.config.Config` settings for the ``agar.auth`` library.
    Settings are under the ``agar_auth`` namespace.

    The following settings (and defaults) are provided::

        agar_auth_AUTHENTICATION_PROPERTY = 'user'
        def agar_auth_authenticate(request):
            return None

    To override ``agar.auth`` settings, define values in the ``appengine_config.py`` file in the root of your project.
    """
    _prefix = 'agar_auth'

    #: The property name under which to place the authentication object on the request.
    AUTHENTICATION_PROPERTY = 'user'

    def authenticate(request):
        """
        The authenticate function. It takes a single `webapp2.Request`_ argument, and returns a non-``None`` value if
        the request can be authenticated. If the request can not be authenticated, the function should return ``None``.
        The type of the returned value can be anything, but it should be a type that your `webapp2.RequestHandler`_ expects.
        The default implementation always returns ``None``.

        :param request: The `webapp2.Request`_ object to authenticate.
        :return: A non-``None`` value if the request can be authenticated. If the request can not be authenticated, the
            function should return ``None``.
        """
        return None

#: The configuration object for ``agar.auth`` settings.
config = AuthConfig.get_config()

def authentication_required(authenticate=None, require_https=False):
    """
    A decorator to authenticate a `RequestHandler <http://webapp-improved.appspot.com/api.html#webapp2.RequestHandler>`_.
    If the authenticate function returns a non-``None`` value, it will assign it to the request ``user`` attribute
    (or any re-configured name), that is passed to the decorated handler. If the authenticate function returns ``None``,
    it will call the `webapp2.RequestHandler.abort`_ method with a status of ``403``.

    :param authenticate: The authenticate function to use to authenticate a request. The function should take a single
        `webapp2.Request`_ argument, and return a non-``None`` value if the request can be authenticated. If the request
        can not be authenticated, the function should return ``None``. The type of the returned value can be anything,
        but it should be a type that your `webapp2.RequestHandler`_ expects.
        If ``None``, the config function :py:meth:`~agar.auth.AuthConfig.authenticate` will be used.
    :param require_https: If ``True``, this will enforce that a request was made via HTTPS, otherwise a ``403`` response
        will be returned.
    """
    if authenticate is None:
        authenticate = config.authenticate
    def decorator(request_method):
        @wraps(request_method)
        def wrapped(self, *args, **kwargs):
            if require_https:
                import urlparse
                from agar.env import on_server
                scheme, netloc, path, query, fragment = urlparse.urlsplit(self.request.url)
                if on_server and scheme and scheme.lower() != 'https':
                    self.abort(403)
            authentication = authenticate(self.request)
            if authentication is not None:
                setattr(self.request, config.AUTHENTICATION_PROPERTY, authentication)
                request_method(self, *args, **kwargs)
            else:
                self.abort(403)
        return wrapped
    return decorator

def https_authentication_required(authenticate=None):
    """
    A decorator to authenticate a secure request to a `webapp2.RequestHandler`_.
    This decorator uses the :py:func:`~agar.auth.https_authenticate` authenticate function.
    """
    return authentication_required(authenticate=authenticate, require_https=True)
<commit_msg>Update docs<commit_after>"""
The ``agar.auth`` module contains classes, functions, and decorators to help secure a `webapp2.Requesthandler`_.
"""
from functools import wraps

from agar.config import Config

class AuthConfig(Config):
    """
    :py:class:`~agar.config.Config` settings for the ``agar.auth`` library.
    Settings are under the ``agar_auth`` namespace.

    The following settings (and defaults) are provided::

        agar_auth_AUTHENTICATION_PROPERTY = 'user'
        def agar_auth_authenticate(request):
            return None

    To override ``agar.auth`` settings, define values in the ``appengine_config.py`` file in the root of your project.
    """
    _prefix = 'agar_auth'

    #: The property name under which to place the authentication object on the request.
    AUTHENTICATION_PROPERTY = 'user'

    def authenticate(request):
        """
        The authenticate function. It takes a single `webapp2.Request`_ argument, and returns a non-``None`` value if
        the request can be authenticated. If the request can not be authenticated, the function should return ``None``.
        The type of the returned value can be anything, but it should be a type that your `webapp2.RequestHandler`_ expects.
        The default implementation always returns ``None``.

        :param request: The `webapp2.Request`_ object to authenticate.
        :return: A non-``None`` value if the request can be authenticated. If the request can not be authenticated, the
            function should return ``None``.
        """
        return None

#: The configuration object for ``agar.auth`` settings.
config = AuthConfig.get_config()

def authentication_required(authenticate=None, require_https=False):
    """
    A decorator to authenticate a `RequestHandler <http://webapp-improved.appspot.com/api.html#webapp2.RequestHandler>`_.
    If the authenticate function returns a non-``None`` value, it will assign it to the request ``user`` attribute
    (or any re-configured name), that is passed to the decorated handler. If the authenticate function returns ``None``,
    it will call the `webapp2.RequestHandler.abort`_ method with a status of ``403``.

    :param authenticate: The authenticate function to use to authenticate a request. The function should take a single
        `webapp2.Request`_ argument, and return a non-``None`` value if the request can be authenticated. If the request
        can not be authenticated, the function should return ``None``. The type of the returned value can be anything,
        but it should be a type that your `webapp2.RequestHandler`_ expects.
        If ``None``, the config function :py:meth:`~agar.auth.AuthConfig.authenticate` will be used.
    :param require_https: If ``True``, this will enforce that a request was made via HTTPS, otherwise a ``403`` response
        will be returned.
    """
    if authenticate is None:
        authenticate = config.authenticate
    def decorator(request_method):
        @wraps(request_method)
        def wrapped(self, *args, **kwargs):
            if require_https:
                import urlparse
                from agar.env import on_server
                scheme, netloc, path, query, fragment = urlparse.urlsplit(self.request.url)
                if on_server and scheme and scheme.lower() != 'https':
                    self.abort(403)
            authentication = authenticate(self.request)
            if authentication is not None:
                setattr(self.request, config.AUTHENTICATION_PROPERTY, authentication)
                request_method(self, *args, **kwargs)
            else:
                self.abort(403)
        return wrapped
    return decorator

def https_authentication_required(authenticate=None):
    """
    A decorator to authenticate a secure request to a `webapp2.RequestHandler`_.
    """
    return authentication_required(authenticate=authenticate, require_https=True)
<|endoftext|>


<commit_before># coding: utf-8

import numpy as np
import torch
from torch import nn
from torch import optim
from torch import cuda
from torch.autograd import Variable

from .core import Agent
from .replay_buffer import ReplayBuffer

# TODO
class PPOAgent(Agent):
    def __init__(self):
        raise NotImplementedError

    def act(self, observation, reward, done):
        raise NotImplementedError

    def save(self):
        raise NotImplementedError

    def load(self):
        raise NotImplementedError
<commit_msg>Update ppo.py<commit_after># coding: utf-8

import numpy as np
import torch
from torch import nn
from torch import optim
from torch import cuda
from torch.autograd import Variable
from visualdl import LogWriter

from .core import Agent
from .replay_buffer import ReplayBuffer

# TODO
class PPOAgent(Agent):
    def __init__(self):
        raise NotImplementedError

    def act(self, observation, reward, done):
        raise NotImplementedError

    def save(self):
        raise NotImplementedError

    def load(self):
        raise NotImplementedError
<|endoftext|>


<commit_before>import math

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

from utils import AddBias

# TODO: In order to make this code faster:
# 1) Implement _extract_patches as a single cuda kernel
# 2) Compute QR decomposition in a separate process
# 3) Actually make a general KFAC optimizer so it fits PyTorch

def _extract_patches(x, kernel_size, stride, padding):
    if padding[0] + padding[1] > 0:
        x = F.pad(x, (padding[1], padding[1], padding[0],
                      padding[0])).data  # Actually check dims
    x = x.unfold(2, kernel_size[0], stride[0])
    x = x.unfold(3, kernel_size[1], stride[1])
    x = x.transpose_(1, 2).transpose_(2, 3).contiguous()
    x = x.view(
        x.size(0), x.size(1), x.size(2),
        x.size(3) * x.size(4) * x.size(5))
    return x

def compute_cov_a(a, classname, layer_info, fast_cnn):
    batch_size = a.size(0)

    if classname == 'Conv2d':
        if fast_cnn:
            a = _extract_patches(a, *layer_info)
            a = a.view(a.size(0), -1, a.size(-1))
            a = a.mean(1)
        else:
            a = _extract_patches(a, *layer_info)
            a = a.view(-1, a.size(-1)).div_(a.size(1)).div_(a.size(2))
    elif classname == 'AddBias':
        is_cuda = a.is_cuda
        a = torch.ones(a.size(0), 1)
        if is_cuda:
            a = a.cuda()

    return a.t() @ (a / batch_size)

def compute_cov_g(g, classname, layer_info, fast_cnn):
    batch_size = g.size(0)

    if classname == 'Conv2d':
        if fast_cnn:
            g = g.view(g.size(0), g.size(1), -1)
            g = g.sum(-1)
        else:
            g = g.transpose(1, 2).transpose(2, 3).contiguous()
            g = g.view(-1, g.size(-1)).mul_(g.size(1)).mul_(g.size(2))
    elif classname == 'AddBias':
        g = g.view(g.size(0), g.size(1), -1)
        g = g.sum(-1)

    g_ = g * batch_size
    return g_.t() @ (g_ / g.size(0))

def update_running_stat(aa, m_aa, momentum):
    # Do the trick to keep aa unchanged and not create any additional tensors
    m_aa *= momentum / (1 - momentum)
    m_aa += aa
    m_aa *= (1 - momentum)

class SplitBias(nn.Module):
    def __init__(self, module):
        super(SplitBias, self).__init__()
        self.module = module
        self.add_bias = AddBias(module.bias.data)
        self.module.bias = None

    def forward(self, input):
        x = self.module(input)
        x = self.add_bias(x)
        return x

class KFACOptimizer(optim.Optimizer):
    def __init__(self,
                 model,
                 lr=0.25,
                 momentum=0.9,
                 stat_decay=0.99,
                 kl_clip=0.001,
                 damping=1e-2,
                 weight_decay=0,
                 fast_cnn=False,
                 Ts=1,
                 Tf=10):
        defaults = dict()

        def split_bias(module):
            for mname, child in module.named_children():
                if hasattr(child, 'bias') and child.bias is not None:
                    module._modules[mname] = SplitBias(child)
                else:
                    split_bias(child)

        split_bias(model)

        super(KFACOptimizer, self).__init__(model.parameters(), defaults)

        self.known_modules = {'Linear', 'Conv2d', 'AddBias'}

        self.modules = []
        self.grad_outputs = {}

        self.model = model
        self._prepare_model()

        self.steps = 0

        self.m_aa, self.m_gg = {}, {}
        self.Q_a, self.Q_g = {}, {}
        self.d_a, self.d_g = {}, {}

        self.momentum = momentum
        self.stat_decay = stat_decay

        self.lr = lr
        self.kl_clip = kl_clip
        self.damping = damping
        self.weight_decay = weight_decay

        self.fast_cnn = fast_cnn

        self.Ts = Ts
        self.Tf = Tf

        self.optim = optim.SGD(
            model.parameters(),
            lr=self.lr * (1 - self.momentum),
            momentum=self.momentum)

    def _save_input(self, module, input):
        if torch.is_grad_enabled() and self.steps % self.Ts == 0:
            classname = module.__class__.__name__
            layer_info = None
            if classname == 'Conv2d':
                layer_info = (module.kernel_size, module.stride,
                              module.padding)

            aa = compute_cov_a(input[0].data, classname, layer_info,
                               self.fast_cnn)

            # Initialize buffers
            if self.steps == 0:
                self.m_aa[module] = aa.clone()

            update_running_stat(aa, self.m_aa[module], self.stat_decay)

    def _save_grad_output(self, module, grad_input, grad_output):
        if self.acc_stats:
            classname = module.__class__.__name__
            layer_info = None
            if classname == 'Conv2d':
                layer_info = (module.kernel_size, module.stride,
                              module.padding)

            gg = compute_cov_g(grad_output[0].data, classname, layer_info,
                               self.fast_cnn)

            # Initialize buffers
            if self.steps == 0:
                self.m_gg[module] = gg.clone()

            update_running_stat(gg, self.m_gg[module], self.stat_decay)

    def _prepare_model(self):
        for module in self.model.modules():
            classname = module.__class__.__name__
            if classname in self.known_modules:
                assert not ((classname in ['Linear', 'Conv2d']) and module.bias is not None), \
                                    "You must have a bias as a separate layer"

                self.modules.append(module)
                module.register_forward_pre_hook(self._save_input)
                module.register_backward_hook(self._save_grad_output)

    def step(self):
        # Add weight decay
        if self.weight_decay > 0:
            for p in self.model.parameters():
                p.grad.data.add_(self.weight_decay, p.data)

        updates = {}
        for i, m in enumerate(self.modules):
            assert len(list(m.parameters())
                       ) == 1, "Can handle only one parameter at the moment"
            classname = m.__class__.__name__
            p = next(m.parameters())

            la = self.damping + self.weight_decay

            if self.steps % self.Tf == 0:
                # My asynchronous implementation exists, I will add it later.
                # Experimenting with different ways to this in PyTorch.
                self.d_a[m], self.Q_a[m] = torch.symeig(
                    self.m_aa[m], eigenvectors=True)
                self.d_g[m], self.Q_g[m] = torch.symeig(
                    self.m_gg[m], eigenvectors=True)

                self.d_a[m].mul_((self.d_a[m] > 1e-6).float())
                self.d_g[m].mul_((self.d_g[m] > 1e-6).float())

            if classname == 'Conv2d':
                p_grad_mat = p.grad.data.view(p.grad.data.size(0), -1)
            else:
                p_grad_mat = p.grad.data

            v1 = self.Q_g[m].t() @ p_grad_mat @ self.Q_a[m]
            v2 = v1 / (
                self.d_g[m].unsqueeze(1) * self.d_a[m].unsqueeze(0) + la)
            v = self.Q_g[m] @ v2 @ self.Q_a[m].t()

            v = v.view(p.grad.data.size())
            updates[p] = v

        vg_sum = 0
        for p in self.model.parameters():
            v = updates[p]
            vg_sum += (v * p.grad.data * self.lr * self.lr).sum()

        nu = min(1, math.sqrt(self.kl_clip / vg_sum))

        for p in self.model.parameters():
            v = updates[p]
            p.grad.data.copy_(v)
            p.grad.data.mul_(nu)

        self.optim.step()
        self.steps += 1
<commit_msg>Update kfac.py<commit_after>import math

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

from utils import AddBias

# TODO: In order to make this code faster:
# 1) Implement _extract_patches as a single cuda kernel
# 2) Compute QR decomposition in a separate process
# 3) Actually make a general KFAC optimizer so it fits PyTorch

def _extract_patches(x, kernel_size, stride, padding):
    if padding[0] + padding[1] > 0:
        x = F.pad(x, (padding[1], padding[1], padding[0],
                      padding[0])).data  # Actually check dims
    x = x.unfold(2, kernel_size[0], stride[0])
    x = x.unfold(3, kernel_size[1], stride[1])
    x = x.transpose_(1, 2).transpose_(2, 3).contiguous()
    x = x.view(
        x.size(0), x.size(1), x.size(2),
        x.size(3) * x.size(4) * x.size(5))
    return x

def compute_cov_a(a, classname, layer_info, fast_cnn):
    batch_size = a.size(0)

    if classname == 'Conv2d':
        if fast_cnn:
            a = _extract_patches(a, *layer_info)
            a = a.view(a.size(0), -1, a.size(-1))
            a = a.mean(1)
        else:
            a = _extract_patches(a, *layer_info)
            a = a.view(-1, a.size(-1)).div_(a.size(1)).div_(a.size(2))
    elif classname == 'AddBias':
        is_cuda = a.is_cuda
        a = torch.ones(a.size(0), 1)
        if is_cuda:
            a = a.cuda()

    return a.t() @ (a / batch_size)

def compute_cov_g(g, classname, layer_info, fast_cnn):
    batch_size = g.size(0)

    if classname == 'Conv2d':
        if fast_cnn:
            g = g.view(g.size(0), g.size(1), -1)
            g = g.sum(-1)
        else:
            g = g.transpose(1, 2).transpose(2, 3).contiguous()
            g = g.view(-1, g.size(-1)).mul_(g.size(1)).mul_(g.size(2))
    elif classname == 'AddBias':
        g = g.view(g.size(0), g.size(1), -1)
        g = g.sum(-1)

    g_ = g * batch_size
    return g_.t() @ (g_ / g.size(0))

def update_running_stat(aa, m_aa, momentum):
    # Do the trick to keep aa unchanged and not create any additional tensors
    m_aa *= momentum / (1 - momentum)
    m_aa += aa
    m_aa *= (1 - momentum)

class SplitBias(nn.Module):
    def __init__(self, module):
        super(SplitBias, self).__init__()
        self.module = module
        self.add_bias = AddBias(module.bias.data)
        self.module.bias = None

    def forward(self, input):
        x = self.module(input)
        x = self.add_bias(x)
        return x

class KFACOptimizer(optim.Optimizer):
    def __init__(self,
                 model,
                 lr=0.25,
                 momentum=0.9,
                 stat_decay=0.99,
                 kl_clip=0.001,
                 damping=1e-2,
                 weight_decay=0,
                 fast_cnn=False,
                 Ts=1,
                 Tf=10):
        defaults = dict()

        def split_bias(module):
            for mname, child in module.named_children():
                if hasattr(child, 'bias') and child.bias is not None:
                    module._modules[mname] = SplitBias(child)
                else:
                    split_bias(child)

        split_bias(model)

        super(KFACOptimizer, self).__init__(model.parameters(), defaults)

        self.known_modules = {'Linear', 'Conv2d', 'AddBias'}

        self.modules = []
        self.grad_outputs = {}

        self.model = model
        self._prepare_model()

        self.steps = 0

        self.m_aa, self.m_gg = {}, {}
        self.Q_a, self.Q_g = {}, {}
        self.d_a, self.d_g = {}, {}

        self.momentum = momentum
        self.stat_decay = stat_decay

        self.lr = lr
        self.kl_clip = kl_clip
        self.damping = damping
        self.weight_decay = weight_decay

        self.fast_cnn = fast_cnn

        self.Ts = Ts
        self.Tf = Tf

        self.optim = optim.SGD(
            model.parameters(),
            lr=self.lr * (1 - self.momentum),
            momentum=self.momentum)

    def _save_input(self, module, input):
        if torch.is_grad_enabled() and self.steps % self.Ts == 0:
            classname = module.__class__.__name__
            layer_info = None
            if classname == 'Conv2d':
                layer_info = (module.kernel_size, module.stride,
                              module.padding)

            aa = compute_cov_a(input[0].data, classname, layer_info,
                               self.fast_cnn)

            # Initialize buffers
            if self.steps == 0:
                self.m_aa[module] = aa.clone()

            update_running_stat(aa, self.m_aa[module], self.stat_decay)

    def _save_grad_output(self, module, grad_input, grad_output):
        # Accumulate statistics for Fisher matrices
        if self.acc_stats:
            classname = module.__class__.__name__
            layer_info = None
            if classname == 'Conv2d':
                layer_info = (module.kernel_size, module.stride,
                              module.padding)

            gg = compute_cov_g(grad_output[0].data, classname, layer_info,
                               self.fast_cnn)

            # Initialize buffers
            if self.steps == 0:
                self.m_gg[module] = gg.clone()

            update_running_stat(gg, self.m_gg[module], self.stat_decay)

    def _prepare_model(self):
        for module in self.model.modules():
            classname = module.__class__.__name__
            if classname in self.known_modules:
                assert not ((classname in ['Linear', 'Conv2d']) and module.bias is not None), \
                                    "You must have a bias as a separate layer"

                self.modules.append(module)
                module.register_forward_pre_hook(self._save_input)
                module.register_backward_hook(self._save_grad_output)

    def step(self):
        # Add weight decay
        if self.weight_decay > 0:
            for p in self.model.parameters():
                p.grad.data.add_(self.weight_decay, p.data)

        updates = {}
        for i, m in enumerate(self.modules):
            assert len(list(m.parameters())
                       ) == 1, "Can handle only one parameter at the moment"
            classname = m.__class__.__name__
            p = next(m.parameters())

            la = self.damping + self.weight_decay

            if self.steps % self.Tf == 0:
                # My asynchronous implementation exists, I will add it later.
                # Experimenting with different ways to this in PyTorch.
                self.d_a[m], self.Q_a[m] = torch.symeig(
                    self.m_aa[m], eigenvectors=True)
                self.d_g[m], self.Q_g[m] = torch.symeig(
                    self.m_gg[m], eigenvectors=True)

                self.d_a[m].mul_((self.d_a[m] > 1e-6).float())
                self.d_g[m].mul_((self.d_g[m] > 1e-6).float())

            if classname == 'Conv2d':
                p_grad_mat = p.grad.data.view(p.grad.data.size(0), -1)
            else:
                p_grad_mat = p.grad.data

            v1 = self.Q_g[m].t() @ p_grad_mat @ self.Q_a[m]
            v2 = v1 / (
                self.d_g[m].unsqueeze(1) * self.d_a[m].unsqueeze(0) + la)
            v = self.Q_g[m] @ v2 @ self.Q_a[m].t()

            v = v.view(p.grad.data.size())
            updates[p] = v

        vg_sum = 0
        for p in self.model.parameters():
            v = updates[p]
            vg_sum += (v * p.grad.data * self.lr * self.lr).sum()

        nu = min(1, math.sqrt(self.kl_clip / vg_sum))

        for p in self.model.parameters():
            v = updates[p]
            p.grad.data.copy_(v)
            p.grad.data.mul_(nu)

        self.optim.step()
        self.steps += 1
<|endoftext|>


<commit_before># mapped alternative names
alternative_names = {
    'drboom' : 'drbalance',
    'fierywaraxe' : 'fierywinaxe',
    'lordjaraxxus' : 'jaraxxus',
    'alexstrasza' : 'alex',
    'auchenaisoulpriest' : ['auchenai', 'soulpriest'],
    'malygos' : 'maly',
    'pyroblast' : 'pyro',
    'mysteriouschallenger' : 'mc',
    'ragnarosthefirelord' : ['ragnaros', 'rag'],
    'tirionfordring' : 'tirion',
    'archmageantonidas' : 'antonidas',
    'ancientofwar' : 'aow',
    'ancientoflore' : 'aol',
    'prophetvelen' : 'velen',
    'emperorthaurissan' : ['emperor', 'thaurissan'],
    'illidanstormrage' : 'illidan',
    'sylvanaswindrunner' : 'sylvanas',
    'antiquehealbot' : 'healbot',
    'grimpatron' : 'patron',
    'sludgebelcher' : 'belcher',
    'truesilverchampion' : 'truesilver',
    'polymorph' : 'sheep',
    'hex' : 'frog',
    'kezanmystic' : 'kezan',
    'biggamehunter' : 'bgh',
    'mindcontroltech' : 'mct',
    'lorewalkercho' : 'cho',
    'millhousemanastorm' : 'millhouse',
    'shieldedminibot' : 'minibot',
    'sorcerersapprentice' : 'apprentice',
    'northshirecleric' : 'northshire',
    'defenderofargus' : 'argus',
    'pilotedshredder' : 'shredder',
    'hemetnesingwary' : 'hemet',
    'ironbeakowl' : 'owl',
    'sunfuryprotector' : 'sunfury',
    'senjinshieldmasta' : 'tazdingo',
    'preparation' : 'prep',
    'theblackknight' : ['tbk', 'blackknight'],
    'elitetaurenchieftain' : 'etc',
    'harrisonjones' : ['harrison', 'jones'],
    'kelthuzad' : 'kel',
    'eydisdarkbane' : ['darkbane', 'eydis'],
    'fjolalightbane' : ['lightbane', 'fjola'],
    'velenschosen' : 'velens',
    'varianwrynn' : 'varian',
    'polymorphboar' : 'boar',
    'themistcaller' : 'mistcaller',
    'confessorpaletress' : ['confessor', 'paletress'],
    'theskeletonknight' : 'skeletonknight',
    'nexuschampionsaraad' : ['sarad', 'saraad'],
    'edwinvancleef' : ['edwin', 'vancleef'],
    'gormoktheimpaler' : 'gormok',
    'eadricthepure' : 'eadric',
    'anubarak' : 'anub',
    'wilfredfizzlebang' : ['wilfred', 'fizzlebang'],
    'summoningportal' : 'portal',
    'cairnebloodhoof' : 'cairne',
    'draeneitotemcarver' : 'totemcarver',
    'justicartrueheart' : ['justicar', 'trueheart'],
    'sneedsoldshredder' : 'sneeds',
    'alakirthewindlord' : 'alakir',
    'bolvarfordragon' : 'bolvar',
    'troggzortheearthinator' : 'troggzor',

    'archthiefrafaam' : 'rafaam',
    'renojackson' : ['reno', 'jackson'],
    'djinniofzephyrs' : ['genie', 'djinni'],
    'elisestarseeker' : ['elise', 'starseeker'],
    'brannbronzebeard' : 'brann',
    'forgottentorch' : 'torch',
    'explorershat' : 'hat',
    'maptothegoldenmonkey' : 'map',
    'sirfinleymrrgglton' : ['sirfinley', 'finley'],

    'hallazealtheascended' : 'hallazeal',
    'tentacleofnzoth' : 'tentacle',
    'nzoththecorruptor' : 'nzoth',
    'theboogeymonster' : 'boogeymonster',
    'xarilpoisonedmind' : 'xaril',
    'muklatyrantofthevale' : 'muklatyrant',
    'yshaarjrageunbound' : 'yshaarj',
    'yoggsaronhopesend' : ['yogg', 'yoggsaron', 'rngesus'],

    'medivhtheguardian' : 'medivh',
    'princemalchezaar' : 'malchezaar',
    'thecurator': 'curator',

    'patchesthepirate' : 'patches',
    'finjatheflyingstar' : 'finja',
    'wickerflameburnbristle' : ['wickerflame', 'burnbristle'],
    'kuntheforgottenking' : 'kun',
    'mistressofmixtures' : 'mom',
    'auctionmasterbeardo' : 'beardo',
    'shakuthecollector' : 'shaku',
    'patchesthepirate' : 'patches',
    'auctionmasterbeardo' : 'beardo',
    'donhancho' : 'don',
    'inkmastersolia' : ['inkmaster', 'solia'],
    'razathechained' : 'raza'
}

# reverse dict
translations = {}
for org, alts in alternative_names.items():
    if isinstance(alts, list):
        for alt in alts:
            translations[alt] = org
    else:
        translations[alts] = org
<commit_msg>added: new alternative card names<commit_after># mapped alternative names
alternative_names = {
    'drboom' : 'drbalance',
    'fierywaraxe' : 'fierywinaxe',
    'lordjaraxxus' : 'jaraxxus',
    'alexstrasza' : 'alex',
    'auchenaisoulpriest' : ['auchenai', 'soulpriest'],
    'malygos' : 'maly',
    'pyroblast' : 'pyro',
    'mysteriouschallenger' : 'mc',
    'ragnarosthefirelord' : ['ragnaros', 'rag'],
    'tirionfordring' : 'tirion',
    'archmageantonidas' : 'antonidas',
    'ancientofwar' : 'aow',
    'ancientoflore' : 'aol',
    'prophetvelen' : 'velen',
    'emperorthaurissan' : ['emperor', 'thaurissan'],
    'illidanstormrage' : 'illidan',
    'sylvanaswindrunner' : 'sylvanas',
    'antiquehealbot' : 'healbot',
    'grimpatron' : 'patron',
    'sludgebelcher' : 'belcher',
    'truesilverchampion' : 'truesilver',
    'polymorph' : 'sheep',
    'hex' : 'frog',
    'kezanmystic' : 'kezan',
    'biggamehunter' : 'bgh',
    'mindcontroltech' : 'mct',
    'lorewalkercho' : 'cho',
    'millhousemanastorm' : 'millhouse',
    'shieldedminibot' : 'minibot',
    'sorcerersapprentice' : 'apprentice',
    'northshirecleric' : 'northshire',
    'defenderofargus' : 'argus',
    'pilotedshredder' : 'shredder',
    'hemetnesingwary' : 'hemet',
    'ironbeakowl' : 'owl',
    'sunfuryprotector' : 'sunfury',
    'senjinshieldmasta' : 'tazdingo',
    'preparation' : 'prep',
    'theblackknight' : ['tbk', 'blackknight'],
    'elitetaurenchieftain' : 'etc',
    'harrisonjones' : ['harrison', 'jones'],
    'kelthuzad' : 'kel',
    'eydisdarkbane' : ['darkbane', 'eydis'],
    'fjolalightbane' : ['lightbane', 'fjola'],
    'velenschosen' : 'velens',
    'varianwrynn' : 'varian',
    'polymorphboar' : 'boar',
    'themistcaller' : 'mistcaller',
    'confessorpaletress' : ['confessor', 'paletress'],
    'theskeletonknight' : 'skeletonknight',
    'nexuschampionsaraad' : ['sarad', 'saraad'],
    'edwinvancleef' : ['edwin', 'vancleef'],
    'gormoktheimpaler' : 'gormok',
    'eadricthepure' : 'eadric',
    'anubarak' : 'anub',
    'wilfredfizzlebang' : ['wilfred', 'fizzlebang'],
    'summoningportal' : 'portal',
    'cairnebloodhoof' : 'cairne',
    'draeneitotemcarver' : 'totemcarver',
    'justicartrueheart' : ['justicar', 'trueheart'],
    'sneedsoldshredder' : 'sneeds',
    'alakirthewindlord' : 'alakir',
    'bolvarfordragon' : 'bolvar',
    'troggzortheearthinator' : 'troggzor',

    'archthiefrafaam' : 'rafaam',
    'renojackson' : ['reno', 'jackson'],
    'djinniofzephyrs' : ['genie', 'djinni'],
    'elisestarseeker' : ['elise', 'starseeker'],
    'brannbronzebeard' : 'brann',
    'forgottentorch' : 'torch',
    'explorershat' : 'hat',
    'maptothegoldenmonkey' : 'map',
    'sirfinleymrrgglton' : ['sirfinley', 'finley'],

    'hallazealtheascended' : 'hallazeal',
    'tentacleofnzoth' : 'tentacle',
    'nzoththecorruptor' : 'nzoth',
    'theboogeymonster' : 'boogeymonster',
    'xarilpoisonedmind' : 'xaril',
    'muklatyrantofthevale' : 'muklatyrant',
    'yshaarjrageunbound' : 'yshaarj',
    'yoggsaronhopesend' : ['yogg', 'yoggsaron', 'rngesus'],

    'medivhtheguardian' : 'medivh',
    'princemalchezaar' : 'malchezaar',
    'thecurator': 'curator',

    'patchesthepirate' : 'patches',
    'finjatheflyingstar' : 'finja',
    'wickerflameburnbristle' : ['wickerflame', 'burnbristle'],
    'kuntheforgottenking' : 'kun',
    'mistressofmixtures' : 'mom',
    'auctionmasterbeardo' : 'beardo',
    'shakuthecollector' : 'shaku',
    'patchesthepirate' : 'patches',
    'auctionmasterbeardo' : 'beardo',
    'donhancho' : 'don',
    'inkmastersolia' : 'solia',
    'razathechained' : 'raza',
    'hobartgrapplehammer': 'hobart',
    'genzotheshark' : 'genzo',
    'krultheunshackled':'krul'
}

# reverse dict
translations = {}
for org, alts in alternative_names.items():
    if isinstance(alts, list):
        for alt in alts:
            translations[alt] = org
    else:
        translations[alts] = org
<|endoftext|>


<commit_before>#!/usr/bin/env python
'''
Predict Server
Create a server to accept image inputs and run them against a trained neural network.
This then sends the steering output back to the client.
Author: Tawn Kramer
'''
import os
import argparse
import sys
import numpy as np
import h5py
import json
from keras.models import model_from_json
#import matplotlib.pyplot as plt
import time
import asyncore
import json
import socket
from PIL import Image

class SteeringServer(asyncore.dispatcher):
    """Receives connections and establishes handlers for each client.
    """

    def __init__(self, address, model):
        asyncore.dispatcher.__init__(self)
        self.create_socket(socket.AF_INET, socket.SOCK_STREAM)
        self.bind(address)
        self.address = self.socket.getsockname()
        print 'binding to', self.address
        self.listen(1)
        self.model = model
        return

    def handle_accept(self):
        # Called when a client connects to our socket
        client_info = self.accept()
        #self.logger.debug('handle_accept() -> %s', client_info[1])
        print 'got a new client', client_info[1]
        h = SteeringHandler(sock=client_info[0], chunk_size=8*1024, model=self.model)
        return

    def handle_close(self):
        #self.logger.debug('handle_close()')
        self.close()
        return

class SteeringHandler(asyncore.dispatcher):
    """Handles echoing messages from a single client.
    """
    IDLE = 1
    GETTING_IMG = 2
    SENDING_STEERING = 3

    def __init__(self, sock, chunk_size=256, model=None):
        self.model = model
        self.chunk_size = chunk_size
        asyncore.dispatcher.__init__(self, sock=sock)
        self.data_to_write = []
        self.image_byes = []
        self.mode = self.IDLE
        return

    def writable(self):
        """We want to write if we have received data."""
        response = bool(self.data_to_write)
        return response

    def handle_write(self):
        """Write as much as possible of the most recent message we have received."""
        data = self.data_to_write.pop()
        sent = self.send(data[:self.chunk_size])
        if sent < len(data):
            remaining = data[sent:]
            self.data.to_write.append(remaining)
        elif self.mode == self.SENDING_STEERING:
          self.mode = self.IDLE

    def handle_read(self):
        """Read an incoming message from the client and put it into our outgoing queue."""
        data = self.recv(self.chunk_size)
        view_image = False
        #print 'got', len(data), 'bytes'
        if len(data) == 0:
          self.handle_close()
        elif self.mode == self.IDLE:
          try:
            jsonObj = json.loads(data)
            self.num_bytes = jsonObj['num_bytes']
            self.width = jsonObj['width']
            self.height = jsonObj['height']
            self.num_channels = jsonObj['num_channels']
            self.format = jsonObj['format']
            self.flip_y = jsonObj['flip_y']
            self.data_to_write.insert(0, "{ 'response' : 'ready_for_image' }")
            self.mode = self.GETTING_IMG
            self.image_byes = []
            self.num_read = 0
          except:
            self.mode = self.IDLE
            print 'failed to read json from: ', data
        elif self.mode == self.GETTING_IMG:
          self.image_byes.append(data)
          self.num_read += len(data)
          if self.num_read == self.num_bytes:
            lin_arr = np.fromstring(''.join(self.image_byes), dtype=np.uint8)
            self.mode = self.SENDING_STEERING
            if self.format == 'array_of_pixels':
              img = lin_arr.reshape(self.width, self.height, self.num_channels)
              if self.flip_y:
                img = np.flipud(img)
              img = img.transpose()
            else: #assumed to be ArrayOfChannels
              img = lin_arr.reshape(self.num_channels, self.width, self.height)
            if view_image:
              #this can be useful when validating that you have your images coming in correctly.
              vis_img = Image.fromarray(img.transpose(), 'RGB')
              vis_img.show()
              #this launches too many windows if we leave it up.
              self.handle_close()
            steering = self.model.predict(img[None, :, :, :])
            reply = '{ "steering" : "%f" }' % steering
            self.data_to_write.append(reply)
        else:
            print "wasn't prepared to recv request!"

    def handle_close(self):
        self.close()

# ***** main loop *****
if __name__ == "__main__":
  parser = argparse.ArgumentParser(description='prediction server')
  parser.add_argument('model', type=str, help='model name. no json or keras.')
  parser.add_argument('--model-path', dest='path', default='../outputs/steering_model', help='model dir')
  args = parser.parse_args()

  model_json = os.path.join(args.path, args.model +'.json')

  with open(model_json, 'r') as jfile:
    model = model_from_json(json.load(jfile))

  model.compile("sgd", "mse")
  weights_file = model_json.replace('json', 'keras')
  model.load_weights(weights_file)

  address = ('0.0.0.0', 9090)
  s = SteeringServer(address, model)
  try:
    asyncore.loop()
  except KeyboardInterrupt:
    print 'stopping'

<commit_msg>tweak to predict server so it can be launched as a sub-process<commit_after>#!/usr/bin/env python
'''
Predict Server
Create a server to accept image inputs and run them against a trained neural network.
This then sends the steering output back to the client.
Author: Tawn Kramer
'''
import os
import argparse
import sys
import numpy as np
import h5py
import json
from keras.models import model_from_json
#import matplotlib.pyplot as plt
import time
import asyncore
import json
import socket
from PIL import Image

class SteeringServer(asyncore.dispatcher):
    """Receives connections and establishes handlers for each client.
    """

    def __init__(self, address, model):
        asyncore.dispatcher.__init__(self)
        self.create_socket(socket.AF_INET, socket.SOCK_STREAM)
        self.bind(address)
        self.address = self.socket.getsockname()
        print 'binding to', self.address
        self.listen(1)
        self.model = model
        return

    def handle_accept(self):
        # Called when a client connects to our socket
        client_info = self.accept()
        #self.logger.debug('handle_accept() -> %s', client_info[1])
        print 'got a new client', client_info[1]
        h = SteeringHandler(sock=client_info[0], chunk_size=8*1024, model=self.model)
        return

    def handle_close(self):
        #self.logger.debug('handle_close()')
        self.close()
        return

class SteeringHandler(asyncore.dispatcher):
    """Handles echoing messages from a single client.
    """
    IDLE = 1
    GETTING_IMG = 2
    SENDING_STEERING = 3

    def __init__(self, sock, chunk_size=256, model=None):
        self.model = model
        self.chunk_size = chunk_size
        asyncore.dispatcher.__init__(self, sock=sock)
        self.data_to_write = []
        self.image_byes = []
        self.mode = self.IDLE
        return

    def writable(self):
        """We want to write if we have received data."""
        response = bool(self.data_to_write)
        return response

    def handle_write(self):
        """Write as much as possible of the most recent message we have received."""
        data = self.data_to_write.pop()
        sent = self.send(data[:self.chunk_size])
        if sent < len(data):
            remaining = data[sent:]
            self.data.to_write.append(remaining)
        elif self.mode == self.SENDING_STEERING:
          self.mode = self.IDLE

    def handle_read(self):
        """Read an incoming message from the client and put it into our outgoing queue."""
        data = self.recv(self.chunk_size)
        view_image = False
        #print 'got', len(data), 'bytes'
        if len(data) == 0:
          self.handle_close()
        elif self.mode == self.IDLE:
          try:
            jsonObj = json.loads(data)
            self.num_bytes = jsonObj['num_bytes']
            self.width = jsonObj['width']
            self.height = jsonObj['height']
            self.num_channels = jsonObj['num_channels']
            self.format = jsonObj['format']
            self.flip_y = jsonObj['flip_y']
            self.data_to_write.insert(0, "{ 'response' : 'ready_for_image' }")
            self.mode = self.GETTING_IMG
            self.image_byes = []
            self.num_read = 0
          except:
            self.mode = self.IDLE
            print 'failed to read json from: ', data
        elif self.mode == self.GETTING_IMG:
          self.image_byes.append(data)
          self.num_read += len(data)
          if self.num_read == self.num_bytes:
            lin_arr = np.fromstring(''.join(self.image_byes), dtype=np.uint8)
            self.mode = self.SENDING_STEERING
            if self.format == 'array_of_pixels':
              img = lin_arr.reshape(self.width, self.height, self.num_channels)
              if self.flip_y:
                img = np.flipud(img)
              img = img.transpose()
            else: #assumed to be ArrayOfChannels
              img = lin_arr.reshape(self.num_channels, self.width, self.height)
            if view_image:
              #this can be useful when validating that you have your images coming in correctly.
              vis_img = Image.fromarray(img.transpose(), 'RGB')
              vis_img.show()
              #this launches too many windows if we leave it up.
              self.handle_close()
            steering = self.model.predict(img[None, :, :, :])
            reply = '{ "steering" : "%f" }' % steering
            self.data_to_write.append(reply)
        else:
            print "wasn't prepared to recv request!"

    def handle_close(self):
        self.close()

def go(model_json, address):
  with open(model_json, 'r') as jfile:
      model = model_from_json(json.load(jfile))

  model.compile("sgd", "mse")
  weights_file = model_json.replace('json', 'keras')
  model.load_weights(weights_file)

  s = SteeringServer(address, model)
  try:
    asyncore.loop()
  except KeyboardInterrupt:
    print 'stopping'

# ***** main loop *****
if __name__ == "__main__":
  parser = argparse.ArgumentParser(description='prediction server')
  parser.add_argument('model', type=str, help='model name. no json or keras.')
  parser.add_argument('--model-path', dest='path', default='../outputs/steering_model', help='model dir')
  args = parser.parse_args()

  model_json = os.path.join(args.path, args.model +'.json')
  address = ('0.0.0.0', 9090)
  go(model_json, address)

<|endoftext|>