Fix sg names

Author: Octogonapus

axon/client.py

@@ -160,7 +160,7 @@ def ensure_ecs_security_group(region):
     :param region: The region, or `None` to pull the region from the environment.
     :return: The GroupId of the SecurityGroup.
     """
-    sg_name = "axon-autogenerated-sg-ecs"
+    sg_name = "axon-autogenerated-ecs-sg"
     client = make_client("ec2", region)
     sg_id = get_single_security_group(client, sg_name, "Axon autogenerated for ECS.")
     ensure_ecs_gress(sg_id, region)
@@ -173,7 +173,7 @@ def ensure_ec2_security_group(region):
     :param region: The region, or `None` to pull the region from the environment.
     :return: The GroupId of the SecurityGroup.
     """
-    sg_name = "axon-autogenerated-sg-ec2"
+    sg_name = "axon-autogenerated-ec2-sg"
     client = make_client("ec2", region)
     sg_id = get_single_security_group(client, sg_name, "Axon autogenerated for EC2.")
     ensure_ec2_gress(sg_id, region)
@@ -631,46 +631,47 @@ def impl_download_trained_model_file(model_path, bucket_name, region):
     print("Downloaded from: {}\n".format(key))
 
 
-def impl_download_training_script(script_name, bucket_name, region):
+def impl_download_training_script(script_path, bucket_name, region):
     """
     Downloads a training script from S3.
 
-    :param script_name: The filename of the script to download (must be in the current directory).
+    :param script_path: The file path to download to, ending with the name of the script.
     :param bucket_name: The S3 bucket name.
     :param region: The region, or `None` to pull the region from the environment.
     """
     client = make_client("s3", region)
-    remote_path = "axon-training-scripts/" + os.path.basename(script_name)
-    client.download_file(bucket_name, remote_path, script_name)
-    print("Downloaded from: {}\n".format(remote_path))
+    key = "axon-training-scripts/" + os.path.basename(script_path)
+    client.download_file(bucket_name, key, script_path)
+    print("Downloaded from: {}\n".format(key))
 
 
-def impl_upload_dataset(dataset_name, bucket_name, region):
+def impl_upload_dataset(dataset_path, bucket_name, region):
     """
     Uploads a dataset to S3.
 
-    :param dataset_name: The filename of the dataset to upload (must be in the current directory).
+    :param dataset_path: The file path to the dataset to upload, ending with the name of the
+    dataset.
     :param bucket_name: The S3 bucket name.
     :param region: The region, or `None` to pull the region from the environment.
     """
     client = make_client("s3", region)
-    remote_path = "axon-datasets/" + os.path.basename(dataset_name)
-    client.upload_file(dataset_name, bucket_name, remote_path)
-    print("Uploaded to: {}\n".format(remote_path))
+    key = "axon-datasets/" + os.path.basename(dataset_path)
+    client.upload_file(dataset_path, bucket_name, key)
+    print("Uploaded to: {}\n".format(key))
 
 
-def impl_download_dataset(dataset_name, bucket_name, region):
+def impl_download_dataset(dataset_path, bucket_name, region):
     """
     Downloads a dataset from S3.
 
-    :param dataset_name: The filename of the dataset to download (must be in the current directory).
+    :param dataset_path: The file path to download to, ending with the name of the dataset.
     :param bucket_name: The S3 bucket name.
     :param region: The region, or `None` to pull the region from the environment.
     """
     client = make_client("s3", region)
-    remote_path = "axon-datasets/" + os.path.basename(dataset_name)
-    client.download_file(bucket_name, remote_path, dataset_name)
-    print("Downloaded from: {}\n".format(remote_path))
+    key = "axon-datasets/" + os.path.basename(dataset_path)
+    client.download_file(bucket_name, key, dataset_path)
+    print("Downloaded from: {}\n".format(key))
 
 
 def impl_update_training_progress(model_name, dataset_name, progress_text, bucket_name, region):
@@ -701,15 +702,24 @@ def cli():
     return
 
 
-# TODO: Don't set a default value for region in any of these
+region_choices = ['us-east-1', 'us-east-2', 'us-west-1', 'us-west-2', 'ca-central-1',
+                  'eu-central-1', 'eu-west-1', 'eu-west-2', 'eu-west-3',
+                  'eu-north-1', 'ap-east-1', 'ap-south-1', 'ap-northeast-1',
+                  'ap-northeast-2', 'ap-northeast-3', 'ap-southeast-1',
+                  'ap-southeast-2', 'me-south-1', 'sa-east-1']
+
 
 @cli.command(name="start-axon")
-@click.argument("cluster-name")
-@click.argument("task-family")
+@click.option("--region", help="The region to connect to.",
+              type=click.Choice(region_choices))
 @click.option("--revision", default=None,
               help="The revision of the task. Set to None to use the latest revision.")
