Let the region be None to select from the env

Author: Octogonapus

axon/client.py

@@ -4,6 +4,7 @@
 import ipify
 import webbrowser
 import os.path
+import axon.progress_reporter
 
 all_perm = {
     "FromPort": -1,
@@ -22,16 +23,23 @@
 }
 
 
+def make_client(name, region):
+    if region is None:
+        return boto3.client(name)
+    else:
+        return boto3.client(name, region_name=region)
+
+
 def ensure_log_group(group_name, region):
     """
     Ensures that a log group is present. If there is a matching log group, nothing is created.
     If there is no matching log group, one is created.
 
     :param group_name: The name of the log group.
-    :param region: The region.
+    :param region: The region, or `None` to pull the region from the environment.
     :return: Nothing.
     """
-    client = boto3.client("logs", region_name=region)
+    client = make_client("logs", region)
     matching_log_groups = client.describe_log_groups(
         logGroupNamePrefix=group_name
     )
@@ -66,7 +74,7 @@ def ensure_ecs_gress(sg_id, region):
     egress permissions are revoked. The permissions that Axon needs are authorized.
 
     :param sg_id: The SecurityGroup's GroupId.
-    :param region: The region.
+    :param region: The region, or `None` to pull the region from the environment.
     :return: Nothing.
     """
     ec2 = boto3.resource('ec2', region_name=region)
@@ -92,7 +100,7 @@ def ensure_ec2_gress(sg_id, region):
     egress permissions are revoked. The permissions that Axon needs are authorized.
 
     :param sg_id: The SecurityGroup's GroupId.
-    :param region: The region.
+    :param region: The region, or `None` to pull the region from the environment.
     :return: Nothing.
     """
     ec2 = boto3.resource('ec2', region_name=region)
@@ -144,12 +152,11 @@ def get_single_security_group(client, sg_name, desc):
 def ensure_ecs_security_group(region):
     """
     Ensures that the ECS SecurityGroup exists.
-
-    :param region: The region.
+    :param region: The region, or `None` to pull the region from the environment.
     :return: The GroupId of the SecurityGroup.
     """
     sg_name = "axon-ecs-autogenerated"
-    client = boto3.client("ec2", region_name=region)
+    client = make_client("ec2", region)
     sg_id = get_single_security_group(client, sg_name, "Axon autogenerated for ECS.")
     ensure_ecs_gress(sg_id, region)
     return sg_id
@@ -158,12 +165,11 @@ def ensure_ecs_security_group(region):
 def ensure_ec2_security_group(region):
     """
     Ensures that the EC2 SecurityGroup exists.
-
-    :param region: The region.
+    :param region: The region, or `None` to pull the region from the environment.
     :return: The GroupId of the SecurityGroup.
     """
     sg_name = "axon-ec2-autogenerated"
-    client = boto3.client("ec2", region_name=region)
+    client = make_client("ec2", region)
     sg_id = get_single_security_group(client, sg_name, "Axon autogenerated for EC2.")
     ensure_ec2_gress(sg_id, region)
     return sg_id
@@ -172,12 +178,11 @@ def ensure_ec2_security_group(region):
 def select_subnet(region):
     """
     Picks the first available subnet.
-
-    :param region: The region.
+    :param region: The region, or `None` to pull the region from the environment.
     :return: The SubnetId.
     """
-    ec2 = boto3.client("ec2", region_name=region)
-    return ec2.describe_subnets(Filters=[])["Subnets"][0]["SubnetId"]
+    client = make_client("ec2", region)
+    return client.describe_subnets(Filters=[])["Subnets"][0]["SubnetId"]
 
 
 def ensure_role(client, role_name):
@@ -207,12 +212,11 @@ def ensure_task_role(region):
 
     TODO: Fix this.
     This method does not check that a matching role has the correct policies.
-
-    :param region: The region.
+    :param region: The region, or `None` to pull the region from the environment.
     :return: The role Arn.
     """
     role_name = "axon-ecs-autogenerated-task-role"
-    client = boto3.client("iam", region_name=region)
+    client = make_client("iam", region)
     role_arn = ensure_role(client, role_name)
     if role_arn is None:
         # Need to create the role
@@ -250,7 +254,7 @@ def ensure_task_role(region):
 
 def ensure_ec2_role(region):
     role_name = "axon-ec2-role-manual"
-    client = boto3.client("iam", region_name=region)
+    client = make_client("iam", region)
     role_arn = ensure_role(client, role_name)
     if role_arn is None:
         # Need to create the role
@@ -301,7 +305,7 @@ def ensure_task(ecs_client, task_family, region, vcpu, memory):
 
     :param ecs_client: The ECS client to use.
     :param task_family: The task family name.
-    :param region: The region.
+    :param region: The region, or `None` to pull the region from the environment.
     :param vcpu: The amount of cpu in vcpu units.
     :param memory: The amount of memory in MB.
     :return: The task definition's Arn.
@@ -357,10 +361,10 @@ def wait_for_task_to_start(task_arn, cluster, region):
 
     :param task_arn: The Arn of the task to wait for.
     :param cluster: The simple name of the cluster the task is in.
-    :param region: The region.
+    :param region: The region, or `None` to pull the region from the environment.
     :return: Nothing.
     """
