def compute_zone_weights(self, worker_cores: int, local_ssd_data_disk: bool, data_disk_size_gb: int) -> List[ZoneWeight]: weights = [] for r in self._region_info.values(): quota_remaining = { q['metric']: q['limit'] - q['usage'] for q in r['quotas'] } remaining = quota_remaining['PREEMPTIBLE_CPUS'] / worker_cores if local_ssd_data_disk: specific_disk_type_quota = quota_remaining[ 'LOCAL_SSD_TOTAL_GB'] else: specific_disk_type_quota = quota_remaining['SSD_TOTAL_GB'] # FIXME: data_disk_size_gb is assumed to be constant across all instances, but it is # passed as a variable parameter to this function!! remaining = min(remaining, specific_disk_type_quota / data_disk_size_gb) weight = max(remaining / len(r['zones']), 1) for z in r['zones']: zone_name = url_basename(z) weights.append(ZoneWeight(zone_name, weight)) log.info(f'zone_weights {weights}') return weights
def zone_weights(self, worker_cores, worker_local_ssd_data_disk, worker_pd_ssd_data_disk_size_gb): if not self.region_info: return None _zone_weights = [] for r in self.region_info.values(): quota_remaining = { q['metric']: q['limit'] - q['usage'] for q in r['quotas'] } remaining = quota_remaining['PREEMPTIBLE_CPUS'] / worker_cores if worker_local_ssd_data_disk: remaining = min(remaining, quota_remaining['LOCAL_SSD_TOTAL_GB'] / 375) else: remaining = min( remaining, quota_remaining['SSD_TOTAL_GB'] / worker_pd_ssd_data_disk_size_gb) weight = max(remaining / len(r['zones']), 1) for z in r['zones']: zone_name = url_basename(z) _zone_weights.append(ZoneWeight(zone_name, weight)) log.info(f'zone_weights {_zone_weights}') return _zone_weights
async def fetch_region_quotas( compute_client: aiogoogle.GoogleComputeClient, regions: Set[str]) -> Tuple[Dict[str, Dict[str, Any]], List[str]]: region_info = { name: await compute_client.get(f'/regions/{name}') for name in regions } zones = [url_basename(z) for r in region_info.values() for z in r['zones']] return region_info, zones
async def _full_dest(self): if self.dest_type_task: dest_type = await self.dest_type_task else: dest_type = None if (self.treat_dest_as == Transfer.DEST_DIR or (self.treat_dest_as == Transfer.INFER_DEST and dest_type == AsyncFS.DIR)): # We know dest is a dir, but we're copying to # dest/basename(src), and we don't know its type. return url_join(self.dest, url_basename(self.src.rstrip('/'))), None if (self.treat_dest_as == Transfer.DEST_IS_TARGET and self.dest.endswith('/')): dest_type = AsyncFS.DIR return self.dest, dest_type
async def _full_dest(self): dest_type = await self.dest_type_task if (self.treat_dest_as == Transfer.TARGET_DIR or self.dest.endswith('/') or (self.treat_dest_as == Transfer.INFER_TARGET and dest_type == AsyncFS.DIR)): if dest_type is None: raise FileNotFoundError(self.dest) if dest_type == AsyncFS.FILE: raise NotADirectoryError(self.dest) assert dest_type == AsyncFS.DIR # We know dest is a dir, but we're copying to # dest/basename(src), and we don't know its type. return url_join(self.dest, url_basename(self.src.rstrip('/'))), None assert not self.dest.endswith('/') return self.dest, dest_type
async def on_startup(app): db = Database() await db.async_init() app['db'] = db app['client_session'] = httpx.client_session() aiogoogle_credentials = aiogoogle.GoogleCredentials.from_file( '/billing-monitoring-gsa-key/key.json') bigquery_client = aiogoogle.GoogleBigQueryClient( 'broad-ctsa', credentials=aiogoogle_credentials) app['bigquery_client'] = bigquery_client compute_client = aiogoogle.GoogleComputeClient( PROJECT, credentials=aiogoogle_credentials) app['compute_client'] = compute_client query_billing_event = asyncio.Event() app['query_billing_event'] = query_billing_event region_info = { name: await compute_client.get(f'/regions/{name}') for name in BATCH_GCP_REGIONS } zones = [url_basename(z) for r in region_info.values() for z in r['zones']] app['zones'] = zones app['task_manager'] = aiotools.BackgroundTaskManager() app['task_manager'].ensure_future( retry_long_running('polling_loop', polling_loop, app)) app['task_manager'].ensure_future( retry_long_running('query_billing_loop', run_if_changed_idempotent, query_billing_event, query_billing_body, app)) app['task_manager'].ensure_future(periodically_call( 60, monitor_disks, app)) app['task_manager'].ensure_future( periodically_call(60, monitor_instances, app))
def test_url_basename(): assert url_basename('/path/to/file') == 'file' assert url_basename('https://hail.is/path/to/file') == 'file'
async def update_region_quotas(self): self.region_info = {name: await self.compute_client.get(f'/regions/{name}') for name in BATCH_GCP_REGIONS} self.zones = [url_basename(z) for r in self.region_info.values() for z in r['zones']] log.info('updated region quotas')