def _sync_code(host: str) -> None: if ".git" not in os.listdir(): nprint( "No git repo found. Syncing all python and bash files as a fallback.", style="warning", ) nprint("Please consider using git to track the files to sync.", style="warning") subprocess.Popen( f"rsync -avm -e 'ssh -i {CONFIG.instance_key}' " f"--include '*/' --include '*.py' --include '*.ipynb' --include '*.sh' " f"--exclude '*' " f". ubuntu@{host}:/home/ubuntu/project", shell=True, ).communicate() else: output, error = subprocess.Popen("git ls-tree -r HEAD --name-only", stdout=subprocess.PIPE, shell=True).communicate() git_tracked_files = output.decode("utf-8").strip().splitlines() include_files = [ f"--include '{file_name}'" for file_name in git_tracked_files ] include_string = " ".join(include_files) subprocess.Popen( f"rsync -amr -e 'ssh -i {CONFIG.instance_key}' " f"--include '*/' {include_string} --exclude '*' " f". ubuntu@{host}:/home/ubuntu/project", shell=True, ).communicate()
def ls_spot_gpu_prices(dry_run=False) -> None: if dry_run: return ec2 = CONFIG.get_session().client("ec2") string = AwsUtils._format_price_string("InstanceType", "Price ($/hour)", "GPUs", "CPUs", "Mem (Gb)") print() nprint(string, style="bold") for instance_type in AwsUtils._instance_types(): response = ec2.describe_spot_price_history( InstanceTypes=[instance_type], Filters=[{ "Name": "product-description", "Values": ["Linux/UNIX"] }], ) price = float(response["SpotPriceHistory"][0]["SpotPrice"]) num_gpus, gpu_type, mem, cpus = INSTANCE_GPU_MAP[instance_type] string = AwsUtils._format_price_string(instance_type, round(price, 2), f"{num_gpus} x {gpu_type}", cpus, mem) print(string) print()
def _create_policy(client, policy_name, policy_json): try: client.create_policy(PolicyName=policy_name, PolicyDocument=json.dumps(policy_json)) except botocore.exceptions.ClientError as e: if e.response["Error"]["Code"] == "EntityAlreadyExists": nprint(f"Policy {policy_name} already exists. Skipping.", style="warning") else: raise
def _create_group(client, group_name): try: client.create_group(GroupName=group_name) except botocore.exceptions.ClientError as e: if e.response["Error"]["Code"] == "EntityAlreadyExists": nprint( f"User group {group_name} already exists. Skipping.", style="warning", ) else: raise
def decorated(*args, **kwargs): try: CONFIG.assert_required_config_exists(*cases) return func(*args, **kwargs) except AssertionError as e: nprint(e, style="error") sys.exit(1) except FileNotFoundError as e: # Happens when nimbo config file is not found nprint(e, style="error") sys.exit(1)
def decorated(*args, **kwargs): if IS_TEST_ENV: return func(*args, **kwargs) else: try: return func(*args, **kwargs) except botocore.errorfactory.ClientError as e: nprint(e, style="error") sys.exit(1) except ValueError as e: nprint(e, style="error") sys.exit(1) except KeyboardInterrupt: print("Aborting...") sys.exit(1)
def allow_ingress_current_ip(target: str, dry_run=False) -> None: ec2 = CONFIG.get_session().client("ec2") try: response = ec2.describe_security_groups(GroupNames=[target], DryRun=dry_run) security_group_id = response["SecurityGroups"][0]["GroupId"] except botocore.exceptions.ClientError as e: if e.response["Error"]["Code"] == "InvalidGroup.NotFound": nprint( f"Security group {target} not found. Please use an existing" " security group or create a new one in the AWS console.", style="error", ) sys.exit(1) elif e.response["Error"]["Code"] == "UnauthorizedOperation": return else: raise my_public_ip = requests.get( "https://checkip.amazonaws.com").text.strip() try: ec2.authorize_security_group_ingress( GroupId=security_group_id, IpPermissions=[{ "IpProtocol": "tcp", "FromPort": 22, "ToPort": 22, "IpRanges": [{ "CidrIp": f"{my_public_ip}/16" }], }], ) except botocore.exceptions.ClientError as e: if e.response["Error"]["Code"] == "InvalidPermission.Duplicate": return elif e.response["Error"]["Code"] == "UnauthorizedOperation": return else: raise
def ls_active_instances(dry_run=False) -> None: ec2 = CONFIG.get_session().client("ec2") try: response = ec2.describe_instances( Filters=[{ "Name": "instance-state-name", "Values": ["running", "pending"] }] + AwsInstance._make_instance_filters(), DryRun=dry_run, ) for reservation in response["Reservations"]: for inst in reservation["Instances"]: nprint( f"Id: [bright_green]{inst['InstanceId']}[/bright_green]\n" f"Status: {inst['State']['Name']}\n" f"Launch Time: {inst['LaunchTime']}\n" f"InstanceType: {inst['InstanceType']}\n" f"IP Address: {inst['PublicIpAddress']}\n") except botocore.exceptions.ClientError as e: if "DryRunOperation" not in str(e): raise
def mk_bucket(bucket_name: str, dry_run=False) -> None: """Create an S3 bucket in a specified region :param bucket_name: Bucket to create :param dry_run :return: True if bucket created, else False """ try: session = CONFIG.get_session() s3 = session.client("s3") location = {"LocationConstraint": session.region_name} s3.create_bucket(Bucket=bucket_name, CreateBucketConfiguration=location) except botocore.exceptions.ClientError as e: if e.response["Error"]["Code"] == "BucketAlreadyOwnedByYou": nprint("Bucket nimbo-main-bucket already exists.", style="warning") else: nprint(e, style="error") return print("Bucket %s created." % bucket_name)
def _create_role_and_instance_profile(client, role_name): try: client.create_role( RoleName=role_name, AssumeRolePolicyDocument=json.dumps(ASSUME_ROLE_POLICY), ) except botocore.exceptions.ClientError as e: if e.response["Error"]["Code"] == "EntityAlreadyExists": nprint(f"Role {role_name} already exists. Skipping.", style="warning") else: raise try: client.create_instance_profile(InstanceProfileName=role_name, Path="/") except botocore.exceptions.ClientError as e: if e.response["Error"]["Code"] == "EntityAlreadyExists": nprint( f"Instance profile for role {role_name} already exists. Skipping.", style="warning", ) else: raise try: client.add_role_to_instance_profile(InstanceProfileName=role_name, RoleName=role_name) except botocore.exceptions.ClientError as e: if e.response["Error"]["Code"] == "LimitExceeded": nprint( f"Instance profile {role_name} already has a role. Skipping.", style="warning", ) else: raise
def setup(profile: str, full_s3_access=False) -> None: session = boto3.Session(profile_name=profile) account = session.client("sts").get_caller_identity()["Account"] iam = session.client("iam") nprint_header(f"Creating user group {NIMBO_USER_GROUP}...") AwsPermissions._create_group(iam, NIMBO_USER_GROUP) nprint_header(f"Creating policy {EC2_POLICY_NAME}...") AwsPermissions._create_policy(iam, EC2_POLICY_NAME, EC2_POLICY_JSON) nprint_header( f"Attaching policy {EC2_POLICY_NAME} to user group {NIMBO_USER_GROUP}..." ) iam.attach_group_policy( GroupName=NIMBO_USER_GROUP, PolicyArn=f"arn:aws:iam::{account}:policy/{EC2_POLICY_NAME}", ) if full_s3_access: nprint_header(f"Creating role {S3_ACCESS_ROLE_NAME}...") AwsPermissions._create_role_and_instance_profile( iam, S3_ACCESS_ROLE_NAME) nprint_header( f"Attaching AmazonS3FullAccess policy to role {S3_ACCESS_ROLE_NAME}..." ) iam.attach_role_policy( PolicyArn="arn:aws:iam::aws:policy/AmazonS3FullAccess", RoleName=S3_ACCESS_ROLE_NAME, ) nprint_header(f"Creating policy {PASS_ROLE_POLICY_NAME}...") pass_role_policy_json = { "Version": "2012-10-17", "Statement": [{ "Sid": "NimboPassRolePolicy", "Effect": "Allow", "Action": "iam:PassRole", "Resource": f"arn:aws:iam::*:role/{S3_ACCESS_ROLE_NAME}", }], } AwsPermissions._create_policy(iam, PASS_ROLE_POLICY_NAME, pass_role_policy_json) nprint_header(f"Attaching policy {PASS_ROLE_POLICY_NAME}" f" to user group {NIMBO_USER_GROUP}...") iam.attach_group_policy( GroupName=NIMBO_USER_GROUP, PolicyArn= f"arn:aws:iam::{account}:policy/{PASS_ROLE_POLICY_NAME}", ) nprint_header(f"Attaching policy AmazonS3FullAccess" f" to user group {NIMBO_USER_GROUP}...") iam.attach_group_policy( GroupName=NIMBO_USER_GROUP, PolicyArn="arn:aws:iam::aws:policy/AmazonS3FullAccess", ) else: nprint( "\nSince you chose not to give full S3 access to the Nimbo user group" " and instance role,\nwe recommend that you create a role with the" " necessary S3 permissions in the AWS console.\nOnce you do this, give" " the role name to the people using Nimbo so that they can set\n" "the 'role' field in the nimbo-config.yml to this value.", style="warning", ) print() nprint_header("Done.") nprint_header("To add users to the NimboUserGroup, simply" " run 'nimbo add-user USERNAME YOUR_AWS_PROFILE'.\n" "For more info use 'nimbo add-user --help'")
def run(job_cmd: str, dry_run=False) -> Dict[str, str]: if dry_run: return {"message": job_cmd + "_dry_run"} # Launch instance with new volume for anaconda telemetry.record_event("run") start_t = time.monotonic() instance_id = AwsInstance._start_instance() try: # Wait for the instance to be running AwsInstance._block_until_instance_running(instance_id) end_t = time.monotonic() nprint_header( f"Instance running. ({round((end_t - start_t), 2)} s)") nprint_header(f"InstanceId: [green]{instance_id}[/green]") print() time.sleep(5) host = AwsInstance._get_host_from_instance_id(instance_id) AwsInstance._block_until_ssh_ready(host) if job_cmd == "_nimbo_launch": nprint_header( f"Run [cyan]nimbo ssh {instance_id}[/cyan] to log onto the instance" ) return { "message": job_cmd + "_success", "instance_id": instance_id } ssh = ( f"ssh -i {CONFIG.instance_key} -o 'StrictHostKeyChecking no'" " -o ServerAliveInterval=5 ") scp = f"scp -i {CONFIG.instance_key} -o 'StrictHostKeyChecking no'" local_env = "/tmp/local_env.yml" user_conda_yml = CONFIG.conda_env # TODO: Replace this with shutil subprocess.check_output(f"cp {user_conda_yml} {local_env}", shell=True) # Send conda env yaml and setup scripts to instance print() nprint_header(f"Syncing conda, config, and setup files...") AwsInstance._write_nimbo_vars() # Create project folder and send env and config files there subprocess.check_output(f"{ssh} ubuntu@{host} mkdir project", shell=True) subprocess.check_output( f"{scp} {local_env} {CONFIG.nimbo_config_file} {NIMBO_VARS}" f" ubuntu@{host}:/home/ubuntu/project/", shell=True, ) # Sync code with instance print() nprint_header(f"Syncing code...") AwsInstance._sync_code(host) nprint_header(f"Running setup code on the instance from here on.") # Run remote_setup script on instance AwsInstance._run_remote_script(ssh, scp, host, instance_id, job_cmd, "remote_setup.sh") if job_cmd == "_nimbo_notebook": subprocess.Popen( f"{ssh} -o 'ExitOnForwardFailure yes' " f"ubuntu@{host} -NfL 57467:localhost:57467 >/dev/null 2>&1", shell=True, ).communicate() nprint_header( "Make sure to run 'nimbo sync-notebooks <instance_id>' frequently " "to sync the notebook to your local folder, as the remote notebooks" " will be lost once the instance is terminated.") return { "message": job_cmd + "_success", "instance_id": instance_id } except BaseException as e: if (type(e) != KeyboardInterrupt and type(e) != subprocess.CalledProcessError): nprint(e, style="error") if not CONFIG.persist: nprint_header( f"Deleting instance {instance_id} (from local)... ") AwsInstance.delete_instance(instance_id) return {"message": job_cmd + "_error", "instance_id": instance_id}
def run_access_test(dry_run=False) -> None: if dry_run: return CONFIG.instance_type = "t3.medium" CONFIG.run_in_background = False CONFIG.persist = False try: # Send test file to s3 results path and delete it profile = CONFIG.aws_profile region = CONFIG.region_name results_path = CONFIG.s3_results_path subprocess.check_output( "echo 'Hello World' > nimbo-access-test.txt", shell=True) command = AwsStorage.mk_s3_command("cp", "nimbo-access-test.txt", results_path) subprocess.check_output(command, shell=True) command = f"aws s3 ls {results_path} --profile {profile} --region {region}" subprocess.check_output(command, shell=True) command = (f"aws s3 rm {results_path}/nimbo-access-test.txt " f"--profile {profile} --region {region}") subprocess.check_output(command, shell=True) print("You have the necessary S3 read/write " "permissions from your computer \u2713") except subprocess.CalledProcessError as e: nprint(e, style="error") sys.exit(1) # Launch instance with new volume for anaconda print("Launching test instance... ") instance_id = AwsInstance._start_instance() try: # Wait for the instance to be running AwsInstance._block_until_instance_running(instance_id) print(f"Instance running. Instance creation allowed \u2713") print(f"InstanceId: {instance_id}") print() print("Trying to delete this instance...") AwsInstance.delete_instance(instance_id) print("Instance deletion allowed \u2713") print("\nLaunching another instance...") instance_id = AwsInstance._start_instance() print(f"Instance running. InstanceId: {instance_id}") time.sleep(5) host = AwsInstance._get_host_from_instance_id(instance_id) ssh = ( f"ssh -i {CONFIG.instance_key} -o 'StrictHostKeyChecking no' " "-o ServerAliveInterval=20") scp = f"scp -i {CONFIG.instance_key} -o 'StrictHostKeyChecking no'" AwsInstance._block_until_ssh_ready(host) print("Instance key allows ssh access to remote instance \u2713") print("Security group allows ssh access to remote instance \u2713") AwsInstance._write_nimbo_vars() subprocess.check_output( f"{scp} {CONFIG.nimbo_config_file} {NIMBO_VARS} " + f"ubuntu@{host}:/home/ubuntu/", shell=True, ) AwsInstance._run_remote_script(ssh, scp, host, instance_id, "", "remote_s3_test.sh") except BaseException as e: if (type(e) != KeyboardInterrupt and type(e) != subprocess.CalledProcessError): nprint(e, style="error") if not CONFIG.persist: nprint_header( f"Deleting instance {instance_id} (from local)...") AwsInstance.delete_instance(instance_id) sys.exit(1)
def ls_gpu_prices(dry_run=False) -> None: if dry_run: return full_region_name = FULL_REGION_NAMES[CONFIG.region_name] pricing = CONFIG.get_session().client("pricing", region_name="us-east-1") string = AwsUtils._format_price_string("InstanceType", "Price ($/hour)", "GPUs", "CPUs", "Mem (Gb)") print() nprint(string, style="bold") for instance_type in AwsUtils._instance_types(): response = pricing.get_products( ServiceCode="AmazonEC2", MaxResults=100, FormatVersion="aws_v1", Filters=[ { "Type": "TERM_MATCH", "Field": "instanceType", "Value": instance_type, }, { "Type": "TERM_MATCH", "Field": "location", "Value": full_region_name, }, { "Type": "TERM_MATCH", "Field": "operatingSystem", "Value": "Linux", }, { "Type": "TERM_MATCH", "Field": "capacitystatus", "Value": "Used" }, { "Type": "TERM_MATCH", "Field": "preInstalledSw", "Value": "NA" }, { "Type": "TERM_MATCH", "Field": "tenancy", "Value": "shared" }, ], ) inst = json.loads(response["PriceList"][0]) inst = inst["terms"]["OnDemand"] inst = list(inst.values())[0] inst = list(inst["priceDimensions"].values())[0] inst = inst["pricePerUnit"] currency = list(inst.keys())[0] price = float(inst[currency]) num_gpus, gpu_type, mem, cpus = INSTANCE_GPU_MAP[instance_type] string = AwsUtils._format_price_string(instance_type, round(price, 2), f"{num_gpus} x {gpu_type}", cpus, mem) print(string) print()