pybiolib 0.2.951__py3-none-any.whl → 1.2.1890__py3-none-any.whl

biolib/jobs/job_result.py

@@ -0,0 +1,185 @@
+import time
+from pathlib import Path
+
+from biolib._internal.utils import PathFilter, filter_lazy_loaded_files
+from biolib.biolib_binary_format import ModuleOutputV2
+from biolib.biolib_binary_format.remote_endpoints import RemoteJobStorageEndpoint
+from biolib.biolib_binary_format.remote_stream_seeker import StreamSeeker
+from biolib.biolib_binary_format.utils import LazyLoadedFile, RemoteIndexableBuffer
+from biolib.biolib_errors import BioLibError
+from biolib.biolib_logging import logger
+from biolib.typing_utils import Dict, List, Optional
+
+
+class JobResult:
+    def __init__(
+        self,
+        job_uuid: str,
+        job_auth_token: str,
+        module_output: Optional[ModuleOutputV2] = None,
+    ):
+        self._job_uuid: str = job_uuid
+        self._job_auth_token: str = job_auth_token
+
+        self._module_output: Optional[ModuleOutputV2] = module_output
+
+    def get_stdout(self) -> bytes:
+        return self._get_module_output().get_stdout()
+
+    def get_stderr(self) -> bytes:
+        return self._get_module_output().get_stderr()
+
+    def get_exit_code(self) -> int:
+        return self._get_module_output().get_exit_code()
+
+    def save_files(
+        self,
+        output_dir: str,
+        path_filter: Optional[PathFilter] = None,
+        skip_file_if_exists: bool = False,
+        overwrite: bool = False,
+        flat: bool = False,
+    ) -> None:
+        module_output = self._get_module_output()
+        output_files = module_output.get_files()
+        filtered_output_files = filter_lazy_loaded_files(output_files, path_filter) if path_filter else output_files
+
+        if len(filtered_output_files) == 0:
+            logger.debug('No output files to save')
+            return
+
+        if flat:
+            basename_to_paths: Dict[str, List[str]] = {}
+            for file in filtered_output_files:
+                basename = Path(file.path).name
+                if basename not in basename_to_paths:
+                    basename_to_paths[basename] = []
+                basename_to_paths[basename].append(file.path)
+
+            duplicates = {basename: paths for basename, paths in basename_to_paths.items() if len(paths) > 1}
+
+            if duplicates:
+                max_shown = 3
+                error_parts = []
+                sorted_basenames = sorted(duplicates.keys())
+
+                for basename in sorted_basenames[:max_shown]:
+                    paths = duplicates[basename]
+                    error_parts.append(f'  {basename}: ({", ".join(paths)})')
+
+                error_message = 'Cannot save files in flat mode: duplicate filenames detected:\n' + '\n'.join(
+                    error_parts
+                )
+
+                if len(duplicates) > max_shown:
+                    remaining = len(duplicates) - max_shown
+                    error_message += f'\n  (and {remaining} more)'
+
+                raise BioLibError(error_message)
+
+        major_gap_threshold = 50_000
+        n = len(filtered_output_files)
+
+        next_break_end = [0] * n
+        if n > 0:
+            next_break_end[n - 1] = filtered_output_files[n - 1].start + filtered_output_files[n - 1].length
+            for i in range(n - 2, -1, -1):
+                end_i = filtered_output_files[i].start + filtered_output_files[i].length
+                gap = filtered_output_files[i + 1].start - end_i
+                if gap >= major_gap_threshold:
+                    next_break_end[i] = end_i
+                else:
+                    next_break_end[i] = next_break_end[i + 1]
+
+        total_files_data_to_download_in_bytes = 0
+        file_read_ahead_map = {}
+        for i, file in enumerate(filtered_output_files):
+            total_files_data_to_download_in_bytes += file.length
+            end_i = file.start + file.length
+            read_ahead_bytes = max(0, next_break_end[i] - end_i)
+            file_read_ahead_map[i] = read_ahead_bytes
+
+        # Assume files are in order
+        first_file = filtered_output_files[0]
+        last_file = filtered_output_files[len(filtered_output_files) - 1]
+        stream_seeker = StreamSeeker(
+            files_data_start=first_file.start,
+            files_data_end=last_file.start + last_file.length,
+            max_chunk_size=min(total_files_data_to_download_in_bytes, 10_000_000),
+            upstream_buffer=module_output.buffer,
+        )
+
+        logger.info(f'Saving {len(filtered_output_files)} files to {output_dir}...')
+        for file_index, file in enumerate(filtered_output_files):
+            if flat:
+                destination_file_path = Path(output_dir) / Path(file.path).name
+            else:
+                # Remove leading slash of file_path
+                destination_file_path = Path(output_dir) / Path(file.path.lstrip('/'))
+
+            if destination_file_path.exists():
+                if skip_file_if_exists:
+                    logger.info(f'Skipping {destination_file_path} as a file with that name already exists locally.')
+                    continue
+                elif not overwrite:
+                    raise BioLibError(f'File {destination_file_path} already exists. Set overwrite=True to overwrite.')
+                else:
+                    destination_file_path.rename(
+                        f'{destination_file_path}.biolib-renamed.{time.strftime("%Y%m%d%H%M%S")}'
+                    )
+
+            dir_path = destination_file_path.parent
+            if dir_path:
+                dir_path.mkdir(parents=True, exist_ok=True)
+
+            # write content to temporary (partial) file
+            partial_path = destination_file_path.with_suffix(
+                destination_file_path.suffix + f'.{self._job_uuid}.partial_biolib_download'
+            )
+            file_start = file.start
+            data_to_download = file.length
+            if partial_path.exists():
+                data_already_downloaded = partial_path.stat().st_size
+                file_start += data_already_downloaded
+                data_to_download -= data_already_downloaded
+
+            read_ahead_bytes = file_read_ahead_map[file_index]
+
+            with open(partial_path, mode='ab') as partial_file:
+                for chunk in stream_seeker.seek_and_read(
+                    file_start=file_start, file_length=data_to_download, read_ahead_bytes=read_ahead_bytes
+                ):
+                    partial_file.write(chunk)
+
+            # rename partial file to actual file name
+            partial_path.rename(destination_file_path)
+
+    def get_output_file(self, filename) -> LazyLoadedFile:
+        files = self._get_module_output().get_files()
+        filtered_files = filter_lazy_loaded_files(files, path_filter=filename)
+        if not filtered_files:
+            raise BioLibError(f'File {filename} not found in results.')
+
+        if len(filtered_files) != 1:
+            raise BioLibError(f'Found multiple results for filename {filename}.')
+
+        return filtered_files[0]
+
+    def list_output_files(self, path_filter: Optional[PathFilter] = None) -> List[LazyLoadedFile]:
+        files = self._get_module_output().get_files()
+        if not path_filter:
+            return files
+
+        return filter_lazy_loaded_files(files, path_filter)
+
+    def _get_module_output(self) -> ModuleOutputV2:
+        if self._module_output is None:
+            remote_job_storage_endpoint = RemoteJobStorageEndpoint(
+                job_auth_token=self._job_auth_token,
+                job_uuid=self._job_uuid,
+                storage_type='output',
+            )
+            buffer = RemoteIndexableBuffer(endpoint=remote_job_storage_endpoint)
+            self._module_output = ModuleOutputV2(buffer)
+
+        return self._module_output

