class FirebaseProvider(KeyValueProvider):

    def __init__(self, db_config, blocking_auth=True):
        self.meta_store = FirebaseStorageHandler(db_config)

        super().__init__(
            db_config,
            self.meta_store,
            blocking_auth)

    def _get(self, key, shallow=False):
        try:
            splitKey = key.split('/')
            key_path = '/'.join(splitKey[:-1])
            key_name = splitKey[-1]
            dbobj = self.app.database().child(key_path).child(key_name)
            return dbobj.get(self.auth.get_token(), shallow=shallow).val() \
                if self.auth else dbobj.get(shallow=shallow).val()
        except Exception as err:
            self.logger.warn(("Getting key {} from a database " +
                              "raised an exception: {}").format(key, err))
            return None

    def _set(self, key, value):
        try:
            splitKey = key.split('/')
            key_path = '/'.join(splitKey[:-1])
            key_name = splitKey[-1]
            dbobj = self.app.database().child(key_path)
            if self.auth:
                dbobj.update({key_name: value}, self.auth.get_token())
            else:
                dbobj.update({key_name: value})
        except Exception as err:
            self.logger.warn(("Putting key {}, value {} into a database " +
                              "raised an exception: {}")
                             .format(key, value, err))

    def _delete(self, key, shallow=True, token=None):
        dbobj = self.app.database().child(key)

        if self.auth:
            dbobj.remove(self.auth.get_token())
        else:
            dbobj.remove() class HTTPProvider(object):
    """Data provider communicating with API server."""

    def __init__(
            self,
            config,
            verbose=10,
            blocking_auth=True,
            compression=None):
        # TODO: implement connection
        self.url = config.get('serverUrl', None)
        self.verbose = get_storage_verbose_level()
        self.logger = logs.get_logger('HTTPProvider')
        self.logger.setLevel(self.verbose)

        self.credentials: Credentials = \
            Credentials.get_credentials(config)

        self.storage_handler = HTTPStorageHandler(
            self.url,
            self.credentials.to_dict() if self.credentials else None,
            compression=compression)

        self.auth = None
        guest = config.get('guest', None)
        if not guest and 'serviceAccount' not in config.keys():
            self.auth = get_auth(
                config.get('authentication', None),
                blocking_auth
            )

        self.compression = compression
        if self.compression is None:
            self.compression = config.get('compression', None)

    def add_experiment(self, experiment, userid=None,
                       compression=None):

        headers = self._get_headers()
        compression = compression if compression else self.compression

        for tag, art in experiment.artifacts.items():
            if not art.is_mutable and art.local_path is not None:
                art.hash = art.get_hash(art.local_path)

        data = {}
        data['experiment'] = experiment.__dict__
        data['compression'] = compression

        def post_request():
            request = requests.post(
                self.url + '/api/add_experiment',
                headers=headers,
                data=json.dumps(data))

            self._raise_detailed_error(request)
            return request

        request = retry(
            post_request,
            no_retries=10,
            sleep_time=10,
            logger=self.logger)

        artifacts = request.json()['artifacts']

        self._update_artifacts(experiment, artifacts)

    def _update_artifacts(self, experiment, artifacts):
        self.logger.debug(str(experiment.artifacts.keys()))
        self.logger.debug(str(artifacts.keys()))

        for tag, art in experiment.artifacts.items():
            target_art = artifacts.get(tag)
            if target_art is not None:
                art.key = target_art.key
                art.remote_path = target_art.remote_path

                if art.local_path is not None:
                    art.upload(art.local_path)

    def delete_experiment(self, experiment):
        if isinstance(experiment, str):
            key = experiment
        else:
            key = experiment.key

        headers = self._get_headers()

        def post_request():
            request = requests.post(self.url + '/api/delete_experiment',
                                    headers=headers,
                                    data=json.dumps({"key": key})
                                    )
            self._raise_detailed_error(request)

        post_request()
        # retry(post_request, sleep_time=10, logger=self.logger)

    def get_experiment(self, experiment, getinfo='True'):
        if isinstance(experiment, str):
            key = experiment
        else:
            key = experiment.key

        headers = self._get_headers()
        try:
            request = requests.post(self.url + '/api/get_experiment',
                                    headers=headers,
                                    data=json.dumps({"key": key})
                                    )

            self._raise_detailed_error(request)
            data = request.json()['experiment']
            return experiment_from_dict(data)
        except BaseException as e:
            check_for_kb_interrupt()
            self.logger.info('error getting experiment {}'.format(key))
            self.logger.info(e)
            return None

    def start_experiment(self, experiment):
        self.checkpoint_experiment(experiment)
        if isinstance(experiment, str):
            key = experiment
        else:
            key = experiment.key

        headers = self._get_headers()

        def post_request():
            request = requests.post(self.url + '/api/start_experiment',
                                    headers=headers,
                                    data=json.dumps({"key": key})
                                    )
            self._raise_detailed_error(request)
        retry(post_request, sleep_time=10, logger=self.logger)

        if not isinstance(experiment, str):
            experiment.time_started = time.time()

    def stop_experiment(self, experiment):
        if isinstance(experiment, str):
            key = experiment
        else:
            key = experiment.key

        headers = self._get_headers()

        def post_request():
            request = requests.post(self.url + '/api/stop_experiment',
                                    headers=headers,
                                    data=json.dumps({"key": key})
                                    )
            self._raise_detailed_error(request)

        retry(post_request, sleep_time=10, logger=self.logger)

    def finish_experiment(self, experiment):
        if isinstance(experiment, str):
            key = experiment
        else:
            key = experiment.key

        headers = self._get_headers()

        def post_request():
            request = requests.post(self.url + '/api/finish_experiment',
                                    headers=headers,
                                    data=json.dumps({"key": key})
                                    )
            self._raise_detailed_error(request)

        retry(post_request, sleep_time=10, logger=self.logger)
        if not isinstance(experiment, str):
            experiment.time_finished = time.time()

    def get_user_experiments(self, user=None, blocking=True):
        headers = self._get_headers()
        user = user if user else self._get_userid()

        response = requests.post(
            self.url + '/api/get_user_experiments',
            headers=headers,
            data=json.dumps({"user": user}))

        self._raise_detailed_error(response)
        data = response.json()['experiments']

        experiments = data

        return experiments

    def get_projects(self):
        headers = self._get_headers()
        response = requests.post(
            self.url + '/api/get_projects',
            headers=headers)

        self._raise_detailed_error(response)
        projects = response.json()['projects']

        return projects

    def get_project_experiments(self, project):
        headers = self._get_headers()
        response = requests.post(
            self.url + '/api/get_project_experiments',
            headers=headers,
            data=json.dumps({"project": project}))

        self._raise_detailed_error(response)
        data = response.json()['experiments']

        experiments = data
        return experiments

    def get_artifacts(self, key):
        return {t: a['url'] for t, a in
                self.get_experiment(key).artifacts.items()}

    def get_artifact(self, artifact,
                     local_path=None, only_newer='True'):

        if isinstance(artifact, str):
            experiment_key = re.match(r'.*(?=/)', artifact).group(0)
            artifact_tag = re.search(r'(?<=/)[^/]*\Z', artifact).group(0)
            experiment = self.get_experiment(experiment_key)
            artifact = experiment.artifacts[artifact_tag]

        return artifact.download(local_path=local_path, only_newer=only_newer)

    def get_users(self):
        headers = self._get_headers()
        response = requests.post(
            self.url + '/api/get_users',
            headers=headers)

        self._raise_detailed_error(response)
        users = response.json()['users']

        return users

    def checkpoint_experiment(self, experiment):
        if isinstance(experiment, str):
            key = experiment
            experiment = self.get_experiment(key)
        else:
            key = experiment.key

        headers = self._get_headers()

        def post_request():
            request = requests.post(self.url + '/api/checkpoint_experiment',
                                    headers=headers,
                                    data=json.dumps({"key": key}))

            self._raise_detailed_error(request)
            artifacts = request.json()['artifacts']
            return artifacts

        artifacts = retry(post_request, sleep_time=10, logger=self.logger)

        self._update_artifacts(experiment, artifacts)
        experiment.time_last_checkpoint = time.time()

    def refresh_auth_token(self, email, refresh_token):
        if self.auth:
            self.auth.refresh_token(email, refresh_token)

    def register_user(self, userid, email):
        pass

    def _get_headers(self):
        headers = {"content-type": "application/json"}
        if self.auth:
            token = self.auth.get_token()
            if token:
                headers["Authorization"] = "Firebase " + token
        return headers

    def _get_userid(self):
        userid = None
        if self.auth:
            userid = self.auth.get_user_id()
        userid = userid if userid else 'guest'
        return userid

    def _raise_detailed_error(self, request):
        if request.status_code != 200:
            raise ValueError(str(request.__dict__))

        data = request.json()
        if 'status' in data.keys():
            if data['status'] == 'ok':
                return

            raise ValueError(data['status'])
        else:
            raise ValueError(json.dumps(data))

    def __enter__(self):
        return self

    def __exit__(self, *args):
        pass class PubsubQueue(object):
    def __init__(self, queue_name, sub_name=None, verbose=10):
        from google.cloud import pubsub

        assert 'GOOGLE_APPLICATION_CREDENTIALS' in os.environ.keys()
        with open(os.environ['GOOGLE_APPLICATION_CREDENTIALS']) as f:
            credentials = json.loads(f.read())

        project_name = credentials['project_id']
        self.logger = logs.get_logger(self.__class__.__name__)
        if verbose is not None:
            self.logger.setLevel(parse_verbosity(verbose))

        self.pubclient = pubsub.PublisherClient()
        self.subclient = pubsub.SubscriberClient()

        self.project = project_name
        self.topic_name = self.pubclient.topic_path(project_name, queue_name)
        self.logger.info("Topic name = {}".format(self.topic_name))
        try:
            self.pubtopic = self.pubclient.get_topic(self.topic_name)
        except BaseException as e:
            check_for_kb_interrupt()
            self.pubtopic = self.pubclient.create_topic(self.topic_name)
            self.logger.info('topic {} created'.format(self.topic_name))

        sub_name = sub_name if sub_name else queue_name + "_sub"
        self.logger.info("Topic name = {}".format(queue_name))
        self.logger.info("Subscription name = {}".format(sub_name))

        self.sub_name = self.subclient.subscription_path(
            project_name, sub_name)
        try:
            self.subclient.get_subscription(self.sub_name)
        except BaseException as e:
            check_for_kb_interrupt()
            self.logger.warn(e)
            self.subclient.create_subscription(self.sub_name, self.topic_name,
                                               ack_deadline_seconds=20)

        self.logger.info('subscription {} created'.format(sub_name))

    def clean(self, timeout=0):
        while True:
            msg = self.dequeue(timeout=timeout)
            if not msg:
                break

    def get_name(self):
        return self.subclient.match_topic_from_topic_name(self.topic_name)

    def has_next(self):
        raise NotImplementedError(
            'Using has_next with distributed queue ' +
            'such as pubsub will bite you in the ass! ' +
            'Use dequeue with timeout instead')

        no_retries = 5
        for i in range(no_retries):
            response = self.subclient.api.pull(
                self.sub_name,
                return_immediately=True, max_messages=1)
            messages = response.received_messages

            retval = any(messages)
            if retval:
                break
            else:
                self.logger.debug('has_next has not found any messages, ' +
                                  'retrying (attempt {})'.format(i))
                time.sleep(5)

        for m in messages:
            self.hold(m.ack_id, 0)

        return retval

    def enqueue(self, data):
        self.logger.debug('Sending message with data \n {}'.format(data))
        data = data.encode('utf-8')
        msg_id = self.pubclient.publish(self.topic_name, data)
        self.logger.debug('Message with id {} published'.format(msg_id))

    def dequeue(self, acknowledge=True, timeout=0):
        wait_step = 1
        for waited in range(0, timeout + wait_step, wait_step):
            response = self.subclient.api.pull(
                self.sub_name,
                return_immediately=True, max_messages=1)
            msgs = response.received_messages
            if any(msgs):
                break
            elif waited == timeout:
                return None
            else:
                self.logger.info(
                    ('No messages found, sleeping for {} ' +
                     ' (total sleep time {})').format(wait_step, waited))
                time.sleep(wait_step)

        retval = msgs[0]

        if acknowledge:
            self.acknowledge(retval.ack_id)
            self.logger.debug("Message {} received and acknowledged"
                              .format(retval.message.message_id))

            return sixdecode(retval.message.data)
        else:
            self.logger.debug(
                "Message {} received, ack_id {}" .format(
                    retval.message.message_id,
                    retval.ack_id))
            return (sixdecode(retval.message.data), retval.ack_id)

    def hold(self, ack_key, delay=5):
        self.logger.debug(
            ("Message acknoledgment deadline is extended by {} " +
             "min for {}").format(
                delay,
                ack_key))
        self.subclient.modify_ack_deadline(
            self.sub_name, [ack_key], int(delay * 60))

    def acknowledge(self, ack_key):
        self.logger.debug("Message with key {} acknowledged".format(ack_key))
        retry(lambda: self.subclient.acknowledge(self.sub_name, [ack_key]),
              sleep_time=10,
              logger=self.logger,
              exception_class=RetryError)

    def delete(self):
        self.logger.debug("Deleting pubsub queue with topic" + self.topic_name)
        try:
            self.pubclient.delete_topic(self.topic_name)
            self.subclient.delete_subscription(self.sub_name)
        except RetryError as e:
            self.logger.error('Queue deletion failed with exception:')
            self.logger.exception(e)

    def shutdown(self, delete_queue=True):
        self.delete() class GCloudWorkerManager(object):
    def __init__(self, zone='us-east1-c',
                 auth_cookie=None, verbose=10, branch=None,
                 user_startup_script=None):
        assert 'GOOGLE_APPLICATION_CREDENTIALS' in os.environ.keys()
        with open(os.environ['GOOGLE_APPLICATION_CREDENTIALS'], 'r') as f:
            credentials_dict = json.loads(f.read())

        self.compute = googleapiclient.discovery.build('compute', 'v1')

        self.startup_script_file = os.path.join(
            os.path.dirname(__file__),
            'scripts/gcloud_worker_startup.sh')

        self.install_studio_script = os.path.join(
            os.path.dirname(__file__),
            'scripts/install_studio.sh')

        self.zone = zone
        self.projectid = credentials_dict['project_id']
        self.logger = logs.get_logger("GCloudWorkerManager")
        self.logger.setLevel(verbose)
        self.auth_cookie = auth_cookie
        self.user_startup_script = user_startup_script
        self.repo_url = git_util.get_my_repo_url()
        self.branch = branch if branch else git_util.get_my_checkout_target()
        self.log_bucket = "studioml-logs"

        if user_startup_script:
            self.logger.warn('User startup script argument is deprecated')

    def start_worker(
            self,
            queue_name,
            resources_needed={},
            blocking=False,
            ssh_keypair=None,
            timeout=300,
            ports=[]):

        if ssh_keypair is not None:
            self.logger.warn('ssh keypairs are not supported ' +
                             'for google workers')

        if resources_needed is None:
            resources_needed = {}

        name = self._generate_instance_name()

        config = self._get_instance_config(
            resources_needed, queue_name, timeout=timeout)

        config['name'] = name

        op = self.compute.instances().insert(
            project=self.projectid,
            zone=self.zone,
            body=config).execute()

        if blocking:
            self._wait_for_operation(op['name'])
            self.logger.debug('worker {} created'.format(name))
            return name
        else:
            return (name, op['name'])

    def start_spot_workers(
            self,
            queue_name,
            bid=None,
            resources_needed={},
            ssh_keypair=None,
            queue_upscaling=True,
            start_workers=1,
            max_workers=100,
            timeout=300,
            ports=[]):

        if resources_needed is None:
            resources_needed = {}

        if queue_upscaling is not False:
            self.logger.warn("autoscaling on the queue is not " +
                             "supported for google workers yet")

        if ssh_keypair is not None:
            self.logger.warn('ssh keypairs are not supported ' +
                             'for google workers')

        if bid is not None:
            self.logger.warn("bidding is not supported for " +
                             "google spot instances")

        template_name = self._generate_template_name()
        group_name = self._generate_group_name()

        config = self._get_instance_config(
            resources_needed, queue_name, timeout=timeout)
        config['scheduling']['preemptible'] = True
        config['machineType'] = config['machineType'].split('/')[-1]
        config['metadata']['items'].append(
            {'key': 'groupname', 'value': group_name})

        template_config = {
            'name': template_name,
            'properties': config
        }

        op = self.compute.instanceTemplates() \
            .insert(project=self.projectid, body=template_config) \
            .execute()

        self._wait_for_operation(op['name'], 'global')

        self.logger.info('instance template {} added'.format(template_name))

        self.compute.instanceGroupManagers() .insert(
            project=self.projectid,
            zone=self.zone,
            body={
                "baseInstanceName": self._generate_instance_name(),
                "instanceTemplate": 'global/instanceTemplates/' +
                template_name,
                "name": group_name,
                "targetSize": start_workers}) .execute()

        self.logger.info('Managed groupd {} created'.format(group_name))

    def _get_instance_config(self, resources_needed, queue_name, timeout=300):
        # image_response = self.compute.images().getFromFamily(
        #    project='studio-ed756', family='studioml').execute()

        image_response = None

        if image_response is None:
            image_response = self.compute.images().getFromFamily(
                project='ubuntu-os-cloud', family='ubuntu-1604-lts').execute()

        source_disk_image = image_response['selfLink']

        # Configure the machine
        machine_type = self._generate_machine_type(resources_needed)
        self.logger.debug('Machine type = {}'.format(machine_type))

        with open(self.startup_script_file, 'r') as f:
            startup_script = f.read()

        with open(self.install_studio_script) as f:
            install_studio_script = f.read()

        startup_script = insert_user_startup_script(
            self.user_startup_script,
            startup_script, self.logger)

        startup_script = startup_script.replace(
            '{install_studio}', install_studio_script)
        startup_script = startup_script.format(
            studioml_branch=self.branch,
            repo_url=self.repo_url,
            log_bucket=self.log_bucket,
            use_gpus=resources_needed.get('gpus', 0)
        )

        self.logger.info('Startup script:')
        self.logger.info(startup_script)

        with open(os.environ['GOOGLE_APPLICATION_CREDENTIALS'], 'r') as f:
            credentials = f.read()

        if self.auth_cookie is not None:
            auth_key = os.path.basename(self.auth_cookie)
            with open(self.auth_cookie, 'r') as f:
                auth_data = f.read()
        else:
            auth_key = None
            auth_data = None

        config = {
            'machineType': machine_type,

            # Specify the boot disk and the image to use as a source.
            'disks': [
                {
                    'boot': True,
                    'autoDelete': True,
                    'initializeParams': {
                        'sourceImage': source_disk_image,
                    }
                }
            ],

            # Specify a network interface with NAT to access the public
            # internet.
            'networkInterfaces': [{
                'network': 'global/networks/default',
                'accessConfigs': [
                    {'type': 'ONE_TO_ONE_NAT', 'name': 'External NAT'}
                ]
            }],

            # Allow the instance to access cloud storage and logs.
            'serviceAccounts': [{
                'email': 'default',
                'scopes': [
                    'https://www.googleapis.com/auth/cloud-platform',
                ]
            }],

            # Metadata is readable from the instance and allows you to
            # pass configuration from deployment scripts to instances.
            'metadata': {
                'items': [{
                    'key': 'startup-script',
                    'value': startup_script
                }, {
                    'key': 'credentials',
                    'value': credentials
                }, {
                    'key': 'queue_name',
                    'value': queue_name
                }, {
                    'key': 'auth_key',
                    'value': auth_key
                }, {
                    'key': 'auth_data',
                    'value': auth_data
                }, {
                    'key': 'timeout',
                    'value': str(timeout)
                }]
            },
            "scheduling": {
                "preemptilble": False
            }
        }

        if 'hdd' in resources_needed.keys():
            config['disks'][0]['initializeParams']['diskSizeGb'] = \
                memstr2int(resources_needed['hdd']) / memstr2int('1Gb')

        if resources_needed['gpus'] > 0:
            gpu_type = "nvidia-tesla-k80"
            config['guestAccelerators'] = [
                {
                    "acceleratorType":
                        "projects/{}/zones/{}/acceleratorTypes/{}"
                        .format(self.projectid, self.zone, gpu_type),
                    "acceleratorCount": resources_needed['gpus']
                }
            ]

            config["scheduling"]['onHostMaintenance'] = "TERMINATE"
            config["automaticRestart"] = True

        return config

    def _stop_worker(self, worker_id, blocking=True):
        op = self.compute.instances().delete(
            project=self.projectid,
            zone=self.zone,
            instance=worker_id).execute()

        if blocking:
            self._wait_for_operation(op['name'])
        else:
            return op['name']

    def _generate_instance_name(self):
        return "worker-" + str(uuid.uuid4())

    def _generate_group_name(self):
        return "group-" + str(uuid.uuid4())

    def _generate_template_name(self):
        return "template-" + str(uuid.uuid4())

    def _generate_machine_type(self, resources_needed={}):
        if not any(resources_needed):
            machine_type = "zones/{}/machineTypes/n1-standard-1".format(
                self.zone)
        else:
            cpus = int(resources_needed['cpus'])
            default_ram_per_cpu = 4096
            ram = default_ram_per_cpu * cpus

            if 'ram' in resources_needed.keys():
                ram = memstr2int(resources_needed['ram']) / memstr2int('1Mb')
                ram = int(math.ceil(ram / 256.0) * 256)

            ram_per_cpu = ram / cpus
            assert 1024 <= ram_per_cpu and ram_per_cpu <= 6192, \
                "RAM per cpu should be between 0.9 and 6.5 Gb"

            machine_type = "zones/{}/machineTypes/custom-{}-{}".format(
                self.zone, cpus, ram)

        return machine_type

    def _wait_for_operation(self, operation, locality='zone'):
        self.logger.debug('Waiting for operation {} to finish...'.
                          format(operation))
        while True:
            if locality == 'zone':
                result = self.compute.zoneOperations().get(
                    project=self.projectid,
                    zone=self.zone,
                    operation=operation).execute()
            elif locality == 'global':
                result = self.compute.globalOperations().get(
                    project=self.projectid,
                    operation=operation).execute()
            else:
                raise ValueError(('Unknown locality {} ' +
                                  'should be zone or global'.format(locality)))

            if result['status'] == 'DONE':
                self.logger.debug("done.")
                if 'error' in result:
                    raise Exception(result['error'])
                return result

            time.sleep(1) class EC2WorkerManager(object):

    def __init__(self, auth_cookie=None, verbose=10, branch=None,
                 user_startup_script=None):
        self.startup_script_file = os.path.join(
            os.path.dirname(__file__),
            'scripts/ec2_worker_startup.sh')

        self.install_studio_script = os.path.join(
            os.path.dirname(__file__),
            'scripts/install_studio.sh')

        self.client = boto3.client('ec2')
        self.asclient = boto3.client('autoscaling')
        self.cwclient = boto3.client('cloudwatch')

        self.region = self.client._client_config.region_name

        self.logger = logs.get_logger('EC2WorkerManager')
        self.logger.setLevel(verbose)
        self.auth_cookie = auth_cookie

        self.prices = self._get_ondemand_prices(_instance_specs.keys())

        self.repo_url = git_util.get_my_repo_url()
        self.branch = branch if branch else git_util.get_my_checkout_target()
        self.user_startup_script = user_startup_script

        if user_startup_script:
            self.logger.warn('User startup script argument is deprecated')

    def _get_image_id(self):
        price_path = os.path.join(
            os.path.dirname(__file__),
            'aws/aws_amis.yaml')
        with open(price_path) as f:
            ami_dict = yaml.load(f.read(), Loader=yaml.FullLoader)

        region = self.client._client_config.region_name
        image_type = 'ubuntu16.04'
        return ami_dict[image_type][region]

    def _get_block_device_mappings(self, resources_needed):
        return [{
            'DeviceName': '/dev/sda1',
            'Ebs': {
                'DeleteOnTermination': True,
                'VolumeSize': int(memstr2int(resources_needed['hdd']) /
                                  memstr2int('1g')),
                'VolumeType': 'standard'
            }
        }]

    def start_worker(
            self,
            queue_name,
            resources_needed={},
            blocking=True,
            ssh_keypair=None,
            timeout=300,
            ports=[]):

        imageid = self._get_image_id()

        name = self._generate_instance_name()

        instance_type = self._select_instance_type(resources_needed)

        startup_script = self._get_startup_script(
            resources_needed, queue_name, timeout=timeout)

        self.logger.info(
            'Starting EC2 instance of type {}'.format(instance_type))
        kwargs = {
            'BlockDeviceMappings':
                self._get_block_device_mappings(resources_needed),
            'ImageId': imageid,
            'InstanceType': instance_type,
            'MaxCount': 1,
            'MinCount': 1,
            'UserData': startup_script,
            'InstanceInitiatedShutdownBehavior': 'terminate',
            'TagSpecifications': [{
                'ResourceType': 'instance',
                'Tags': [
                    {
                        'Key': 'Name',
                        'Value': name
                    },
                    {
                        'Key': 'Team',
                        'Value': 'tools'
                    }]
            }]
        }
        ports = set(ports)
        if ssh_keypair is not None:
            ports.add(22)  # ssh port
        if any(ports):
            groupid = self._get_security_group(ports)
            kwargs['SecurityGroupIds'] = [groupid]
            if ssh_keypair is not None:
                kwargs['KeyName'] = ssh_keypair

        response = self.client.run_instances(**kwargs)
        instance_id = response['Instances'][0]['InstanceId']
        self.logger.info(
            'Starting instance {}'.format(instance_id))

        if blocking:
            while True:
                try:
                    response = self.client.describe_instances(
                        InstanceIds=[instance_id]
                    )
                    instance_data = response['Reservations'][0]['Instances'][0]
                    ip_addr = instance_data.get('PublicIpAddress')
                    if ip_addr:
                        print("ip address: {}".format(ip_addr))
                        return
                except BaseException as e:
                    pass

    def _select_instance_type(self, resources_needed):
        sorted_specs = sorted(_instance_specs.items(),
                              key=lambda x: self.prices[x[0]])
        for instance in sorted_specs:
            if int(
                instance[1]['cpus']) >= int(
                resources_needed['cpus']) and memstr2int(
                instance[1]['ram']) >= memstr2int(
                resources_needed['ram']) and int(
                    instance[1]['gpus']) >= int(
                        resources_needed['gpus']):
                return instance[0]

        raise ValueError('No instances that satisfy requirements {} '
                         'can be found'.format(resources_needed))

    def _get_startup_script(
            self,
            resources_needed,
            queue_name,
            autoscaling_group=None,
            timeout=300):
        if self.auth_cookie is not None:
            auth_key = os.path.basename(self.auth_cookie)
            with open(self.auth_cookie, 'r') as f:
                auth_data = f.read()
        else:
            auth_key = None
            auth_data = None

        credentials = ''

        if 'GOOGLE_APPLICATION_CREDENTIALS' in os.environ.keys():
            with open(os.environ['GOOGLE_APPLICATION_CREDENTIALS'], 'r') as f:
                credentials = f.read()
        else:
            self.logger.info('credentials NOT found')

        with open(self.startup_script_file) as f:
            startup_script = f.read()

        with open(self.install_studio_script) as f:
            install_studio_script = f.read()

        startup_script = startup_script.replace(
            '{install_studio}', install_studio_script)

        startup_script = startup_script.format(
            auth_key=auth_key if auth_key else "",
            queue_name=queue_name,
            auth_data=base64.b64encode(auth_data.encode('utf-8'))
            .decode('utf-8') if auth_data else "",
            google_app_credentials=base64.b64encode(
                credentials.encode('utf-8')).decode('utf-8'),
            aws_access_key=self.client._request_signer._credentials.access_key,
            aws_secret_key=self.client._request_signer._credentials.secret_key,
            autoscaling_group=autoscaling_group if autoscaling_group else "",
            region=self.region,
            use_gpus=0 if resources_needed['gpus'] == 0 else 1,
            timeout=timeout,
            repo_url=self.repo_url,
            studioml_branch=self.branch,
        )

        startup_script = insert_user_startup_script(
            self.user_startup_script,
            startup_script, self.logger)

        self.logger.info('Startup script:')
        self.logger.info(startup_script)

        return startup_script

    def _generate_instance_name(self):
        return 'studioml_worker_' + str(uuid.uuid4())

    def _get_security_group(self, ports):
        ports = sorted([p for p in set(ports)])
        group_name = "studioml_{}".format(
            hashlib.sha256(
                pickle.dumps(sorted(ports))
            ).hexdigest()
        )

        try:
            response = self.client.describe_security_groups(
                GroupNames=[group_name]
            )
            groupid = response['SecurityGroups'][0]['GroupId']

        except botocore.exceptions.ClientError as e:

            self.logger.error("Error creating security group!")
            self.logger.exception(e)
            response = self.client.create_security_group(
                GroupName=group_name,
                Description='opens ports {} in studioml workers'
                            .format(",".join([str(p) for p in ports]))
            )

            groupid = response['GroupId']

            ip_permissions = []
            for port in ports:
                ip_permissions.append({
                    'IpProtocol': 'tcp',
                    'FromPort': int(port),
                    'ToPort': int(port),
                    'IpRanges': [{
                        'CidrIp': '0.0.0.0/0'
                    }]
                })

            response = self.client.authorize_security_group_ingress(
                GroupId=groupid,
                GroupName=group_name,
                IpPermissions=ip_permissions
            )

        return groupid

    def start_spot_workers(
            self,
            queue_name,
            bid_price,
            resources_needed={},
            ssh_keypair=None,
            queue_upscaling=True,
            start_workers=1,
            max_workers=100,
            timeout=300,
            ports=[],
            autoscaling_group_name=None):

        # TODO should be able to put bid price as None,
        # which means price of on-demand instance
        # or maybe specify bid_price as a fraction of on-demand price

        instance_type = self._select_instance_type(resources_needed)

        asg_name = "studioml_autoscaling_" + queue_name or \
                   autoscaling_group_name
        launch_config_name = asg_name + "_launch_config"

        startup_script = self._get_startup_script(
            resources_needed, queue_name, asg_name, timeout=timeout)

        if bid_price.endswith('%'):
            bid_price = str(self.prices[instance_type] *
                            float(bid_price.replace('%', '')) / 100)

        self.logger.info('Price bid for instance type {} : {}'
                         .format(instance_type, bid_price))

        launch_config = {
            "ImageId": self._get_image_id(),
            "UserData": startup_script,
            "InstanceType": instance_type,
            "BlockDeviceMappings": self._get_block_device_mappings(
                resources_needed),
            "InstanceMonitoring": {
                'Enabled': False},
            "SpotPrice": bid_price,
        }

        launch_config_name = 'studioml_launch_config_' + \
            hashlib.sha256(
                json.dumps(launch_config, sort_keys=True).encode('utf-8')
            ).hexdigest()

        if ssh_keypair is not None:
            ports.append(22)
        if any(ports):
            groupid = self._get_security_group(ports)
            launch_config['SecurityGroups'] = [groupid]
            if ssh_keypair is not None:
                launch_config['KeyName'] = ssh_keypair

        try:
            response = self.asclient.create_launch_configuration(
                LaunchConfigurationName=launch_config_name, **launch_config)
            self.logger.debug(
                "create_launch_configuration response:\n {}".format(response))
        except self.asclient.exceptions.AlreadyExistsFault:
            self.logger.debug('Launch config {} already exists'
                              .format(launch_config_name))

        asg_config = {
            "LaunchConfigurationName": launch_config_name,
            "MinSize": 0,
            "MaxSize": max_workers,
            "DesiredCapacity": int(start_workers),
            "LoadBalancerNames": [],
            "AvailabilityZones": [self.region + "a"],
            "TerminationPolicies": ['NewestInstance'],
            "DefaultCooldown": 0,
            "NewInstancesProtectedFromScaleIn": True
        }

        self.logger.debug("Creating auto-scaling group " + asg_name)

        try:
            response = self.asclient.create_auto_scaling_group(
                AutoScalingGroupName=asg_name, **asg_config)
        except self.asclient.exceptions.AlreadyExistsFault:
            self.logger.debug('Autoscaling group {} already exists'
                              .format(asg_name))

        if queue_upscaling:
            max_steps = 20
            stepsize = - (- max_workers // max_steps)

            step_adjustments = [{
                'MetricIntervalLowerBound': i,
                'MetricIntervalUpperBound': i + stepsize,
                'ScalingAdjustment': i
            } for i in range(0, max_workers, stepsize)]

            del step_adjustments[-1]['MetricIntervalUpperBound']

            scaleup_policy_response = self.asclient.put_scaling_policy(
                AutoScalingGroupName=asg_name,
                PolicyName=asg_name + "_scaleup",
                PolicyType='StepScaling',
                AdjustmentType="ChangeInCapacity",
                StepAdjustments=step_adjustments
            )

            self.cwclient.put_metric_alarm(
                AlarmName=asg_name + "_scaleup_alarm",
                MetricName="ApproximateNumberOfMessagesVisible",
                Namespace="AWS/SQS",
                Statistic="Average",
                Period=300,
                Threshold=1,
                ComparisonOperator="GreaterThanOrEqualToThreshold",
                EvaluationPeriods=1,
                AlarmActions=[scaleup_policy_response['PolicyARN']],
                Dimensions=[{
                    'Name': 'QueueName',
                    'Value': queue_name
                }]
            )

    def _get_ondemand_prices(self, instances=_instance_specs.keys()):

        # TODO un-hardcode the us-east as a region
        # so that prices are being read for a correct region

        price_path = os.path.join(
            os.path.dirname(__file__),
            'aws/aws_prices.yaml')
        with open(price_path, 'r') as f:
            data = yaml.load(f.read(), Loader=yaml.FullLoader)

        return {i: data[i] for i in instances}

        price_path = os.path.join(os.path.expanduser('~'), '.studioml',
                                  'awsprices.json')
        offer_file_lock = filelock.FileLock(price_path + '.lock')

        try:
            self.logger.info('Reading AWS prices from cache...')
            with offer_file_lock:
                with open(price_path, 'r') as f:
                    offer_dict = json.load(f)

        except IOError:
            self.logger.info(
                'Getting prices info from AWS (this may take a moment...)')

            with offer_file_lock:
                storage_util.download_file(
                    'https://pricing.us-east-1.amazonaws.com/offers/v1.0/'
                    'aws/AmazonEC2/current/index.json',
                    price_path)

                with open(price_path, 'r') as f:
                    offer_dict = json.load(f)

        self.logger.info('Done!')

        region_name = 'US East (N. Virginia)'

        prices = {}

        for instance_type in instances:
            product_sku = [
                k for k, v in six.iteritems(offer_dict['products'])
                if v['attributes'].get('instanceType') == instance_type and
                v['attributes']['tenancy'] == 'Shared' and
                v['attributes']['operatingSystem'] == 'Linux' and
                v['attributes']['location'] == region_name
            ]

            assert len(product_sku) == 1, \
                'Either no or too many products found for {}!' \
                .format(instance_type)

            prices[instance_type] = float(
                list(
                    list(
                        offer_dict['terms']['OnDemand']
                        [product_sku[0]].values()
                    )[0]['priceDimensions'].values()
                )[0]['pricePerUnit']['USD'])

        return prices def parse_verbosity(verbosity=None):
    if verbosity is None:
        return parse_verbosity('info')

    if verbosity == 'True':
        return parse_verbosity('info')

    logger_levels = {
        'debug': 10,
        'info': 20,
        'warn': 30,
        'error': 40,
        'crit': 50
    }

    if isinstance(verbosity, str) and \
            verbosity in logger_levels.keys():
        return logger_levels[verbosity]
    return int(verbosity) def get_auth(
        config,
        blocking=True,
        verbose=logs.DEBUG):

    global _auth_singleton
    if _auth_singleton is None:
        if isinstance(config, dict):
            auth_class = get_auth_class(config['type'])
        else:
            auth_class = get_auth_class(config)

        if auth_class:
            _auth_singleton = auth_class(
                config,
                blocking=blocking,
                verbose=verbose
            )
        else:
            return None
    return _auth_singleton def get_config(config_file=None):
                def replace_with_env(config):
def get_artifact_store(config) -> StorageHandler:
def get_db_provider(config=None, blocking_auth=True): def get_queue(
        queue_name=None,
        cloud=None,
        config=None,
        logger=None,
        close_after=None,
        verbose=10):
def shutdown_queue(queue, logger=None, delete_queue=True): def setup_storage(db_provider, artifact_store):
    global _storage_setup
    _storage_setup = { DB_KEY: db_provider, STORE_KEY: artifact_store } def get_storage_db_provider():
    global _storage_setup
    if _storage_setup is None:
        return None
    return _storage_setup.get(DB_KEY, None) def reset_storage():
    global _storage_setup
    _storage_setup = None def set_storage_verbose_level(level: int):
    global _storage_verbose_level
    _storage_verbose_level = level DEBUG = logging.DEBUG
INFO = logging.INFO
ERROR = logging.ERROR
def get_logger(name):
def debug(line):
def error(line):
def info(line):
import uuid
from studio.firebase_provider import FirebaseProvider
from studio.http_provider import HTTPProvider
from studio.pubsub_queue import PubsubQueue
from studio.gcloud_worker import GCloudWorkerManager
from studio.ec2cloud_worker import EC2WorkerManager
from studio.util.util import parse_verbosity
from studio.auth import get_auth
from studio.db_providers import db_provider_setup
from studio.queues import queues_setup
from studio.storage.storage_setup import setup_storage, get_storage_db_provider,\
    reset_storage, set_storage_verbose_level
from studio.util import logs
"""Data providers."""



def reset_storage_providers():
    reset_storage()


def get_config(config_file=None):
    return db_provider_setup.get_config(config_file=config_file)


def get_db_provider(config=None, blocking_auth=True):

    db_provider = get_storage_db_provider()
    if db_provider is not None:
        return db_provider

    if config is None:
        config = get_config()
    verbose = parse_verbosity(config.get('verbose', None))

    # Save this verbosity level as global for the whole experiment job:
    set_storage_verbose_level(verbose)

    logger = logs.get_logger("get_db_provider")
    logger.setLevel(verbose)
    logger.debug('Choosing db provider with config:')
    logger.debug(config)

    if 'storage' in config.keys():
        artifact_store = db_provider_setup.get_artifact_store(config['storage'])
    else:
        artifact_store = None

    assert 'database' in config.keys()
    db_config = config['database']
    if db_config['type'].lower() == 'firebase':
        db_provider = FirebaseProvider(db_config,
                            blocking_auth=blocking_auth)

    elif db_config['type'].lower() == 'http':
        db_provider = HTTPProvider(db_config,
                            verbose=verbose,
                            blocking_auth=blocking_auth)
    else:
        db_provider = db_provider_setup.get_db_provider(
            config=config, blocking_auth=blocking_auth)

    setup_storage(db_provider, artifact_store)
    return db_provider

def get_queue(
        queue_name=None,
        cloud=None,
        config=None,
        logger=None,
        close_after=None,
        verbose=10):
    queue = queues_setup.get_queue(queue_name=queue_name,
                                   cloud=cloud,
                                   config=config,
                                   logger=logger,
                                   close_after=close_after,
                                   verbose=verbose)
    if queue is None:
        queue = PubsubQueue(queue_name, verbose=verbose)
    return queue

def shutdown_queue(queue, logger=None, delete_queue=True):

    queues_setup.shutdown_queue(queue, logger=logger, delete_queue=delete_queue)