<s> import boto . ec2 <EOL> import boto . exception <EOL> import datetime <EOL> import itertools <EOL> import os <EOL> import time <EOL> from pyfora import __version__ as pyfora_version <EOL> from pyfora import TimeoutError <EOL> def timestamp ( ) : <EOL> return datetime . datetime . now ( ) . strftime ( '<STR_LIT>' ) <EOL> def get_aws_access_key_id ( ) : <EOL> return os . getenv ( '<STR_LIT>' ) <EOL> def get_aws_secret_access_key ( ) : <EOL> return os . getenv ( '<STR_LIT>' ) <EOL> def get_aws_credentials_docker_env ( ) : <EOL> key = get_aws_access_key_id ( ) <EOL> secret = get_aws_secret_access_key ( ) <EOL> if key and secret : <EOL> return '<STR_LIT>' % ( key , secret ) <EOL> return '<STR_LIT>' <EOL> user_data_file = '''<STR_LIT>''' <EOL> class Launcher ( object ) : <EOL> ufora_ssh_security_group_name = '<STR_LIT>' <EOL> ufora_open_security_group_name = '<STR_LIT>' <EOL> ufora_security_group_description = '<STR_LIT>' <EOL> ubuntu_images_paravirtual_ssd = { <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : '<STR_LIT>' <EOL> } <EOL> ubuntu_images_hvm_ssd = { <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : '<STR_LIT>' <EOL> } <EOL> def getUserDataTemplate ( self ) : <EOL> user_data_template = ( <EOL> "<STR_LIT>" + <EOL> "<STR_LIT>" + <EOL> user_data_file + "<STR_LIT>" <EOL> ) <EOL> user_data_template += '<STR_LIT>' <EOL> user_data_template += "<STR_LIT>" <EOL> user_data_template += "<STR_LIT>" <EOL> return user_data_template <EOL> def __init__ ( self , <EOL> region , <EOL> vpc_id = None , <EOL> subnet_id = None , <EOL> security_group_id = None , <EOL> instance_type = None , <EOL> open_public_port = False , <EOL> commit_to_build = None ) : <EOL> assert vpc_id is None or subnet_id is not None <EOL> self . region = region <EOL> self . vpc_id = vpc_id <EOL> self . subnet_id = subnet_id <EOL> self . security_group_id = security_group_id <EOL> self . instance_type = instance_type <EOL> self . open_public_port = open_public_port <EOL> self . ec2 = None <EOL> self . commit_to_build = commit_to_build <EOL> def launch_manager ( self , ssh_key_name , spot_request_bid_price = None , callback = None ) : <EOL> if not self . connected : <EOL> self . connect ( ) <EOL> user_data = self . user_data_for_manager ( ) <EOL> instances = self . launch_instances ( <NUM_LIT:1> , <EOL> ssh_key_name , <EOL> user_data , <EOL> spot_request_bid_price , <EOL> callback = callback ) <EOL> assert len ( instances ) == <NUM_LIT:1> <EOL> self . wait_for_instances ( instances , timeout = <NUM_LIT> , callback = callback ) <EOL> instance = instances [ <NUM_LIT:0> ] <EOL> if instance . state == '<STR_LIT>' : <EOL> self . tag_instance ( instance , "<STR_LIT>" ) <EOL> else : <EOL> raise Exception ( "<STR_LIT>" + instance . id ) <EOL> return instance <EOL> def launch_workers ( self , <EOL> count , <EOL> ssh_key_name , <EOL> manager_instance_id , <EOL> spot_request_bid_price = None , <EOL> callback = None ) : <EOL> if not self . connected : <EOL> self . connect ( ) <EOL> instances = self . launch_instances ( count , <EOL> ssh_key_name , <EOL> self . user_data_for_worker ( manager_instance_id ) , <EOL> spot_request_bid_price , <EOL> callback = callback ) <EOL> self . wait_for_instances ( instances , callback = callback ) <EOL> for instance in instances : <EOL> self . tag_instance ( instance , '<STR_LIT>' ) <EOL> return instances <EOL> def wait_for_services ( self , instances , callback = None ) : <EOL> assert self . connected <EOL> pending = set ( instances ) <EOL> ready = [ ] <EOL> failed = [ ] <EOL> while len ( pending ) : <EOL> toRemove = [ ] <EOL> for i in pending : <EOL> i . update ( ) <EOL> if '<STR_LIT:status>' not in i . tags : <EOL> continue <EOL> if i . tags [ '<STR_LIT:status>' ] . endswith ( '<STR_LIT>' ) : <EOL> toRemove . append ( i ) <EOL> failed . append ( i ) <EOL> elif i . tags [ '<STR_LIT:status>' ] == '<STR_LIT>' : <EOL> toRemove . append ( i ) <EOL> ready . append ( i ) <EOL> for i in toRemove : <EOL> pending . remove ( i ) <EOL> status = { } <EOL> for i in itertools . chain ( pending , ready , failed ) : <EOL> if '<STR_LIT:status>' not in i . tags : <EOL> status_name = '<STR_LIT>' <EOL> else : <EOL> status_name = i . tags [ '<STR_LIT:status>' ] <EOL> if status_name not in status : <EOL> status [ status_name ] = [ ] <EOL> status [ status_name ] . append ( i ) <EOL> if callback : <EOL> callback ( status ) <EOL> if len ( pending ) : <EOL> time . sleep ( <NUM_LIT:1> ) <EOL> return len ( failed ) == <NUM_LIT:0> <EOL> def get_reservations ( self ) : <EOL> if not self . connected : <EOL> self . connect ( ) <EOL> filters = { '<STR_LIT>' : '<STR_LIT>' } <EOL> reservations = self . ec2 . get_all_reservations ( filters = filters ) <EOL> return reservations <EOL> def launch_instances ( self , <EOL> count , <EOL> ssh_key_name , <EOL> user_data , <EOL> spot_request_bid_price = None , <EOL> callback = None ) : <EOL> assert self . instance_type is not None <EOL> assert self . connected <EOL> network_interfaces = self . create_network_interfaces ( ) <EOL> classic_security_groups = None if self . vpc_id else [ self . get_security_group_name ( ) ] <EOL> if self . instance_type [ : <NUM_LIT:2> ] in ( "<STR_LIT>" , "<STR_LIT>" ) : <EOL> image_id = self . ubuntu_images_paravirtual_ssd [ self . region ] <EOL> else : <EOL> image_id = self . ubuntu_images_hvm_ssd [ self . region ] <EOL> dev_sda1 = boto . ec2 . blockdevicemapping . BlockDeviceType ( ) <EOL> dev_sda1 . size = <NUM_LIT:15> <EOL> bdm = boto . ec2 . blockdevicemapping . BlockDeviceMapping ( ) <EOL> bdm [ '<STR_LIT>' ] = dev_sda1 <EOL> request_args = { <EOL> '<STR_LIT>' : image_id , <EOL> '<STR_LIT>' : self . instance_type , <EOL> '<STR_LIT>' : ssh_key_name , <EOL> '<STR_LIT>' : user_data , <EOL> '<STR_LIT>' : classic_security_groups , <EOL> '<STR_LIT>' : network_interfaces , <EOL> '<STR_LIT>' : bdm <EOL> } <EOL> if spot_request_bid_price : <EOL> ts = timestamp ( ) <EOL> spot_requests = self . ec2 . request_spot_instances ( <EOL> price = spot_request_bid_price , <EOL> count = count , <EOL> availability_zone_group = ts , <EOL> launch_group = ts , <EOL> ** request_args ) <EOL> fulfilled , unfulfilled = self . wait_for_spot_instance_requests ( spot_requests , <EOL> callback = callback ) <EOL> assert len ( fulfilled ) == count or len ( unfulfilled ) == count <EOL> if len ( fulfilled ) == <NUM_LIT:0> : <EOL> return [ ] <EOL> return self . ec2 . get_only_instances ( instance_ids = [ r . instance_id for r in fulfilled ] ) <EOL> else : <EOL> reservation = self . ec2 . run_instances ( <EOL> min_count = count , <EOL> max_count = count , <EOL> ** request_args <EOL> ) <EOL> return reservation . instances <EOL> def get_format_args ( self , updates ) : <EOL> assert self . connected <EOL> args = { <EOL> '<STR_LIT>' : self . ec2 . provider . access_key , <EOL> '<STR_LIT>' : self . ec2 . provider . secret_key , <EOL> '<STR_LIT>' : self . region , <EOL> '<STR_LIT>' : get_aws_credentials_docker_env ( ) , <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : pyfora_version , <EOL> '<STR_LIT>' : str ( self . commit_to_build ) if self . commit_to_build is not None else "<STR_LIT>" , <EOL> "<STR_LIT>" : "<STR_LIT:1>" if self . instance_type [ : <NUM_LIT:2> ] == "<STR_LIT>" else "<STR_LIT>" <EOL> } <EOL> args . update ( updates ) <EOL> return args <EOL> def user_data_for_manager ( self ) : <EOL> format_args = self . get_format_args ( { <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> } ) <EOL> return self . getUserDataTemplate ( ) . format ( ** format_args ) <EOL> def user_data_for_worker ( self , manager_instance_id ) : <EOL> manager_address = self . get_instance_internal_ip ( manager_instance_id ) <EOL> format_args = self . get_format_args ( { <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : '<STR_LIT>' + manager_address , <EOL> '<STR_LIT>' : '<STR_LIT>' <EOL> } ) <EOL> return self . getUserDataTemplate ( ) . format ( ** format_args ) <EOL> @ staticmethod <EOL> def wait_for_instances ( instances , timeout = None , callback = None ) : <EOL> t0 = time . time ( ) <EOL> while any ( i . state == '<STR_LIT>' for i in instances ) : <EOL> if timeout and time . time ( ) > t0 + timeout : <EOL> raise TimeoutError ( <EOL> "<STR_LIT>" + [ i . id for i in instances ] <EOL> ) <EOL> time . sleep ( <NUM_LIT:1> ) <EOL> status = { } <EOL> for i in instances : <EOL> if not i . state in status : <EOL> status [ i . state ] = [ ] <EOL> status [ i . state ] . append ( i . id ) <EOL> if i . state == '<STR_LIT>' : <EOL> i . update ( ) <EOL> if callback : <EOL> callback ( status ) <EOL> return True <EOL> @ staticmethod <EOL> def tag_instance ( instance , tag_prefix ) : <EOL> instance . add_tag ( '<STR_LIT:Name>' , "<STR_LIT>" % ( tag_prefix , timestamp ( ) ) ) <EOL> @ property <EOL> def connected ( self ) : <EOL> return self . ec2 is not None <EOL> def connect ( self ) : <EOL> if self . region is None : <EOL> self . ec2 = boto . connect_ec2 ( ) <EOL> else : <EOL> self . ec2 = boto . ec2 . connect_to_region ( self . region ) <EOL> assert self . ec2 is not None <EOL> if self . security_group_id is None : <EOL> self . security_group_id = self . find_or_create_ufora_security_group ( ) <EOL> def create_network_interfaces ( self ) : <EOL> if self . vpc_id is None : <EOL> return None <EOL> interface = boto . ec2 . networkinterface . NetworkInterfaceSpecification ( <EOL> subnet_id = self . subnet_id , <EOL> groups = [ self . security_group_id ] , <EOL> associate_public_ip_address = True <EOL> ) <EOL> return boto . ec2 . networkinterface . NetworkInterfaceCollection ( interface ) <EOL> def get_security_group_name ( self ) : <EOL> groups = self . ec2 . get_all_security_groups ( group_ids = [ self . security_group_id ] ) <EOL> assert len ( groups ) == <NUM_LIT:1> <EOL> return groups [ <NUM_LIT:0> ] . name <EOL> def wait_for_spot_instance_requests ( self , requests , timeout = None , callback = None ) : <EOL> t0 = time . time ( ) <EOL> pending_states = [ '<STR_LIT>' , '<STR_LIT>' ] <EOL> pending = { r . id : r for r in requests } <EOL> fulfilled = { } <EOL> unfulfilled = { } <EOL> while len ( pending ) > <NUM_LIT:0> : <EOL> if timeout and time . time ( ) > t0 + timeout : <EOL> raise TimeoutError ( <EOL> "<STR_LIT>" , <EOL> fulfilled = fulfilled , <EOL> unfulfilled = unfulfilled , <EOL> pending = pending <EOL> ) <EOL> time . sleep ( <NUM_LIT:1> ) <EOL> spot_requests = self . ec2 . get_all_spot_instance_requests ( <EOL> request_ids = pending . keys ( ) <EOL> ) <EOL> status = { } <EOL> for request in spot_requests : <EOL> status_key = '<STR_LIT>' % ( request . state , request . status . code ) <EOL> if status_key not in status : <EOL> status [ status_key ] = [ ] <EOL> status [ status_key ] . append ( request . id ) <EOL> if request . state == '<STR_LIT>' : <EOL> fulfilled [ request . id ] = request <EOL> del pending [ request . id ] <EOL> elif request . state != '<STR_LIT>' or request . status . code not in pending_states : <EOL> unfulfilled [ request . id ] = request <EOL> del pending [ request . id ] <EOL> if callback : <EOL> callback ( status ) <EOL> return ( fulfilled . values ( ) , unfulfilled . values ( ) ) <EOL> def get_instance_internal_ip ( self , instance_id ) : <EOL> reservations = self . ec2 . get_all_instances ( instance_ids = [ instance_id ] ) <EOL> assert len ( reservations ) == <NUM_LIT:1> <EOL> assert len ( reservations [ <NUM_LIT:0> ] . instances ) == <NUM_LIT:1> <EOL> instance = reservations [ <NUM_LIT:0> ] . instances [ <NUM_LIT:0> ] <EOL> return instance . private_ip_address <EOL> def find_or_create_ufora_security_group ( self ) : <EOL> security_group = self . find_ufora_security_group ( ) <EOL> if security_group is None : <EOL> security_group = self . create_ufora_security_group ( ) <EOL> return security_group . id <EOL> def find_ufora_security_group ( self ) : <EOL> filters = None <EOL> if self . vpc_id is not None : <EOL> filters = { '<STR_LIT>' : self . vpc_id } <EOL> try : <EOL> groups = self . ec2 . get_all_security_groups ( <EOL> groupnames = [ self . security_group_name ] , <EOL> filters = filters <EOL> ) <EOL> except boto . exception . EC2ResponseError as e : <EOL> if e . error_code == '<STR_LIT>' : <EOL> return None <EOL> raise <EOL> assert len ( groups ) < <NUM_LIT:2> , "<STR_LIT>" <EOL> if len ( groups ) == <NUM_LIT:0> : <EOL> return None <EOL> return groups [ <NUM_LIT:0> ] <EOL> @ property <EOL> def security_group_name ( self ) : <EOL> return self . ufora_open_security_group_name if self . open_public_port else self . ufora_ssh_security_group_name <EOL> def create_ufora_security_group ( self ) : <EOL> security_group = self . ec2 . create_security_group ( <EOL> name = self . security_group_name , <EOL> description = self . ufora_security_group_description , <EOL> vpc_id = self . vpc_id <EOL> ) <EOL> security_group . authorize ( src_group = security_group , <EOL> ip_protocol = '<STR_LIT>' , <EOL> from_port = <NUM_LIT> , <EOL> to_port = <NUM_LIT> ) <EOL> security_group . authorize ( cidr_ip = '<STR_LIT>' , <EOL> ip_protocol = '<STR_LIT>' , <EOL> from_port = <NUM_LIT> , <EOL> to_port = <NUM_LIT> ) <EOL> if self . open_public_port : <EOL> security_group . authorize ( cidr_ip = '<STR_LIT>' , <EOL> ip_protocol = '<STR_LIT>' , <EOL> from_port = <NUM_LIT> , <EOL> to_port = <NUM_LIT> ) <EOL> return security_group </s>
