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

biolib/biolib_api_client/lfs_types.py

@@ -1,35 +1,19 @@
-from biolib.typing_utils import TypedDict, List
+from biolib.typing_utils import TypedDict
 
 
-class LargeFileSystemVersion(TypedDict):
+class DataRecordVersion(TypedDict):
     presigned_download_url: str
     size_bytes: int
     uri: str
     uuid: str
 
 
-class LargeFileSystem(TypedDict):
+class DataRecordInfo(TypedDict):
     uri: str
     uuid: str
 
 
-class LfsVersionPresignedUploadUrlResponse(TypedDict):
-    presigned_upload_url: str
-
-
-class LfsUploadPartMetadata(TypedDict):
-    PartNumber: int
-    ETag: str
-
-
-class ZipProxyFileMetadata(TypedDict):
-    path: str
-    size_bytes: int
-
-
-class ZipProxyFileListResponse(TypedDict):
-    files: List[ZipProxyFileMetadata]
-
-
-class LargeFileSystemVersionMetadata(LargeFileSystemVersion):
-    files: List[ZipProxyFileMetadata]
+class DataRecordVersionInfo(TypedDict):
+    resource_uri: str
+    resource_uuid: str
+    resource_version_uuid: str

biolib/biolib_api_client/user_state.py

@@ -0,0 +1,56 @@
+from biolib.biolib_logging import logger_no_user_data
+from biolib.typing_utils import Optional, TypedDict
+from biolib.utils.cache_state import CacheState
+
+# TODO: Save job keys in the user state instead of a separate state file
+# UuidStr = str
+# class JobStateType(TypedDict):
+#     job_uuid: UuidStr
+#     aes_key: Optional[str]
+
+
+class UserStateType(TypedDict):
+    refresh_token: Optional[str]
+    # jobs: Dict[UuidStr, JobStateType]
+
+
+class UserState(CacheState[UserStateType]):
+    def __init__(self) -> None:
+        super().__init__(fail_fast_on_lock_acquire=True)
+        self._is_in_memory_only: bool = False
+
+    @property
+    def _state_path(self) -> str:
+        return f'{self._user_cache_dir}/user-state.json'
+
+    def _get_default_state(self) -> UserStateType:
+        return UserStateType(refresh_token=None)
+
+    def __enter__(self) -> UserStateType:
+        if self._is_in_memory_only:
+            if self._state is None:
+                self._state = self._get_default_state()
+            return self._state
+        try:
+            return super().__enter__()
+        except Exception as error:
+            logger_no_user_data.warning(
+                f'UserState: Could not access state file, continuing with in-memory state only. '
+                f'Login state will not persist across Python processes. Error: {error}'
+            )
+            self._is_in_memory_only = True
+            if self._state is None:
+                self._state = self._get_default_state()
+            return self._state
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        if self._is_in_memory_only:
+            return
+        try:
+            super().__exit__(exc_type, exc_val, exc_tb)
+        except Exception as error:
+            logger_no_user_data.warning(
+                f'UserState: Could not write state file. '
+                f'Login state will not persist across Python processes. Error: {error}'
+            )
+            self._is_in_memory_only = True

biolib/biolib_binary_format/__init__.py

@@ -1,9 +1,6 @@
 from .module_input import *
-from .module_output import *
-from .aes_encrypted_package import *
-from .attestation_document import *
-from .rsa_encrypted_aes_package import *
 from .system_status_update import *
 from .saved_job import *
 from .system_exception import *
 from .base_bbf_package import *
+from .module_output_v2 import *

biolib/biolib_binary_format/file_in_container.py