-@click.option("--region", help="The region to connect to.")
-def start_axon(cluster_name, task_family, revision, region):
+def start_axon(revision, region):
+    """
+    Starts Axon and opens the ECS server in the default web browser.
+    """
+    cluster_name = "axon-autogenerated-cluster"
+    task_family = "axon-autogenerated-task-family"
     impl_ensure_configuration(cluster_name, task_family, region)
     task_arn = impl_start_task(cluster_name, task_family, revision, region)
     print("Started task: {}".format(task_arn))
@@ -722,22 +732,29 @@ def start_axon(cluster_name, task_family, revision, region):
 
 
 @cli.command(name="ensure-configuration")
-@click.argument("cluster-name")
-@click.argument("task-family")
-@click.option("--region", help="The region to connect to.")
-def ensure_configuration(cluster_name, task_family, region):
-    impl_ensure_configuration(cluster_name, task_family, region)
+@click.option("--region", help="The region to connect to.",
+              type=click.Choice(region_choices))
+def ensure_configuration(region):
+    """
+    Ensures that AWS is configured for Axon.
+    """
+    impl_ensure_configuration("axon-autogenerated-cluster", "axon-autogenerated-task-family",
+                              region)
 
 
 @cli.command(name="start-task")
-@click.argument("cluster-name")
-@click.argument("task-family")
+@click.option("--region", help="The region to connect to.",
+              type=click.Choice(region_choices))
 @click.option("--revision", default=None,
               help="The revision of the task. Set to None to use the latest revision.")
-@click.option("--region", help="The region to connect to.")
 @click.option("--stop-after/--no-stop-after", default=False,
               help="Whether to stop the task immediately after creating it.")
-def start_task(cluster_name, task_family, revision, region, stop_after):
+def start_task(revision, region, stop_after):
+    """
+    Starts a task.
+    """
+    cluster_name = "axon-autogenerated-cluster"
+    task_family = "axon-autogenerated-task-family"
     impl_ensure_configuration(cluster_name, task_family, region)
     task_arn = impl_start_task(cluster_name, task_family, revision, region)
     print("Started task: {}".format(task_arn))
@@ -750,82 +767,138 @@ def start_task(cluster_name, task_family, revision, region, stop_after):
 
 
 @cli.command(name="stop-task")
-@click.argument("cluster-name")
-@click.argument("task")
-@click.option("--region", help="The region to connect to.")
-def stop_task(cluster_name, task, region):
-    impl_stop_task(cluster_name, task, region)
+@click.argument("task-arn")
+@click.option("--region", help="The region to connect to.",
+              type=click.Choice(region_choices))
+def stop_task(task_arn, region):
+    """
+    Stops a task.
+
+    TASK_ARN The ARN of the task to stop. Given to you from `start_task` or `start_axon`.
+    """
+    impl_stop_task("axon-autogenerated-cluster", task_arn, region)
 
 
 @cli.command(name="get-container-ip")
-@click.argument("cluster-name")
-@click.argument("task")
-@click.option("--region", help="The region to connect to.")
-def get_container_ip(cluster_name, task, region):
-    print(impl_get_task_ip(cluster_name, task, region))
+@click.argument("task-arn")
+@click.option("--region", help="The region to connect to.",
+              type=click.Choice(region_choices))
+def get_container_ip(task_arn, region):
+    """
+    Gets the IP of the container the task is running on.
+
+    TASK_ARN The ARN of the task.
+    """
+    print(impl_get_task_ip("axon-autogenerated-cluster", task_arn, region))
 
 
 @cli.command(name="upload-untrained-model-file")
-@click.argument("model-name")
-@click.argument("bucket-name")
-@click.option("--region", help="The region to connect to.")
-def upload_untrained_model_file(model_name, bucket_name, region):
-    impl_upload_untrained_model_file(model_name, bucket_name, region)
+@click.argument("model-path")
+@click.option("--region", help="The region to connect to.",
+              type=click.Choice(region_choices))
+def upload_untrained_model_file(model_path, region):
+    """
+    Uploads an untrained model from a local file.
+
+    MODEL_PATH The path to the model to upload, ending with the name of the model.
+    """
+    impl_upload_untrained_model_file(model_path, ensure_s3_bucket(region), region)
 
 
 @cli.command(name="download-untrained-model-file")
-@click.argument("model-name")
-@click.argument("bucket-name")
-@click.option("--region", help="The region to connect to.")
-def download_untrained_model_file(model_name, bucket_name, region):
-    impl_download_untrained_model_file(model_name, bucket_name, region)
+@click.argument("model-path")
+@click.option("--region", help="The region to connect to.",
+              type=click.Choice(region_choices))
+def download_untrained_model_file(model_path, region):
+    """
+    Downloads an untrained model to a local file.
+
+    MODEL_PATH The path to download the model to, ending with the name of the model.
+    """
+    impl_download_untrained_model_file(model_path, ensure_s3_bucket(region), region)
 
 
 @cli.command(name="upload-trained-model-file")
