biolmai 0.1.4__py2.py3-none-any.whl → 0.1.7__py2.py3-none-any.whl

biolmai/asynch.py

@@ -1,23 +1,15 @@
+import asyncio
+from asyncio import create_task, gather, run
+from itertools import zip_longest
+from typing import Dict, List
+
 import aiohttp.resolver
+from aiohttp import ClientSession
 
+from biolmai.auth import get_user_auth_header
 from biolmai.const import BASE_API_URL, MULTIPROCESS_THREADS
 
 aiohttp.resolver.DefaultResolver = aiohttp.resolver.AsyncResolver
-from aiohttp import ClientSession, TCPConnector
-from typing import List
-import json
-import asyncio
-
-from asyncio import create_task, gather, run, sleep
-
-
-
-async def get_one(session: ClientSession, slug: str, action: str,
-                  payload: dict, response_key: str):
-    pass
-
-
-from aiohttp import ClientSession
 
 
 async def get_one(session: ClientSession, url: str) -> None:
@@ -30,25 +22,31 @@ async def get_one(session: ClientSession, url: str) -> None:
         return text_resp
 
 
-async def get_one_biolm(session: ClientSession,
-                        url: str,
-                        pload: dict,
-                        headers: dict,
-                        response_key: str = None) -> None:
+async def get_one_biolm(
+    session: ClientSession,
+    url: str,
+    pload: dict,
+    headers: dict,
+    response_key: str = None,
+) -> None:
     print("Requesting", url)
-    pload_batch = pload.pop('batch')
-    pload_batch_size = pload.pop('batch_size')
+    pload_batch = pload.pop("batch")
+    pload_batch_size = pload.pop("batch_size")
     t = aiohttp.ClientTimeout(
-        total=1200,
-        # total timeout (time consists connection establishment for a new connection or waiting for a free connection from a pool if pool connection limits are exceeded) default value is 5 minutes, set to `None` or `0` for unlimited timeout
+        total=1600,  # 27 mins
+        # total timeout (time consists connection establishment for
+        # a new connection or waiting for a free connection from a
+        # pool if pool connection limits are exceeded) default value
+        # is 5 minutes, set to `None` or `0` for unlimited timeout
         sock_connect=None,
-        # Maximal number of seconds for connecting to a peer for a new connection, not given from a pool. See also connect.
+        # Maximal number of seconds for connecting to a peer for a
+        # new connection, not given from a pool. See also connect.
         sock_read=None
         # Maximal number of seconds for reading a portion of data from a peer
     )
     async with session.post(url, headers=headers, json=pload, timeout=t) as resp:
         resp_json = await resp.json()
-        resp_json['batch'] = pload_batch
+        resp_json["batch"] = pload_batch
         status_code = resp.status
         expected_root_key = response_key
         to_ret = []