@@ -0,0 +1,105 @@
+import os.path
+import tarfile
+import tempfile
+
+from docker.models.containers import Container  # type: ignore
+
+from biolib.typing_utils import Iterable, Optional
+
+
+class FileInContainer:
+    def __init__(self, container: Container, path_in_container: str, overlay_upper_dir_path: Optional[str]):
+        self._container: Container = container
+        self._path_on_disk: Optional[str] = overlay_upper_dir_path + path_in_container if overlay_upper_dir_path \
+            else None
+
+        self._path_in_container: str = path_in_container
+        self._path: str = path_in_container
+        self._buffered_file_data: Optional[bytes] = None
+
+    def __repr__(self) -> str:
+        return f'FileInContainer({self.path})'
+
+    @property
+    def path(self) -> str:
+        return self._path
+
+    @path.setter
+    def path(self, value: str) -> None:
+        self._path = value
+
+    def is_file(self) -> bool:
+        if self._path_on_disk:
+            return os.path.isfile(self._path_on_disk)
+        else:
+            tmp_file = self._get_temp_file_from_container_via_tar()
+            if tmp_file:
+                os.remove(tmp_file)
+                return True
+            else:
+                return False
+
+    def get_data_size_in_bytes(self) -> int:
+        if self._path_on_disk:
+            return os.path.getsize(self._path_on_disk)
+        else:
+            tmp_file = self._get_temp_file_from_container_via_tar()
+            if tmp_file:
+                file_length = os.path.getsize(tmp_file)
+                os.remove(tmp_file)
+                return file_length
+            else:
+                return 0
+
+    def get_data_stream(self) -> Iterable[bytes]:
+        if self._path_on_disk:
+            with open(self._path_on_disk, mode='rb') as file:
+                while True:
+                    chunk = file.read(1_000_000)
+                    if not chunk:
+                        return
+
+                    yield chunk
+        else:
+            tmp_file = self._get_temp_file_from_container_via_tar()
+            if not tmp_file:
+                yield bytes()
+                return
+            else:
+                file = open(tmp_file, mode='rb')
+                while True:
+                    chunk = file.read(1_000_000)
+                    if not chunk:
+                        file.close()
+                        os.remove(tmp_file)
+                        return
+
+                    yield chunk
+
+    def _get_temp_file_from_container_via_tar(self) -> Optional[str]:
+        with tempfile.NamedTemporaryFile(mode='wb', delete=True) as tmp_io:
+            stream, _ = self._container.get_archive(path=self._path_in_container)
+            for chunk in stream:
+                tmp_io.write(chunk)
+
+            tmp_io.seek(0)
+            with tarfile.open(tmp_io.name) as tar:
+                members = tar.getmembers()
+                file_members = [member for member in members if member.isfile()]
+                if not file_members:
+                    # Path was not a file
+                    return None
+
+                assert len(file_members) == 1
+                file_obj = tar.extractfile(member=file_members[0])
+                if not file_obj:
+                    # Path was not a file
+                    return None
+                tmp_io_2 = tempfile.NamedTemporaryFile(mode='wb', delete=False)
+                while True:
+                    chunk = file_obj.read(1_000_000)
+                    if not chunk:
+                        break
+                    tmp_io_2.write(chunk)
+                tmp_io_2.close()
+                return tmp_io_2.name

biolib/biolib_binary_format/module_input.py

@@ -1,4 +1,12 @@
 from biolib.biolib_binary_format.base_bbf_package import BioLibBinaryFormatBasePackage
+from biolib.biolib_logging import logger
+from biolib.typing_utils import TypedDict, Dict, List
+
+
+class ModuleInputDict(TypedDict):
+    stdin: bytes
+    files: Dict[str, bytes]
+    arguments: List[str]
 
 
 class ModuleInput(BioLibBinaryFormatBasePackage):
@@ -6,7 +14,11 @@ class ModuleInput(BioLibBinaryFormatBasePackage):
         super().__init__(bbf)
         self.package_type = 1
 
-    def serialize(self, stdin, arguments, files):
+    def serialize(self, stdin, arguments, files) -> bytes:
+        for path in files.keys():
+            if '//' in path:
+                raise ValueError(f"File path '{path}' contains double slashes which are not allowed")
+
         bbf_data = bytearray()
         bbf_data.extend(self.version.to_bytes(1, 'big'))
         bbf_data.extend(self.package_type.to_bytes(1, 'big'))
@@ -22,19 +34,21 @@ class ModuleInput(BioLibBinaryFormatBasePackage):
         bbf_data.extend(stdin)
 
         for argument in arguments:
-            bbf_data.extend(len(argument).to_bytes(2, 'big'))
-            bbf_data.extend(argument.encode())
+            encoded_argument = argument.encode()
+            bbf_data.extend(len(encoded_argument).to_bytes(2, 'big'))
+            bbf_data.extend(encoded_argument)
 
         for path, data in files.items():
-            bbf_data.extend(len(path.encode()).to_bytes(4, 'big'))
+            encoded_path = path.encode()
+            bbf_data.extend(len(encoded_path).to_bytes(4, 'big'))
             bbf_data.extend(len(data).to_bytes(8, 'big'))
 
-            bbf_data.extend(path.encode())
+            bbf_data.extend(encoded_path)
             bbf_data.extend(data)
 
         return bbf_data
 
-    def deserialize(self):
+    def deserialize(self) -> ModuleInputDict:
         version = self.get_data(1, output_type='int')
         package_type = self.get_data(1, output_type='int')
         self.check_version_and_type(version=version, package_type=package_type, expected_package_type=self.package_type)
