async def test_update_commits():
deploy_config = get_deploy_config()
headers = service_auth_headers(deploy_config, 'benchmark')
commit_benchmark_url = deploy_config.url(
'benchmark', f'/api/v1alpha/benchmark/commit/{sha}')
async def request(method):
return await utils.request_retry_transient_errors(
session,
method,
f'{commit_benchmark_url}',
headers=headers,
json={'sha': sha})
async with client_session() as session:
await request('DELETE')
resp = await request('GET')
commit = await resp.json()
assert commit['status'] is None, commit
resp = await request('POST')
commit = await resp.json()
while commit['status'] is not None and not commit['status']['complete']:
await asyncio.sleep(5)
resp = await request('GET')
commit = await resp.json()
print(commit['status'])
def test_batch_create_validation(self):
bad_configs = [
# unexpected field fleep
{'billing_project': 'foo', 'n_jobs': 5, 'token': 'baz', 'fleep': 'quam'},
# billing project None/missing
{'billing_project': None, 'n_jobs': 5, 'token': 'baz'},
{'n_jobs': 5, 'token': 'baz'},
# n_jobs None/missing
{'billing_project': 'foo', 'n_jobs': None, 'token': 'baz'},
{'billing_project': 'foo', 'token': 'baz'},
# n_jobs wrong type
{'billing_project': 'foo', 'n_jobs': '5', 'token': 'baz'},
# token None/missing
{'billing_project': 'foo', 'n_jobs': 5, 'token': None},
{'billing_project': 'foo', 'n_jobs': 5},
# attribute key/value None
{'attributes': {'k': None}, 'billing_project': 'foo', 'n_jobs': 5, 'token': 'baz'},
]
url = deploy_config.url('batch', '/api/v1alpha/batches/create')
headers = service_auth_headers(deploy_config, 'batch')
for config in bad_configs:
r = retry_response_returning_functions(
requests.post,
url,
json=config,
allow_redirects=True,
headers=headers)
assert r.status_code == 400, (config, r)
async def test_deploy():
deploy_config = get_deploy_config()
ci_deploy_status_url = deploy_config.url('ci',
'/api/v1alpha/deploy_status')
headers = service_auth_headers(deploy_config, 'ci')
async with in_cluster_ssl_client_session(
raise_for_status=True,
timeout=aiohttp.ClientTimeout(total=60)) as session:
async def wait_forever():
deploy_state = None
failure_information = None
while deploy_state is None:
resp = await utils.request_retry_transient_errors(
session, 'GET', f'{ci_deploy_status_url}', headers=headers)
deploy_statuses = await resp.json()
log.info(
f'deploy_statuses:\n{json.dumps(deploy_statuses, indent=2)}'
)
assert len(deploy_statuses) == 1, deploy_statuses
deploy_status = deploy_statuses[0]
deploy_state = deploy_status['deploy_state']
failure_information = deploy_status.get('failure_information')
await asyncio.sleep(5)
log.info(f'returning {deploy_status} {failure_information}')
return deploy_state, failure_information
deploy_state, failure_information = await asyncio.wait_for(
wait_forever(), timeout=30 * 60)
assert deploy_state == 'success', str(failure_information)
async def test_connect_to_address_on_pod_ip():
ssl_config = _get_ssl_config()
client_ssl_context = ssl.create_default_context(
purpose=ssl.Purpose.SERVER_AUTH, cafile=ssl_config['outgoing_trust'])
client_ssl_context.load_default_certs()
client_ssl_context.load_cert_chain(ssl_config['cert'],
keyfile=ssl_config['key'],
password=None)
client_ssl_context.verify_mode = ssl.CERT_REQUIRED
client_ssl_context.check_hostname = False
async with aiohttp.ClientSession(
connector=aiohttp.TCPConnector(ssl=client_ssl_context),
raise_for_status=True,
timeout=aiohttp.ClientTimeout(total=5),
headers=service_auth_headers(deploy_config, 'address'),
) as session:
async def get():
address, port = await deploy_config.address('address')
session.get(
f'https://{address}:{port}{deploy_config.base_path("address")}/api/address'
)
await retry_transient_errors(get)
def __init__(self,
billing_project: str,
deploy_config: Optional[DeployConfig] = None,
session: Optional[aiohttp.ClientSession] = None,
headers: Optional[Dict[str, str]] = None,
_token: Optional[str] = None,
token_file: Optional[str] = None):
self.billing_project = billing_project
if not deploy_config:
deploy_config = get_deploy_config()
self.url = deploy_config.base_url('batch')
if session is None:
session = client_session()
self._session = session
h: Dict[str, str] = {}
if headers:
h.update(headers)
if _token:
h['Authorization'] = f'Bearer {_token}'
else:
h.update(
service_auth_headers(deploy_config,
'batch',
token_file=token_file))
self._headers = h
async def __init__(self,
billing_project,
deploy_config=None,
session=None,
headers=None,
_token=None):
self.billing_project = billing_project
if not deploy_config:
deploy_config = get_deploy_config()
self.url = deploy_config.base_url('batch')
if session is None:
session = aiohttp.ClientSession(
raise_for_status=True, timeout=aiohttp.ClientTimeout(total=60))
self._session = session
userinfo = await async_get_userinfo(deploy_config)
self.bucket = userinfo['bucket_name']
h = {}
if headers:
h.update(headers)
if _token:
h['Authorization'] = f'Bearer {_token}'
else:
h.update(service_auth_headers(deploy_config, 'batch'))
self._headers = h
async def __aenter__(self):
headers = service_auth_headers(self._deploy_config, 'ci')
self._session = ssl_client_session(
raise_for_status=True,
timeout=aiohttp.ClientTimeout(total=60),
headers=headers)
return self
def __init__(self,
billing_project: str = None,
bucket: str = None,
*,
deploy_config=None,
skip_logging_configuration: bool = False):
if billing_project is None:
billing_project = get_user_config().get('batch',
'billing_project',
fallback=None)
if billing_project is None:
billing_project = os.environ.get('HAIL_BILLING_PROJECT')
if billing_project is None:
raise ValueError(
"No billing project. Call 'init_service' with the billing "
"project, set the HAIL_BILLING_PROJECT environment variable, "
"or run 'hailctl config set batch/billing_project "
"MY_BILLING_PROJECT'")
self._billing_project = billing_project
if bucket is None:
bucket = get_user_config().get('batch', 'bucket', fallback=None)
if bucket is None:
raise ValueError(
f'the bucket parameter of ServiceBackend must be set '
f'or run `hailctl config set batch/bucket '
f'MY_BUCKET`')
self._bucket = bucket
if not deploy_config:
deploy_config = get_deploy_config()
self.url = deploy_config.base_url('query')
self.headers = service_auth_headers(deploy_config, 'query')
self._fs = None
self._logger = PythonOnlyLogger(skip_logging_configuration)
def __init__(self, deploy_config=None, skip_logging_configuration=False):
if not deploy_config:
deploy_config = get_deploy_config()
self.url = deploy_config.base_url('query')
self.headers = service_auth_headers(deploy_config, 'query')
self._fs = None
self._logger = PythonOnlyLogger(skip_logging_configuration)
async def __init__(self,
billing_project,
deploy_config=None,
session=None,
headers=None,
_token=None,
token_file=None):
self.billing_project = billing_project
if not deploy_config:
deploy_config = get_deploy_config()
self.url = deploy_config.base_url('batch')
if session is None:
session = get_context_specific_ssl_client_session(
raise_for_status=True, timeout=aiohttp.ClientTimeout(total=60))
self._session = session
h = {}
if headers:
h.update(headers)
if _token:
h['Authorization'] = f'Bearer {_token}'
else:
h.update(
service_auth_headers(deploy_config,
'batch',
token_file=token_file))
self._headers = h
def __init__(self, deploy_config=None):
from hailtop.config import get_deploy_config
from hailtop.auth import service_auth_headers
if not deploy_config:
deploy_config = get_deploy_config()
self.url = deploy_config.base_url('apiserver')
self.headers = service_auth_headers(deploy_config, 'apiserver')
self._fs = None
async def test_connect_to_address_on_pod_ip():
async with in_cluster_no_hostname_checks_ssl_client_connection(
raise_for_status=True,
timeout=aiohttp.ClientTimeout(total=60),
headers=service_auth_headers(deploy_config, 'address')) as session:
address, port = await deploy_config.address('address')
await request_retry_transient_errors(
session,
'GET',
f'https://{address}:{port}{deploy_config.base_path("address")}/api/address')
async def test_invariants():
deploy_config = get_deploy_config()
url = deploy_config.url('batch-driver', '/check_invariants')
headers = service_auth_headers(deploy_config, 'batch-driver')
async with client_session() as session:
resp = await utils.request_retry_transient_errors(session, 'GET', url, headers=headers)
data = await resp.json()
assert data['check_incremental_error'] is None, data
assert data['check_resource_aggregation_error'] is None, data
async def test_billing_monitoring():
deploy_config = get_deploy_config()
monitoring_deploy_config_url = deploy_config.url('monitoring', '/api/v1alpha/billing')
headers = service_auth_headers(deploy_config, 'monitoring')
async with client_session() as session:
async def wait_forever():
data = None
while data is None:
resp = await utils.request_retry_transient_errors(
session, 'GET', f'{monitoring_deploy_config_url}', headers=headers)
data = await resp.json()
await asyncio.sleep(5)
return data
data = await asyncio.wait_for(wait_forever(), timeout=30 * 60)
assert data['cost_by_service'] is not None, data
async def async_main():
deploy_config = get_deploy_config()
auth_ns = deploy_config.service_ns('auth')
tokens = get_tokens()
if auth_ns not in tokens:
print('Not logged in.')
return
headers = service_auth_headers(deploy_config, 'auth')
async with client_session(headers=headers) as session:
async with session.post(deploy_config.url('auth', '/api/v1alpha/logout')):
pass
auth_ns = deploy_config.service_ns('auth')
del tokens[auth_ns]
tokens.write()
print('Logged out.')
async def create(billing_project: str,
deploy_config: Optional[DeployConfig] = None,
session: Optional[httpx.ClientSession] = None,
headers: Optional[Dict[str, str]] = None,
_token: Optional[str] = None,
token_file: Optional[str] = None):
if not deploy_config:
deploy_config = get_deploy_config()
url = deploy_config.base_url('batch')
if session is None:
session = httpx.client_session()
if headers is None:
headers = dict()
if _token:
headers['Authorization'] = f'Bearer {_token}'
else:
headers.update(service_auth_headers(deploy_config, 'batch', token_file=token_file))
return BatchClient(
billing_project=billing_project,
url=url,
session=session,
headers=headers)
async def async_main():
deploy_config = get_deploy_config()
auth_ns = deploy_config.service_ns('auth')
tokens = get_tokens()
if auth_ns not in tokens:
print('Not logged in.')
return
headers = service_auth_headers(deploy_config, 'auth')
async with ssl_client_session(raise_for_status=True,
timeout=aiohttp.ClientTimeout(total=60),
headers=headers) as session:
async with session.post(
deploy_config.url('auth', '/api/v1alpha/logout')):
pass
auth_ns = deploy_config.service_ns('auth')
del tokens[auth_ns]
tokens.write()
print('Logged out.')
async def test_deploy():
deploy_config = get_deploy_config()
ci_deploy_status_url = deploy_config.url('ci',
'/api/v1alpha/deploy_status')
headers = service_auth_headers(deploy_config, 'ci')
async with aiohttp.ClientSession(
raise_for_status=True,
timeout=aiohttp.ClientTimeout(total=60)) as session:
async def wait_forever():
deploy_state = None
while deploy_state is None:
resp = await utils.request_retry_transient_errors(
session, 'GET', f'{ci_deploy_status_url}', headers=headers)
deploy_statuses = await resp.json()
assert len(deploy_statuses) == 1, deploy_statuses
deploy_status = deploy_statuses[0]
deploy_state = deploy_status['deploy_state']
await asyncio.sleep(5)
return deploy_state
deploy_state = await asyncio.wait_for(wait_forever(), timeout=20 * 60)
assert deploy_state == 'success', deploy_state
async def __aenter__(self):
headers = service_auth_headers(self._deploy_config, 'query')
self._session = client_session(raise_for_status=False, headers=headers)
return self
async def run(args, i):
headers = service_auth_headers(deploy_config, 'workshop', authorize_target=False)
async with client_session() as session:
# make sure notebook is up
async with session.get(deploy_config.url('workshop', ''), headers=headers) as resp:
await resp.text()
log.info(f'{i} loaded notebook home page')
# log in as workshop guest
# get csrf token
async with session.get(deploy_config.url('workshop', '/login'), headers=headers) as resp:
pass
data = aiohttp.FormData()
data.add_field(name='name', value=args.workshop)
data.add_field(name='password', value=args.password)
data.add_field(name='_csrf', value=get_cookie(session, '_csrf'))
async with session.post(deploy_config.url('workshop', '/login'), data=data, headers=headers) as resp:
pass
log.info(f'{i} logged in')
# create notebook
# get csrf token
async with session.get(deploy_config.url('workshop', '/notebook'), headers=headers) as resp:
pass
data = aiohttp.FormData()
data.add_field(name='_csrf', value=get_cookie(session, '_csrf'))
async with session.post(deploy_config.url('workshop', '/notebook'), data=data, headers=headers) as resp:
pass
log.info(f'{i} created notebook')
start = time.time()
# wait for notebook ready
ready = False
attempt = 0
# 5 attempts overkill, should only take 2: Scheduling => Running => Ready
while not ready and attempt < 5:
async with session.ws_connect(
deploy_config.url('workshop', '/notebook/wait', base_scheme='ws'), headers=headers
) as ws:
async for msg in ws:
if msg.data == '1':
ready = True
attempt += 1
end = time.time()
duration = end - start
log.info(f'{i} notebook state {ready} duration {duration}')
# delete notebook
# get csrf token
async with session.get(deploy_config.url('workshop', '/notebook'), headers=headers) as resp:
pass
data = aiohttp.FormData()
data.add_field(name='_csrf', value=get_cookie(session, '_csrf'))
async with session.post(deploy_config.url('workshop', '/notebook/delete'), data=data, headers=headers) as resp:
pass
log.info(f'{i} notebook delete, done.')
return duration, ready
async def __aenter__(self):
headers = service_auth_headers(self._deploy_config, 'query')
self._session = client_session(timeout=aiohttp.ClientTimeout(total=60),
headers=headers)
return self
async def __aenter__(self):
headers = service_auth_headers(self._deploy_config, 'ci')
self._session = get_context_specific_ssl_client_session(
timeout=aiohttp.ClientTimeout(total=60), headers=headers)
return self
async def async_init(self):
if self._session is None:
self._session = client_session()
if 'Authorization' not in self._headers:
self._headers.update(
service_auth_headers(self._deploy_config, 'memory'))
def test_batch_create_validation():
bad_configs = [
# unexpected field fleep
{
'billing_project': 'foo',
'n_jobs': 5,
'token': 'baz',
'fleep': 'quam'
},
# billing project None/missing
{
'billing_project': None,
'n_jobs': 5,
'token': 'baz'
},
{
'n_jobs': 5,
'token': 'baz'
},
# n_jobs None/missing
{
'billing_project': 'foo',
'n_jobs': None,
'token': 'baz'
},
{
'billing_project': 'foo',
'token': 'baz'
},
# n_jobs wrong type
{
'billing_project': 'foo',
'n_jobs': '5',
'token': 'baz'
},
# token None/missing
{
'billing_project': 'foo',
'n_jobs': 5,
'token': None
},
{
'billing_project': 'foo',
'n_jobs': 5
},
# empty gcsfuse bucket name
{
'billing_project':
'foo',
'n_jobs':
5,
'token':
'baz',
'gcsfuse': [{
'bucket': '',
'mount_path': '/bucket',
'read_only': False
}],
},
# empty gcsfuse mount_path name
{
'billing_project': 'foo',
'n_jobs': 5,
'token': 'baz',
'gcsfuse': [{
'bucket': 'foo',
'mount_path': '',
'read_only': False
}],
},
# attribute key/value None
{
'attributes': {
'k': None
},
'billing_project': 'foo',
'n_jobs': 5,
'token': 'baz'
},
]
url = deploy_config.url('batch', '/api/v1alpha/batches/create')
headers = service_auth_headers(deploy_config, 'batch')
session = external_requests_client_session()
for config in bad_configs:
r = retry_response_returning_functions(session.post,
url,
json=config,
allow_redirects=True,
headers=headers)
assert r.status_code == 400, (config, r)
async def session(self) -> aiohttp.ClientSession:
if self._session is None:
self._session = aiohttp.ClientSession(
headers=service_auth_headers(self.deploy_config, 'query'))
return self._session