-@click.argument("model-name")
-@click.argument("bucket-name")
-@click.option("--region", help="The region to connect to.")
-def upload_trained_model_file(model_name, bucket_name, region):
-    impl_upload_trained_model_file(model_name, bucket_name, region)
+@click.argument("model-path")
+@click.option("--region", help="The region to connect to.",
+              type=click.Choice(region_choices))
+def upload_trained_model_file(model_path, region):
+    """
+    Uploads a trained model from a local file.
+
+    MODEL_PATH The path to the model to upload, ending with the name of the model.
+    """
+    impl_upload_trained_model_file(model_path, ensure_s3_bucket(region), region)
 
 
 @cli.command(name="download-trained-model-file")
-@click.argument("model-name")
-@click.argument("bucket-name")
-@click.option("--region", help="The region to connect to.")
-def download_trained_model_file(model_name, bucket_name, region):
-    impl_download_trained_model_file(model_name, bucket_name, region)
+@click.argument("model-path")
+@click.option("--region", help="The region to connect to.",
+              type=click.Choice(region_choices))
+def download_trained_model_file(model_path, region):
+    """
+    Downloads a trained model to a local file.
 
+    MODEL_PATH The path to download the model to, ending with the name of the model.
+    """
+    impl_download_trained_model_file(model_path, ensure_s3_bucket(region), region)
 
-@cli.command(name="download-training-script")
-@click.argument("script-name")
-@click.argument("bucket-name")
-@click.option("--region", help="The region to connect to.")
-def download_training_script(script_name, bucket_name, region):
-    impl_download_training_script(script_name, bucket_name, region)
 
+@cli.command(name="download-training-script")
+@click.argument("script-path")
+@click.option("--region", help="The region to connect to.",
+              type=click.Choice(region_choices))
+def download_training_script(script_path, region):
+    """
+    Downloads a training script.
 
-@cli.command(name="download-dataset")
-@click.argument("dataset-name")
-@click.argument("bucket-name")
-@click.option("--region", help="The region to connect to.")
-def download_dataset(dataset_name, bucket_name, region):
-    impl_download_dataset(dataset_name, bucket_name, region)
+    SCRIPT_PATH The path to download the script to, ending with the name of the script.
+    """
+    impl_download_training_script(script_path, ensure_s3_bucket(region), region)
 
 
 @cli.command(name="upload-dataset")
-@click.argument("dataset-name")
-@click.argument("bucket-name")
-@click.option("--region", help="The region to connect to.")
-def upload_dataset(dataset_name, bucket_name, region):
-    impl_upload_dataset(dataset_name, bucket_name, region)
+@click.argument("dataset-path")
+@click.option("--region", help="The region to connect to.",
+              type=click.Choice(region_choices))
+def upload_dataset(dataset_path, region):
+    """
+    Uploads a dataset.
+
+    DATASET_PATH The path to the dataset to upload, ending with the name of the dataset.
+    """
+    impl_upload_dataset(dataset_path, ensure_s3_bucket(region), region)
+
+
+@cli.command(name="download-dataset")
+@click.argument("dataset-path")
+@click.option("--region", help="The region to connect to.",
+              type=click.Choice(region_choices))
+def download_dataset(dataset_path, region):
+    """
+    Downloads a dataset.
+
+    DATASET_PATH The path to download the dataset to, ending with the name of the dataset.
+    """
+    impl_download_dataset(dataset_path, ensure_s3_bucket(region), region)
 
 
 @cli.command(name="update-training-progress")
 @click.argument("model-name")
 @click.argument("dataset-name")
 @click.argument("progress-text")
-@click.argument("bucket-name")
-@click.option("--region", help="The region to connect to.")
-def update_training_progress(model_name, dataset_name, progress_text, bucket_name, region):
-    impl_update_training_progress(model_name, dataset_name, progress_text, bucket_name, region)
+@click.option("--region", help="The region to connect to.",
+              type=click.Choice(region_choices))
+def update_training_progress(model_name, dataset_name, progress_text, region):
+    """
+    Updates the training progress. Meant to be used while a training script is running to provide
+    progress updates to Axon.
+
+    MODEL_NAME The filename of the model currently being trained.
+
+    DATASET_NAME The name of the dataset currently being trained on.
+
+    PROGRESS_TEXT The text to write to the progress file.
+    """
+    impl_update_training_progress(model_name, dataset_name, progress_text, ensure_s3_bucket(region),
+                                  region)