@@ -58,6 +72,9 @@ class ModuleInput(BioLibBinaryFormatBasePackage):
             data_len = self.get_data(8, output_type='int')
             path = self.get_data(path_len, output_type='str')
             data = self.get_data(data_len)
+            if '//' in path:
+                # TODO: Raise ValueError here once backwards compatibility period is over
+                logger.warning(f"File path '{path}' contains double slashes which are not allowed")
             files[path] = bytes(data)
 
-        return {'stdin': stdin, 'arguments': arguments, 'files': files}
+        return ModuleInputDict(stdin=stdin, arguments=arguments, files=files)

biolib/biolib_binary_format/module_output_v2.py

@@ -0,0 +1,149 @@
+from biolib.biolib_binary_format import BioLibBinaryFormatBasePackage
+from biolib.biolib_binary_format.utils import IndexableBuffer, InMemoryIndexableBuffer, LazyLoadedFile
+from biolib.biolib_binary_format.file_in_container import FileInContainer
+from biolib.typing_utils import TypedDict, List, Optional
+
+
+class Metadata(TypedDict):
+    version: int
+    type: int
+    stdout_length: int
+    stderr_length: int
+    files_info_length: int
+    files_data_length: int
+    exit_code: int
+
+
+class ModuleOutputV2(BioLibBinaryFormatBasePackage):
+    _version = 1
+    _type = 11
+    _metadata_byte_lengths = dict(
+        # Note: order is important
+        version=1,
+        type=1,
+        stdout_length=8,
+        stderr_length=8,
+        files_info_length=8,
+        files_data_length=8,
+        exit_code=2,
+    )
+    _metadata_length = sum(_metadata_byte_lengths.values())
+    _file_path_length_bytes = 4
+    _file_data_length_bytes = 8
+
+    def __init__(self, buffer: IndexableBuffer):
+        super().__init__()
+        self._buffer = buffer
+
+        self._metadata: Optional[Metadata] = None
+        self._stdout: Optional[bytes] = None
+        self._stderr: Optional[bytes] = None
+        self._files: Optional[List[LazyLoadedFile]] = None
+
+    @property
+    def buffer(self) -> IndexableBuffer:
+        return self._buffer
+
+    def get_exit_code(self) -> int:
+        metadata = self._get_metadata()
+        return metadata['exit_code']
+
+    def get_stdout(self) -> bytes:
+        if self._stdout is None:
+            metadata = self._get_metadata()
+            self._stdout = self._buffer.get_data(start=self._metadata_length, length=metadata['stdout_length'])
+
+        return self._stdout
+
+    def get_stderr(self) -> bytes:
+        if self._stderr is None:
+            metadata = self._get_metadata()
+            self._stderr = self._buffer.get_data(
+                start=self._metadata_length + metadata['stdout_length'],
+                length=metadata['stderr_length'],
+            )
+
+        return self._stderr
+
+    def get_files(self) -> List[LazyLoadedFile]:
+        metadata = self._get_metadata()
+        if self._files is None:
+            self._files = []
+            if metadata['files_info_length'] == 0:
+                return self._files
+
+            files_info_start = self._metadata_length + metadata['stdout_length'] + metadata['stderr_length']
+            files_info_buffer = InMemoryIndexableBuffer(
+                data=self._buffer.get_data(start=files_info_start, length=metadata['files_info_length'])
+            )
+
+            files_data_pointer = files_info_start + metadata['files_info_length']
+            while files_info_buffer.pointer < len(files_info_buffer):
+                path_length = files_info_buffer.get_data_with_pointer_as_int(self._file_path_length_bytes)
+                path = files_info_buffer.get_data_with_pointer_as_string(path_length)
+                data_length = files_info_buffer.get_data_with_pointer_as_int(self._file_data_length_bytes)
+
+                data_start = files_data_pointer
+                files_data_pointer += data_length
+
+                self._files.append(LazyLoadedFile(path=path, buffer=self._buffer, start=data_start, length=data_length))
+
+        return self._files
+
+    def _get_metadata(self) -> Metadata:
+        if self._metadata is None:
+            metadata_buffer = InMemoryIndexableBuffer(self._buffer.get_data(start=0, length=self._metadata_length))
+
+            partial_metadata = {}
+            for field_name, field_length in self._metadata_byte_lengths.items():
+                value = metadata_buffer.get_data_with_pointer_as_int(length=field_length)  # type: ignore
+                if field_name == 'version' and value != ModuleOutputV2._version:
+                    raise Exception('Version does not match')
+
+                if field_name == 'type' and value != ModuleOutputV2._type:
+                    raise Exception('Type does not match')
+
+                partial_metadata[field_name] = value
+
+            self._metadata = partial_metadata  # type: ignore
+
+        return self._metadata  # type: ignore
+
+    @staticmethod
+    def write_to_file(
+            output_file_path: str,
+            exit_code: int,
+            files: List[FileInContainer],
+            stderr: bytes,
+            stdout: bytes,
+    ) -> None:
+        with open(output_file_path, mode='wb') as output_file:
+            meta_lengths = ModuleOutputV2._metadata_byte_lengths
+
+            output_file.write(ModuleOutputV2._version.to_bytes(meta_lengths['version'], byteorder='big'))
+            output_file.write(ModuleOutputV2._type.to_bytes(meta_lengths['type'], byteorder='big'))
+            output_file.write(len(stdout).to_bytes(meta_lengths['stdout_length'], byteorder='big'))
+            output_file.write(len(stderr).to_bytes(meta_lengths['stdout_length'], byteorder='big'))
+
+            files_info = bytearray()
+            files_data_size = 0
+
+            for file in files:
+                file_size = file.get_data_size_in_bytes()
+                path_as_bytes = file.path.encode()
+                path_length_as_bytes = len(path_as_bytes).to_bytes(length=4, byteorder='big')
+                data_byte_length_as_bytes = file_size.to_bytes(length=8, byteorder='big')
+
+                files_info.extend(path_length_as_bytes + path_as_bytes + data_byte_length_as_bytes)
+                files_data_size += file_size
+
+            output_file.write(len(files_info).to_bytes(length=8, byteorder='big'))
+            output_file.write(files_data_size.to_bytes(length=8, byteorder='big'))
+            output_file.write(exit_code.to_bytes(length=meta_lengths['exit_code'], byteorder='big'))
+            output_file.write(stdout)
+            output_file.write(stderr)
+            output_file.write(files_info)
+
+            for file in files:
+                for chunk in file.get_data_stream():
+                    output_file.write(chunk)