biolib/jobs/types.py

@@ -0,0 +1,50 @@
+from biolib.typing_utils import List, Literal, Optional, TypedDict
+
+JobState = Literal['in_progress', 'completed', 'failed', 'cancelled']
+
+
+class _BaseCloudJobDict(TypedDict):
+    created_at: str
+    finished_at: Optional[str]
+    uuid: str
+    error_code: int
+
+
+class CloudJobDict(_BaseCloudJobDict):
+    started_at: Optional[str]
+    compute_node_url: Optional[str]
+
+
+class CloudJobStartedDict(_BaseCloudJobDict):
+    started_at: str
+    compute_node_url: str
+
+
+class Result(TypedDict):
+    name: str
+
+
+class JobDict(TypedDict):
+    app_uri: str
+    arguments_override_command: bool
+    auth_token: str
+    created_at: str
+    ended_at: Optional[str]
+    requested_machine: str
+    runtime_seconds: int
+    main_result: Result
+    started_at: str
+    state: JobState
+    uuid: str
+    cloud_job: Optional[CloudJobDict]
+
+
+class BasePaginatedResponse(TypedDict):
+    current_page_number: int
+    object_count: int
+    page_count: int
+    page_size: int
+
+
+class JobsPaginatedResponse(BasePaginatedResponse):
+    results: List[JobDict]

biolib/runtime/__init__.py

@@ -0,0 +1,14 @@
+import warnings
+
+from biolib._runtime.runtime import Runtime as _Runtime
+
+
+def set_main_result_prefix(result_prefix: str) -> None:
+    warnings.warn(
+        'The "biolib.runtime.set_main_result_prefix" function is deprecated. '
+        'It will be removed in future releases from mid 2024. '
+        'Please use "from biolib.sdk import Runtime" and then "Runtime.set_main_result_prefix" instead.',
+        DeprecationWarning,
+        stacklevel=2,
+    )
+    _Runtime.set_main_result_prefix(result_prefix)

