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')