biolib/biolib_binary_format/remote_endpoints.py

@@ -0,0 +1,34 @@
+from datetime import datetime, timedelta, timezone
+
+from biolib.biolib_api_client.biolib_job_api import BiolibJobApi
+from biolib.biolib_binary_format.utils import RemoteEndpoint
+
+# from urllib.parse import urlparse, parse_qs
+from biolib.biolib_logging import logger
+from biolib.typing_utils import Literal
+
+
+class RemoteJobStorageEndpoint(RemoteEndpoint):
+    def __init__(self, job_uuid: str, job_auth_token: str, storage_type: Literal['input', 'output']):
+        self._expires_at = None
+        self._job_auth_token = job_auth_token
+        self._job_uuid = job_uuid
+        self._presigned_url = None
+        self._storage_type: Literal['input', 'output'] = storage_type
+
+    def get_remote_url(self):
+        if not self._presigned_url or not self._expires_at or datetime.now(timezone.utc) > self._expires_at:
+            self._presigned_url = BiolibJobApi.get_job_storage_download_url(
+                job_auth_token=self._job_auth_token,
+                job_uuid=self._job_uuid,
+                storage_type='results' if self._storage_type == 'output' else 'input',
+            )
+            self._expires_at = datetime.now(timezone.utc) + timedelta(minutes=8)
+            # TODO: Use expires at from url
+            # parsed_url = urlparse(self._presigned_url)
+            # query_params = parse_qs(parsed_url.query)
+            # time_at_generation = datetime.datetime.strptime(query_params['X-Amz-Date'][0], '%Y%m%dT%H%M%SZ')
+            # self._expires_at = time_at_generation + timedelta(seconds=int(query_params['X-Amz-Expires'][0]))
+            logger.debug(f'Job "{self._job_uuid}" fetched presigned URL with expiry at {self._expires_at.isoformat()}')
+
+        return self._presigned_url

biolib/biolib_binary_format/remote_stream_seeker.py