biolib/sdk/__init__.py

@@ -0,0 +1,91 @@
+from typing import Any, Dict, List, Optional, Union
+
+# Imports to hide and use as private internal utils
+from biolib._data_record.data_record import DataRecord as _DataRecord
+from biolib._index.index import Index as _Index
+from biolib._index.query_result import IndexQueryResult
+from biolib._index.query_result import query_index as _query_index
+from biolib._internal.push_application import push_application as _push_application
+from biolib._internal.push_application import set_app_version_as_active as _set_app_version_as_active
+from biolib._runtime.runtime import Runtime as _Runtime
+from biolib._session.session import Session as _Session
+from biolib.app import BioLibApp as _BioLibApp
+
+# Classes to expose as public API
+Runtime = _Runtime
+
+
+def get_session(
+    refresh_token: str,
+    base_url: Optional[str] = None,
+    client_type: Optional[str] = None,
+    experiment: Optional[str] = None,
+) -> _Session:
+    return _Session.get_session(
+        refresh_token=refresh_token,
+        base_url=base_url,
+        client_type=client_type,
+        experiment=experiment,
+    )
+
+
+def push_app_version(uri: str, path: str) -> _BioLibApp:
+    push_data = _push_application(
+        app_uri=uri,
+        app_path=path,
+        app_version_to_copy_images_from=None,
+        set_as_active=False,
+        set_as_published=False,
+    )
+    if not push_data:
+        raise Exception('Failed to push application; please check the logs for more details')
+
+    uri = f'{push_data["app_uri"]}:{push_data["sematic_version"]}'
+    return _BioLibApp(uri)
+
+
+def set_app_version_as_default(app_version: _BioLibApp) -> None:
+    app_version_uuid = app_version.version['public_id']
+    _set_app_version_as_active(app_version_uuid)
+
+
+def get_app_version_pytest_plugin(app_version: _BioLibApp):
+    try:
+        import pytest  # type: ignore # pylint: disable=import-outside-toplevel,import-error
+    except BaseException:
+        raise Exception('Failed to import pytest; please make sure it is installed') from None
+
+    class AppVersionFixturePlugin:
+        def __init__(self, app_version_ref):
+            self.app_version_ref = app_version_ref
+
+        @pytest.fixture(scope='session')
+        def app_version(self, request):  # pylint: disable=unused-argument
+            return self.app_version_ref
+
+    return AppVersionFixturePlugin(app_version)
+
+
+def create_data_record(
+    destination: str,
+    data_path: str,
+    name: Optional[str] = None,
+    record_type: Optional[str] = None,
+) -> _DataRecord:
+    return _DataRecord.create(
+        destination=f'{destination}/{name}' if name else destination,
+        data_path=data_path,
+        record_type=record_type,
+    )
+
+
+def get_index(uri: str) -> _Index:
+    return _Index.get_by_uri(uri)
+
+
+def query_index(
+    query: str,
+    data: Optional[Union[List[Dict[str, Any]], bytes]] = None,
+    data_format: str = 'json',
+) -> IndexQueryResult:
+    return _query_index(query=query, data=data, data_format=data_format)

biolib/tables.py

@@ -0,0 +1,34 @@
+from collections import OrderedDict
+
+from rich.console import Console
+from rich.table import Column, Table
+
+from biolib.typing_utils import Any, List
+
+
+class BioLibTable:
+    def __init__(self, columns_to_row_map: OrderedDict, rows: List[Any], title):
+        self.title = title
+        self.rows = rows
+        self.columns_to_row_map = columns_to_row_map
+        self.table = self._create_table()
+
+    def _create_table(self) -> Table:
+        columns = [Column(header=header, **meta['params']) for header, meta in self.columns_to_row_map.items()]
+        table = Table(*columns, title=self.title)
+        for row in self.rows:
+            row_values: List[str] = []
+            for column in self.columns_to_row_map.values():
+                keys = column['key'].split('.')
+                value = row[keys[0]]
+                for key in keys[1:]:
+                    if not value or key not in value:
+                        continue
+                    value = value[key]
+                row_values.append(str(value))
+            table.add_row(*row_values)
+        return table
+
+    def print_table(self):
+        console = Console()
+        console.print(self.table)

biolib/typing_utils.py

