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