@@ -0,0 +1,59 @@
+from biolib.biolib_binary_format.utils import IndexableBuffer
+from biolib.biolib_logging import logger
+from biolib.typing_utils import Iterable
+
+
+class StreamSeeker:
+    def __init__(
+        self,
+        upstream_buffer: IndexableBuffer,
+        files_data_start: int,
+        files_data_end: int,
+        max_chunk_size: int,
+    ):
+        self._upstream_buffer = upstream_buffer
+        self._files_data_end = files_data_end
+        self._max_chunk_size = max_chunk_size
+
+        self._buffer_start = files_data_start
+        self._buffer = bytearray()
+
+    def seek_and_read(self, file_start: int, file_length: int, read_ahead_bytes: int = 0) -> Iterable[bytes]:
+        assert file_start >= self._buffer_start
+        self._buffer = self._buffer[file_start - self._buffer_start :]
+        self._buffer_start = file_start
+
+        while True:
+            file_byte_count_remaining = file_length - (self._buffer_start - file_start)
+            if file_byte_count_remaining <= 0:
+                return
+
+            if len(self._buffer) > 0:
+                take = min(file_byte_count_remaining, len(self._buffer))
+                chunk = self._buffer[:take]
+                if chunk:
+                    yield chunk
+                self._buffer = self._buffer[take:]
+                self._buffer_start += take
+            else:
+                start_of_fetch = self._buffer_start + len(self._buffer)
+                bytes_left_in_stream = self._files_data_end - start_of_fetch
+                if bytes_left_in_stream <= 0:
+                    logger.error(
+                        'StreamSeeker: no bytes left upstream (start_of_fetch=%d, files_data_end=%d)',
+                        start_of_fetch,
+                        self._files_data_end,
+                    )
+                    return
+
+                fetch_size = min(self._max_chunk_size, file_byte_count_remaining + read_ahead_bytes)
+                if fetch_size > bytes_left_in_stream:
+                    logger.error(
+                        'StreamSeeker: fetch_size (%d) > bytes_left_in_stream (%d); clamping',
+                        fetch_size,
+                        bytes_left_in_stream,
+                    )
+                    fetch_size = bytes_left_in_stream
+
+                fetched_data = self._upstream_buffer.get_data(start=start_of_fetch, length=fetch_size)
+                self._buffer.extend(fetched_data)

biolib/biolib_binary_format/saved_job.py

@@ -11,8 +11,9 @@ class SavedJob(BioLibBinaryFormatBasePackage):
         bbf_data.extend(self.version.to_bytes(1, 'big'))
         bbf_data.extend(self.package_type.to_bytes(1, 'big'))
 
-        bbf_data.extend(len(saved_job_json_string).to_bytes(4, 'big'))
-        bbf_data.extend(saved_job_json_string.encode())
+        encoded_saved_job_json_string = saved_job_json_string.encode()
+        bbf_data.extend(len(encoded_saved_job_json_string).to_bytes(4, 'big'))
+        bbf_data.extend(encoded_saved_job_json_string)
 
         return bbf_data
 

biolib/biolib_binary_format/{attestation_document.py → stdout_and_stderr.py}

@@ -1,18 +1,18 @@
 from biolib.biolib_binary_format.base_bbf_package import BioLibBinaryFormatBasePackage
 
 
-class AttestationDocument(BioLibBinaryFormatBasePackage):
+class StdoutAndStderr(BioLibBinaryFormatBasePackage):
     def __init__(self, bbf=None):
         super().__init__(bbf)
-        self.package_type = 4
+        self.package_type = 10
 
-    def serialize(self, attestation_document_bytes):
+    def serialize(self, stdout_and_stderr_bytes):
         bbf_data = bytearray()
         bbf_data.extend(self.version.to_bytes(1, 'big'))
         bbf_data.extend(self.package_type.to_bytes(1, 'big'))
 
-        bbf_data.extend(len(attestation_document_bytes).to_bytes(4, 'big'))
-        bbf_data.extend(attestation_document_bytes)
+        bbf_data.extend(len(stdout_and_stderr_bytes).to_bytes(8, 'big'))
+        bbf_data.extend(stdout_and_stderr_bytes)
 
         return bbf_data
 
@@ -21,7 +21,7 @@ class AttestationDocument(BioLibBinaryFormatBasePackage):
         package_type = self.get_data(1, output_type='int')
         self.check_version_and_type(version=version, package_type=package_type, expected_package_type=self.package_type)
 
-        attestation_document_len = self.get_data(4, output_type='int')
-        attestation_document = self.get_data(attestation_document_len)
+        stdout_and_stderr_length = self.get_data(8, output_type='int')
+        stdout_and_stderr = self.get_data(stdout_and_stderr_length)
 
-        return attestation_document
+        return stdout_and_stderr

biolib/biolib_binary_format/system_status_update.py

@@ -12,8 +12,9 @@ class SystemStatusUpdate(BioLibBinaryFormatBasePackage):
         bbf_data.extend(self.package_type.to_bytes(1, 'big'))
 
         bbf_data.extend(progress.to_bytes(2, 'big'))
-        bbf_data.extend(len(log_message).to_bytes(4, 'big'))
-        bbf_data.extend(log_message.encode())
+        encoded_log_message = log_message.encode()
+        bbf_data.extend(len(encoded_log_message).to_bytes(4, 'big'))
+        bbf_data.extend(encoded_log_message)
 
         return bbf_data