@@ -1,7 +1,2 @@
-import sys
-
-# import and expose everything from the typing module
-from typing import *  # pylint: disable=wildcard-import, unused-wildcard-import
-
-if sys.version_info < (3, 8):
-    from typing_extensions import TypedDict, Literal
+# TODO: Deprecate and later remove this file
+from biolib._shared.types.typing import *  # pylint: disable=wildcard-import, unused-wildcard-import

biolib/user/__init__.py

@@ -0,0 +1 @@
+from .sign_in import sign_in, sign_out

biolib/user/sign_in.py

@@ -0,0 +1,54 @@
+import time
+import webbrowser
+
+from biolib._internal.utils import open_browser_window_from_notebook
+from biolib.biolib_api_client import BiolibApiClient
+from biolib.biolib_api_client.auth import BiolibAuthChallengeApi
+from biolib.biolib_logging import logger_no_user_data
+from biolib.utils import IS_RUNNING_IN_NOTEBOOK
+
+
+def sign_out() -> None:
+    api_client = BiolibApiClient.get(attempt_sign_in=False)
+    api_client.sign_out()
+
+
+def sign_in(open_in_default_browser: bool = False) -> None:
+    if not BiolibApiClient.is_reauthentication_needed():
+        logger_no_user_data.info('Already signed in')
+        return
+
+    api_client = BiolibApiClient.get()
+    auth_challenge = BiolibAuthChallengeApi.create_auth_challenge()
+    auth_challenge_token = auth_challenge['token']
+
+    client_type = 'notebook' if IS_RUNNING_IN_NOTEBOOK else 'cli'
+
+    frontend_sign_in_url = f'{api_client.base_url}/sign-in/request/{client_type}/?token={auth_challenge_token}'
+
+    if IS_RUNNING_IN_NOTEBOOK:
+        print(f'Opening authorization page at: {frontend_sign_in_url}')
+        print('If your browser does not open automatically, click on the link above.')
+        open_browser_window_from_notebook(frontend_sign_in_url)
+    elif open_in_default_browser:
+        print(f'Opening authorization page at: {frontend_sign_in_url}')
+        print('If your browser does not open automatically, click on the link above.')
+        webbrowser.open(frontend_sign_in_url)
+    else:
+        print('Please copy and paste the following link into your browser:')
+        print(frontend_sign_in_url)
+
+    for _ in range(100):
+        time.sleep(3)
+        auth_challenge_status = BiolibAuthChallengeApi.get_auth_challenge_status(token=auth_challenge_token)
+
+        if auth_challenge_status['state'] != 'awaiting':
+            break
+
+    user_tokens = auth_challenge_status.get('user_tokens')
+    if user_tokens:
+        api_client.set_user_tokens(user_tokens)
+        print('Successfully signed in!')
+
+    else:
+        print(f"Sign in failed. Got state: {auth_challenge_status['state']}\nPlease try again")

biolib/utils/__init__.py