-    client = boto3.client("ecs", region_name=region)
+    client = make_client("ecs", region)
     waiter = client.get_waiter("tasks_running")
     waiter.wait(cluster=cluster, tasks=[task_arn])
 
@@ -371,9 +375,9 @@ def impl_ensure_configuration(cluster_name, task_family, region):
 
     :param cluster_name: The simple name of the cluster to start the task in.
     :param task_family: The family of the task to start.
-    :param region: The region.
+    :param region: The region, or `None` to pull the region from the environment.
     """
-    client = boto3.client("ecs", region_name=region)
+    client = make_client("ecs", region)
     ensure_cluster(client, cluster_name)
     ensure_task(client, task_family, region, 2048, 4096)
     ensure_ec2_security_group(region)
@@ -392,10 +396,10 @@ def impl_start_task(cluster_name, task_family, revision, region):
     :param cluster_name: The simple name of the cluster to start the task in.
     :param task_family: The family of the task to start.
     :param revision: A task definition revision number, or None to use the latest revision.
-    :param region: The region.
+    :param region: The region, or `None` to pull the region from the environment.
     :return: The started task's Arn.
     """
-    client = boto3.client("ecs", region_name=region)
+    client = make_client("ecs", region)
 
     impl_ensure_configuration(cluster_name, task_family, region)
 
@@ -431,10 +435,10 @@ def impl_stop_task(cluster_name, task_arn, region):
 
     :param cluster_name: The simple name of the cluster.
     :param task_arn: The Arn of the task.
-    :param region: The region.
+    :param region: The region, or `None` to pull the region from the environment.
     :return: Nothing.
     """
-    client = boto3.client("ecs", region_name=region)
+    client = make_client("ecs", region)
     client.stop_task(cluster=cluster_name, task=task_arn)
 
 
@@ -444,10 +448,10 @@ def impl_get_task_ip(cluster_name, task_arn, region):
 
     :param cluster_name: The simple name of the cluster.
     :param task_arn: The task's Arn.
-    :param region: The region.
+    :param region: The region, or `None` to pull the region from the environment.
     :return: The public IP of the task.
     """
-    client = boto3.client("ecs", region_name=region)
+    client = make_client("ecs", region)
     task_arn = client.describe_tasks(
         cluster=cluster_name,
         tasks=[task_arn]
@@ -458,7 +462,7 @@ def impl_get_task_ip(cluster_name, task_arn, region):
     eni = next(
         x["value"] for x in interface_attachment["details"] if x["name"] == "networkInterfaceId")
 
-    ec2 = boto3.client("ec2", region_name=region)
+    ec2 = make_client("ec2", region)
     nics = ec2.describe_network_interfaces(
         Filters=[
             {
@@ -476,9 +480,9 @@ def impl_upload_model_file(local_file_path, bucket_name, region):
 
     :param local_file_path: The path to the model file on disk.
     :param bucket_name: The S3 bucket name.
-    :param region: The region.
+    :param region: The region, or `None` to pull the region from the environment.
     """
-    client = boto3.client("s3", region_name=region)
+    client = make_client("s3", region)
     remote_path = "axon-uploaded-trained-models/" + os.path.basename(local_file_path)
     client.upload_file(local_file_path, bucket_name, remote_path)
     print("Uploaded to: {}\n".format(remote_path))
@@ -490,9 +494,9 @@ def impl_download_model_file(local_file_path, bucket_name, region):
 
     :param local_file_path: The path to the model file on disk.
     :param bucket_name: The S3 bucket name.
-    :param region: The region.
+    :param region: The region, or `None` to pull the region from the environment.
     """
-    client = boto3.client("s3", region_name=region)
+    client = make_client("s3", region)
     remote_path = "axon-uploaded-trained-models/" + os.path.basename(local_file_path)
     client.download_file(bucket_name, remote_path, local_file_path)
     print("Downloaded from: {}\n".format(remote_path))
@@ -504,9 +508,9 @@ def impl_download_training_script(local_script_path, bucket_name, region):
 
     :param local_script_path: The path to the training script on disk.
     :param bucket_name: The S3 bucket name.
-    :param region: The region.
+    :param region: The region, or `None` to pull the region from the environment.
     """
-    client = boto3.client("s3", region_name=region)
+    client = make_client("s3", region)
     remote_path = "axon-uploaded-training-scripts/" + os.path.basename(local_script_path)
     client.download_file(bucket_name, remote_path, local_script_path)
     print("Downloaded from: {}\n".format(remote_path))
@@ -518,9 +522,9 @@ def impl_download_dataset(local_dataset_path, bucket_name, region):
 
     :param local_dataset_path: The path to the dataset on disk.
     :param bucket_name: The S3 bucket name.
-    :param region: The region.
+    :param region: The region, or `None` to pull the region from the environment.
     """
-    client = boto3.client("s3", region_name=region)
+    client = make_client("s3", region)
     remote_path = "axon-uploaded-datasets/" + os.path.basename(local_dataset_path)
     client.download_file(bucket_name, remote_path, local_dataset_path)
     print("Downloaded from: {}\n".format(remote_path))