@@ -61,9 +59,7 @@ async def get_one_biolm(session: ClientSession,
         else:
             raise ValueError("Unexpected response in parser")
         for idx, item in enumerate(list_of_individual_seq_results):
-            d = {'status_code': status_code,
-                 'batch_id': pload_batch,
-                 'batch_item': idx}
+            d = {"status_code": status_code, "batch_id": pload_batch, "batch_item": idx}
             if not status_code or status_code != 200:
                 d.update(item)  # Put all resp keys at root there
             else:
@@ -77,16 +73,15 @@ async def get_one_biolm(session: ClientSession,
         # await sleep(2)  # for demo purposes
         # text_resp = text.strip().split("\n", 1)[0]
         # print("Got response from", url, text_resp)
-        return j
 
 
 async def async_range(count):
     for i in range(count):
-        yield(i)
+        yield (i)
         await asyncio.sleep(0.0)
 
 
-async def get_all(urls: List[str], num_concurrent: int) -> List:
+async def get_all(urls: List[str], num_concurrent: int) -> list:
     url_iterator = iter(urls)
     keep_going = True
     results = []
@@ -106,22 +101,26 @@ async def get_all(urls: List[str], num_concurrent: int) -> List:
     return results
 
 
-async def get_all_biolm(url: str,
-                        ploads: List[dict],
-                        headers: dict,
-                        num_concurrent: int,
-                        response_key: str = None) -> List:
+async def get_all_biolm(
+    url: str,
+    ploads: List[Dict],
+    headers: dict,
+    num_concurrent: int,
+    response_key: str = None,
+) -> list:
     ploads_iterator = iter(ploads)
     keep_going = True
     results = []
-    connector = aiohttp.TCPConnector(limit=100,
-                                     limit_per_host=50,
-                                     ttl_dns_cache=60)
+    connector = aiohttp.TCPConnector(limit=100, limit_per_host=50, ttl_dns_cache=60)
     ov_tout = aiohttp.ClientTimeout(
         total=None,
-        # total timeout (time consists connection establishment for a new connection or waiting for a free connection from a pool if pool connection limits are exceeded) default value is 5 minutes, set to `None` or `0` for unlimited timeout
+        # total timeout (time consists connection establishment for
+        # a new connection or waiting for a free connection from a
+        # pool if pool connection limits are exceeded) default value
+        # is 5 minutes, set to `None` or `0` for unlimited timeout
         sock_connect=None,
-        # Maximal number of seconds for connecting to a peer for a new connection, not given from a pool. See also connect.
+        # Maximal number of seconds for connecting to a peer for a
+        # new connection, not given from a pool. See also connect.
         sock_read=None
         # Maximal number of seconds for reading a portion of data from a peer
     )
@@ -134,35 +133,31 @@ async def get_all_biolm(url: str,
                 except StopIteration:
                     keep_going = False
                     break
-                new_task = create_task(get_one_biolm(session, url, pload,
-                                                     headers, response_key))
+                new_task = create_task(
+                    get_one_biolm(session, url, pload, headers, response_key)
+                )
                 tasks.append(new_task)
             res = await gather(*tasks)
             results.extend(res)
     return results
 
 
-async def async_main(urls, concurrency) -> List:
+async def async_main(urls, concurrency) -> list:
     return await get_all(urls, concurrency)
 
 
-async def async_api_calls(model_name,
-                          action,
-                          headers,
-                          payloads,
-                          response_key=None):
+async def async_api_calls(model_name, action, headers, payloads, response_key=None):
     """Hit an arbitrary BioLM model inference API."""
     # Normally would POST multiple sequences at once for greater efficiency,
     # but for simplicity sake will do one at at time right now
-    url = f'{BASE_API_URL}/models/{model_name}/{action}/'
+    url = f"{BASE_API_URL}/models/{model_name}/{action}/"
 
     if not isinstance(payloads, (list, dict)):
         err = "API request payload must be a list or dict, got {}"
         raise AssertionError(err.format(type(payloads)))
 
     concurrency = int(MULTIPROCESS_THREADS)
-    return await get_all_biolm(url, payloads, headers, concurrency,
-                               response_key)
+    return await get_all_biolm(url, payloads, headers, concurrency, response_key)
 
     # payload = json.dumps(payload)
     # session = requests_retry_session()
@@ -183,3 +178,45 @@ async def async_api_calls(model_name,
     #     headers = get_user_auth_header()  # Need to re-get these now
     #     response = retry_minutes(session, url, headers, payload, tout, mins=10)
     # return response
+
+
+def async_api_call_wrapper(grouped_df, slug, action, payload_maker, response_key):
+    """Wrap API calls to assist with sequence validation as a pre-cursor to
+    each API call.
+    """
+    model_name = slug
+    # payload = payload_maker(grouped_df)
+    init_ploads = grouped_df.groupby("batch").apply(
+        payload_maker, include_batch_size=True
+    )
+    ploads = init_ploads.to_list()
+    init_ploads = init_ploads.to_frame(name="pload")
+    init_ploads["batch"] = init_ploads.index
+    init_ploads = init_ploads.reset_index(drop=True)
+    assert len(ploads) == init_ploads.shape[0]
+    for inst, b in zip_longest(ploads, init_ploads["batch"].to_list()):
+        if inst is None or b is None:
+            raise ValueError(
+                "ploads and init_ploads['batch'] are not of the same length"
+            )
+        inst["batch"] = b
+
+    headers = get_user_auth_header()  # Need to pull each time
+    # urls = [
+    #     "https://github.com",
+    #     "https://stackoverflow.com",
+    #     "https://python.org",
+    # ]
+    # concurrency = 3
+    api_resp = run(async_api_calls(model_name, action, headers, ploads, response_key))
+    api_resp = [item for sublist in api_resp for item in sublist]
+    api_resp = sorted(api_resp, key=lambda x: x["batch_id"])
+    # print(api_resp)
+    # api_resp = biolmai.api_call(model_name, action, headers, payload,
+    #                             response_key)
+    # resp_json = api_resp.json()
+    # batch_id = int(grouped_df.batch.iloc[0])
+    # batch_size = grouped_df.shape[0]
+    # response = predict_resp_many_in_one_to_many_singles(
+    #     resp_json, api_resp.status_code, batch_id, None, batch_size)
+    return api_resp

biolmai/auth.py

@@ -6,20 +6,19 @@ import stat
 import click
 import requests
 
-from biolmai.const import ACCESS_TOK_PATH, BASE_DOMAIN, GEN_TOKEN_URL, \
-    USER_BIOLM_DIR
+from biolmai.const import ACCESS_TOK_PATH, BASE_DOMAIN, GEN_TOKEN_URL, USER_BIOLM_DIR
 
 
 def validate_user_auth(api_token=None, access=None, refresh=None):
     """Validates an API token, to be used as 'Authorization: Token 1235abc'
     authentication method."""
-    url = f'{BASE_DOMAIN}/api/v1/auth/login-check/'
+    url = f"{BASE_DOMAIN}/api/v1/auth/login-check/"
     if api_token is not None:
-        headers = {'Authorization': f'Token {api_token}'}
+        headers = {"Authorization": f"Token {api_token}"}
     else:
         headers = {
-            'Cookie': 'access={};refresh={}'.format(access, refresh),
-            'Content-Type': 'application/json'
+            "Cookie": f"access={access};refresh={refresh}",
+            "Content-Type": "application/json",
         }
     try:
         r = requests.post(url=url, headers=headers)
@@ -36,28 +35,25 @@ def validate_user_auth(api_token=None, access=None, refresh=None):
 def refresh_access_token(refresh):
     """Attempt to refresh temporary user access token, by using their refresh
     token, which has a longer TTL."""
-    url = f'{BASE_DOMAIN}/api/auth/token/refresh/'
-    headers = {
-        'Cookie': 'refresh={}'.format(refresh),
-        'Content-Type': 'application/json'
-    }
+    url = f"{BASE_DOMAIN}/api/auth/token/refresh/"
+    headers = {"Cookie": f"refresh={refresh}", "Content-Type": "application/json"}
     r = requests.post(url=url, headers=headers)
     json_response = r.json()
-    if r.status_code != 200 or (r.status_code == 200 and 'code' in r.json()):
+    if r.status_code != 200 or (r.status_code == 200 and "code" in r.json()):
         pretty_json = pprint.pformat(json_response, indent=2)
         click.echo(pretty_json)
-        click.echo(f"Token refresh failed! Please login by "
-                   "running `biolmai login`.\n")
+        click.echo(
+            "Token refresh failed! Please login by " "running `biolmai login`.\n"
+        )
         return False
     else:
-        access_refresh_dict = {'access': json_response['access'],
-                               'refresh': refresh}
+        access_refresh_dict = {"access": json_response["access"], "refresh": refresh}
         save_access_refresh_token(access_refresh_dict)
         return True
 
 
 def get_auth_status():
-    environ_token = os.environ.get('BIOLMAI_TOKEN', None)
+    environ_token = os.environ.get("BIOLMAI_TOKEN", None)
     if environ_token:
         msg = "Environment variable BIOLMAI_TOKEN detected. Validating token..."
         click.echo(msg)
@@ -65,12 +61,14 @@ def get_auth_status():
     elif os.path.exists(ACCESS_TOK_PATH):
         msg = f"Credentials file found {ACCESS_TOK_PATH}. Validating token..."
         click.echo(msg)
-        with open(ACCESS_TOK_PATH, 'r') as f:
+        with open(ACCESS_TOK_PATH) as f:
             access_refresh_dict = json.load(f)
-        access = access_refresh_dict.get('access')
-        refresh = access_refresh_dict.get('refresh')
+        access = access_refresh_dict.get("access")
+        refresh = access_refresh_dict.get("refresh")
         resp = validate_user_auth(access=access, refresh=refresh)
-        if resp.status_code != 200 or (resp.status_code == 200 and 'code' in resp.json()):
+        if resp.status_code != 200 or (
+            resp.status_code == 200 and "code" in resp.json()
+        ):
             click.echo("Access token validation failed. Attempting to refresh token...")
             # Attempt to use the 'refresh' token to get a new 'access' token
             if not refresh_access_token(refresh):
@@ -78,9 +76,11 @@ def get_auth_status():
             else:
                 click.echo("Access token refresh was successful.")
     else:
-        msg = f"No https://biolm.ai credentials found. Please " \
-              f"set the environment variable BIOLMAI_TOKEN to a token from {GEN_TOKEN_URL}, or login by " \
-              "running `biolmai login`."
+        msg = (
+            f"No https://biolm.ai credentials found. Please "
+            f"set the environment variable BIOLMAI_TOKEN to a token from "
+            f"{GEN_TOKEN_URL}, or login by running `biolmai login`."
+        )
         click.echo(msg)
 
 
@@ -95,9 +95,9 @@ def generate_access_token(uname, password):
     more permanent auth method for the API, use an API token by setting the
     BIOLMAI_TOKEN environment variable.
     """
-    url = f'{BASE_DOMAIN}/api/auth/token/'
+    url = f"{BASE_DOMAIN}/api/auth/token/"
     try:
-        r = requests.post(url=url, data={'username': uname, 'password': password})
+        r = requests.post(url=url, data={"username": uname, "password": password})
         json_response = r.json()
     except Exception:
         click.echo("Login failed!\n")
@@ -118,10 +118,56 @@ def save_access_refresh_token(access_refresh_dict):
     use."""
     os.makedirs(USER_BIOLM_DIR, exist_ok=True)
     # Save token
-    with open(ACCESS_TOK_PATH, 'w') as f:
+    with open(ACCESS_TOK_PATH, "w") as f:
         json.dump(access_refresh_dict, f)
     os.chmod(ACCESS_TOK_PATH, stat.S_IRUSR | stat.S_IWUSR)
     # Validate token and print user info
-    access = access_refresh_dict.get('access')
-    refresh = access_refresh_dict.get('refresh')
+    access = access_refresh_dict.get("access")
+    refresh = access_refresh_dict.get("refresh")
     validate_user_auth(access=access, refresh=refresh)
+
+
+def get_api_token():
+    """Get a BioLM API token to use with future API requests.
+
+    Copied from https://api.biolm.ai/#d7f87dfd-321f-45ae-99b6-eb203519ddeb.
+    """
+    url = "https://biolm.ai/api/auth/token/"
+
+    payload = json.dumps(
+        {
+            "username": os.environ.get("BIOLM_USER"),
+            "password": os.environ.get("BIOLM_PASSWORD"),
+        }
+    )
+    headers = {"Content-Type": "application/json"}
+
+    response = requests.request("POST", url, headers=headers, data=payload)
+    response_json = response.json()
+
+    return response_json
+
+
+def get_user_auth_header():
+    """Returns a dict with the appropriate Authorization header, either using
+    an API token from BIOLMAI_TOKEN environment variable, or by reading the
+    credentials file at ~/.biolmai/credntials next."""
+    api_token = os.environ.get("BIOLMAI_TOKEN", None)
+    if api_token:
+        headers = {"Authorization": f"Token {api_token}"}
+    elif os.path.exists(ACCESS_TOK_PATH):
+        with open(ACCESS_TOK_PATH) as f:
+            access_refresh_dict = json.load(f)
+        access = access_refresh_dict.get("access")
+        refresh = access_refresh_dict.get("refresh")
+        headers = {
+            "Cookie": f"access={access};refresh={refresh}",
+            "Content-Type": "application/json",
+        }
+    else:
+        err = (
+            "No https://biolm.ai credentials found. Please run "
+            "`biolmai status` to debug."
+        )
+        raise AssertionError(err)
+    return headers

biolmai/biolmai.py

@@ -1,153 +1,5 @@
 """Main module."""
-import json
-import os
-import requests
-import random
-
-import json, os, requests
-import urllib3
-import datetime
-import time
-
-from requests.adapters import HTTPAdapter
-from requests.packages.urllib3.util.retry import Retry
 
 import logging
 
-from biolmai.auth import get_auth_status, refresh_access_token
-from biolmai.const import ACCESS_TOK_PATH, BASE_API_URL
-
-log = logging.getLogger('biolm_util')
-
-
-def requests_retry_session(
-    retries=3,
-    backoff_factor=0.3,
-    status_forcelist=list(range(400, 599)),
-    session=None,
-):
-    session = session or requests.Session()
-    retry = Retry(
-        total=retries,
-        read=retries,
-        connect=retries,
-        backoff_factor=backoff_factor,
-        status_forcelist=status_forcelist
-    )
-    adapter = HTTPAdapter(max_retries=retry)
-    session.mount('http://', adapter)
-    session.mount('https://', adapter)
-    return session
-
-
-def retry_minutes(sess, URL, HEADERS, dat, timeout, mins):
-    """Retry for N minutes."""
-    HEADERS.update({'Content-Type': 'application/json'})
-    attempts, max_attempts = 0, 5
-    try:
-        now = datetime.datetime.now()
-        try_until = now + datetime.timedelta(minutes=mins)
-        while datetime.datetime.now() < try_until and attempts < max_attempts:
-            response = None
-            try:
-                log.info('Trying {}'.format(datetime.datetime.now()))
-                response = sess.post(
-                    URL,
-                    headers=HEADERS,
-                    data=dat,
-                    timeout=timeout
-                )
-                if response.status_code not in (400, 404):
-                    response.raise_for_status()
-                if 'error' in response.json():
-                    raise ValueError(response.json().dumps())
-                else:
-                    break
-            except Exception as e:
-                log.warning(e)
-                if response:
-                    log.warning(response.text)
-                time.sleep(5)  # Wait 5 seconds between tries
-            attempts += 1
-        if response is None:
-            err = "Got Nonetype response"
-            raise ValueError(err)
-        elif 'Server Error' in response.text:
-            err = "Got Server Error"
-            raise ValueError(err)
-    except Exception as e:
-        return response
-    return response
-
-
-def get_user_auth_header():
-    """Returns a dict with the appropriate Authorization header, either using
-    an API token from BIOLMAI_TOKEN environment variable, or by reading the
-    credentials file at ~/.biolmai/credntials next."""
-    api_token = os.environ.get('BIOLMAI_TOKEN', None)
-    if api_token:
-        headers = {'Authorization': f'Token {api_token}'}
-    elif os.path.exists(ACCESS_TOK_PATH):
-        with open(ACCESS_TOK_PATH, 'r') as f:
-            access_refresh_dict = json.load(f)
-        access = access_refresh_dict.get('access')
-        refresh = access_refresh_dict.get('refresh')
-        headers = {
-            'Cookie': 'access={};refresh={}'.format(access, refresh),
-            'Content-Type': 'application/json'
-        }
-    else:
-        err = "No https://biolm.ai credentials found. Please run `biolmai status` to debug."
-        raise AssertionError(err)
-    return headers
-
-
-def get_api_token():
-    """Get a BioLM API token to use with future API requests.
-
-    Copied from https://api.biolm.ai/#d7f87dfd-321f-45ae-99b6-eb203519ddeb.
-    """
-    url = "https://biolm.ai/api/auth/token/"
-
-    payload = json.dumps({
-        "username": os.environ.get("BIOLM_USER"),
-        "password": os.environ.get("BIOLM_PASSWORD")
-    })
-    headers = {
-        'Content-Type': 'application/json'
-    }
-
-    response = requests.request("POST", url, headers=headers, data=payload)
-    response_json = response.json()
-
-    return response_json
-
-
-def api_call(model_name, action, headers, payload, response_key=None):
-    """Hit an arbitrary BioLM model inference API."""
-    # Normally would POST multiple sequences at once for greater efficiency,
-    # but for simplicity sake will do one at at time right now
-    url = f'{BASE_API_URL}/models/{model_name}/{action}/'
-
-    if not isinstance(payload, (list, dict)):
-        err = "API request payload must be a list or dict, got {}"
-        raise AssertionError(err.format(type(payload)))
-    payload = json.dumps(payload)
-    session = requests_retry_session()
-    tout = urllib3.util.Timeout(total=180, read=180)
-    response = retry_minutes(session, url, headers, payload, tout, mins=10)
-    # If token expired / invalid, attempt to refresh.
-    if response.status_code == 401 and os.path.exists(ACCESS_TOK_PATH):
-        # Add jitter to slow down in case we're multiprocessing so all threads
-        # don't try to re-authenticate at once
-        time.sleep(random.random() * 4)
-        with open(ACCESS_TOK_PATH, 'r') as f:
-            access_refresh_dict = json.load(f)
-        refresh = access_refresh_dict.get('refresh')
-        if not refresh_access_token(refresh):
-            err = "Unauthenticated! Please run `biolmai status` to debug or " \
-                  "`biolmai login`."
-            raise AssertionError(err)
-        headers = get_user_auth_header()  # Need to re-get these now
-        response = retry_minutes(session, url, headers, payload, tout, mins=10)
-    return response
+log = logging.getLogger("biolm_util")

biolmai/cli.py

@@ -1,37 +1,43 @@
 """Console script for biolmai."""
+import os
 import sys
+
 import click
-import os
-from biolmai.auth import generate_access_token, \
-    get_auth_status, save_access_refresh_token
-from biolmai.const import ACCESS_TOK_PATH, MULTIPROCESS_THREADS, BASE_API_URL
+
+from biolmai.auth import (
+    generate_access_token,
+    get_auth_status,
+    save_access_refresh_token,
+)
+from biolmai.const import ACCESS_TOK_PATH, BASE_API_URL, MULTIPROCESS_THREADS
 
 
 @click.command()
 def main(args=None):
     """Console script for biolmai."""
-    click.echo("Replace this message by putting your code into "
-               "biolmai.cli.main")
+    click.echo("Replace this message by putting your code into " "biolmai.cli.main")
     click.echo("See click documentation at https://click.palletsprojects.com/")
     return 0
 
 
 @click.group()
-@click.option('--debug/--no-debug', default=False)
+@click.option("--debug/--no-debug", default=False)
 def cli(debug):
     pass
 
 
 def echo_env_vars():
-    env_var_tok = os.environ.get('BIOLMAI_TOKEN', '')[:6]
+    env_var_tok = os.environ.get("BIOLMAI_TOKEN", "")[:6]
     if env_var_tok and len(env_var_tok) == 6:
-        env_var_tok += '*****************'
-    s = '\n'.join([
-        "BIOLMAI_TOKEN={}".format(env_var_tok),
-        f"BIOLMAI_ACCESS_CRED={ACCESS_TOK_PATH}",
-        "BIOLMAI_THREADS={}".format(MULTIPROCESS_THREADS or ''),
-        f"BIOLMAI_BASE_API_URL={BASE_API_URL}"
-    ])
+        env_var_tok += "*****************"
+    s = "\n".join(
+        [
+            f"BIOLMAI_TOKEN={env_var_tok}",
+            f"BIOLMAI_ACCESS_CRED={ACCESS_TOK_PATH}",
+            "BIOLMAI_THREADS={}".format(MULTIPROCESS_THREADS or ""),
+            f"BIOLMAI_BASE_API_URL={BASE_API_URL}",
+        ]
+    )
     click.echo(s)
 
 
@@ -43,17 +49,19 @@ def status():
 
 @cli.command()
 def login():
-    uname = click.prompt("Username", default=None, hide_input=False,
-                  confirmation_prompt=False, type=str)
-    password = click.prompt("Password", default=None, hide_input=True,
-                  confirmation_prompt=False, type=str)
+    uname = click.prompt(
+        "Username", default=None, hide_input=False, confirmation_prompt=False, type=str
+    )
+    password = click.prompt(
+        "Password", default=None, hide_input=True, confirmation_prompt=False, type=str
+    )
     access_refresh_tok_dict = generate_access_token(uname, password)
     try:
-        access = access_refresh_tok_dict.get('access')
-        refresh = access_refresh_tok_dict.get('refresh')
+        assert access_refresh_tok_dict.get("access") is not None
+        assert access_refresh_tok_dict.get("refresh") is not None
         click.echo("Saving new access and refresh token.")
         save_access_refresh_token(access_refresh_tok_dict)
-    except Exception as e:
+    except Exception:
         click.echo("Unhandled login exception!")
         raise