@@ -0,0 +1,162 @@
+import collections.abc
+import multiprocessing
+import os
+import socket
+import sys
+
+from importlib_metadata import version, PackageNotFoundError
+
+from biolib.typing_utils import Optional
+from biolib.utils.seq_util import SeqUtil, SeqUtilRecord
+from biolib._internal.http_client import HttpClient
+from biolib.biolib_logging import logger_no_user_data, logger
+from biolib.typing_utils import Tuple, Iterator
+from .multipart_uploader import MultiPartUploader, get_chunk_iterator_from_bytes
+
+# try fetching version, if it fails (usually when in dev), add default
+try:
+    BIOLIB_PACKAGE_VERSION = version('pybiolib')
+except PackageNotFoundError:
+    BIOLIB_PACKAGE_VERSION = '0.0.0'
+
+IS_DEV = os.getenv('BIOLIB_DEV', '').upper() == 'TRUE'
+
+
+def load_base_url_from_env() -> str:
+    base_url = os.getenv('BIOLIB_BASE_URL')
+    if base_url:
+        return base_url.lower().rstrip('/')
+
+    try:
+        search_list = []
+        with open('/etc/resolv.conf') as file:
+            for line in file:
+                line_trimmed = line.strip()
+                if line_trimmed.startswith('search'):
+                    search_list = line_trimmed.split()[1:]
+                    logger.debug(f'Found search list: {search_list} when resolving base url.')
+                    break
+
+        for search_host in search_list:
+            host_to_try = f'biolib.{search_host}'
+            try:
+                if len(socket.getaddrinfo(host_to_try, 443)) > 0:
+                    return f'https://{host_to_try}'.lower()
+            except BaseException:  # pylint: disable=broad-except
+                pass
+    except BaseException:  # pylint: disable=broad-except
+        pass
+
+    return 'https://biolib.com'
+
+
+BIOLIB_BASE_URL: Optional[str] = None
+BIOLIB_SITE_HOSTNAME: Optional[str] = None
+
+BIOLIB_CLOUD_BASE_URL = os.getenv('BIOLIB_CLOUD_BASE_URL', '').lower()
+
+BIOLIB_PACKAGE_ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+
+BIOLIB_CLOUD_ENVIRONMENT = os.getenv('BIOLIB_CLOUD_ENVIRONMENT', '').lower()
+
+BIOLIB_SECRETS_TMPFS_PATH = os.environ.get('BIOLIB_SECRETS_TMPFS_PATH')
+
+IS_RUNNING_IN_CLOUD = BIOLIB_CLOUD_ENVIRONMENT == 'non-enclave'
+
+BASE_URL_IS_PUBLIC_BIOLIB: Optional[bool] = None
+
+# sys.stdout is an instance of OutStream in Jupyter and Colab which does not have .buffer
+if not hasattr(sys.stdout, 'buffer'):
+    IS_RUNNING_IN_NOTEBOOK = True
+else:
+    IS_RUNNING_IN_NOTEBOOK = False
+
+STREAM_STDOUT = False
+
+if BIOLIB_CLOUD_ENVIRONMENT and not IS_RUNNING_IN_CLOUD:
+    logger_no_user_data.warning((
+        'BIOLIB_CLOUD_ENVIRONMENT defined but does not specify the cloud environment correctly. ',
+        'The compute node will not act as a cloud compute node'
+    ))
+
+ByteRangeTuple = Tuple[int, int]
+DownloadChunkInputTuple = Tuple[ByteRangeTuple, str]
+
+
+def _download_chunk(input_tuple: DownloadChunkInputTuple) -> bytes:
+    byte_range, presigned_url = input_tuple
+    start, end = byte_range
+
+    try:
+        response = HttpClient.request(
+            url=presigned_url,
+            headers={'range': f'bytes={start}-{end}'},
+            timeout_in_seconds=300,  # timeout after 5 min
+            retries=20,
+            retry_on_http_500=True,
+        )
+    except Exception as exception:
+        logger_no_user_data.exception("Hit error downloading chunk")
+        logger_no_user_data.error(exception)
+        raise exception
+    logger_no_user_data.debug(f'Returning raw data for part {start}')
+    return response.content
+
+
+class ChunkIterator(collections.abc.Iterator):
+
+    def __init__(self, file_size: int, chunk_size: int, presigned_url: str):
+        self._semaphore = multiprocessing.BoundedSemaphore(20)  # support 20 chunks to be processed at once
+        self._iterator = self._get_chunk_input_iterator(file_size, chunk_size, presigned_url)
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        if self._semaphore.acquire(timeout=1800):
+            return next(self._iterator)
+        else:
+            raise Exception('Did not receive work within 30 min.')
+
+    def chunk_completed(self) -> None:
+        self._semaphore.release()
+
+    @staticmethod
+    def _get_chunk_input_iterator(
+            file_size: int,
+            chunk_size: int,
+            presigned_url: str,
+    ) -> Iterator[DownloadChunkInputTuple]:
+        for index in range(0, file_size, chunk_size):
+            byte_range: ByteRangeTuple = (index, index + chunk_size - 1)
+            yield byte_range, presigned_url
+
+
+def download_presigned_s3_url(presigned_url: str, output_file_path: str) -> None:
+    chunk_size = 50_000_000
+
+    response = HttpClient.request(url=presigned_url, headers={'range': 'bytes=0-1'})
+    file_size = int(response.headers['Content-Range'].split('/')[1])
+
+    chunk_iterator = ChunkIterator(file_size, chunk_size, presigned_url)
+
+    bytes_written = 0
+    # use 16 cores, unless less is available
+    process_pool = multiprocessing.Pool(processes=min(16, multiprocessing.cpu_count() - 1))
+    try:
+        with open(output_file_path, 'ab') as output_file:
+            for index, data in enumerate(process_pool.imap(_download_chunk, chunk_iterator)):
+                logger_no_user_data.debug(f'Writing part {index} to file...')
+                output_file.write(data)
+
+                bytes_written += chunk_size
+                approx_progress_percent = min(bytes_written / file_size * 100, 100)
+                logger_no_user_data.debug(
+                    f'Wrote part {index} of {file_size} to file, '
+                    f'the approximate progress is {round(approx_progress_percent, 2)}%'
+                )
+                chunk_iterator.chunk_completed()
+    finally:
+        logger_no_user_data.debug('Closing process poll...')
+        process_pool.close()
+        logger_no_user_data.debug('Process poll closed.')