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

biolib/compute_node/job_worker/large_file_system.py

@@ -1,17 +1,19 @@
 import os
 import shutil
-import subprocess
 import time
 import zipfile
 
 import docker.types  # type: ignore
 
+from biolib import utils
 from biolib.biolib_errors import BioLibError
+from biolib.biolib_logging import logger_no_user_data
 from biolib.compute_node.job_worker.cache_state import LfsCacheState
 from biolib.compute_node.job_worker.cache_types import LargeFileSystemCache, StoragePartition
 from biolib.typing_utils import TypedDict, Optional, Callable
 
 from biolib.biolib_api_client import LargeFileSystemMapping
+from biolib.utils import download_presigned_s3_url
 
 
 class StatusUpdate(TypedDict):
@@ -42,22 +44,25 @@ class LargeFileSystem:
             lfs_mapping: LargeFileSystemMapping,
             send_status_update: Callable[[StatusUpdate], None],
     ):
+        if not utils.IS_RUNNING_IN_CLOUD:
+            raise LargeFileSystemError('Large File System is currently not supported in local compute environments')
+
+        if lfs_mapping['size_bytes'] is None:
+            raise LargeFileSystemError('Error: You attempted to attach an LFS without a known size.')
+
         self._job_id: str = job_id
         self._lfs_mapping: LargeFileSystemMapping = lfs_mapping
         self._path_on_disk: Optional[str] = None
+        self._path_on_disk_for_write: Optional[str] = None
         self._send_status_update: Callable[[StatusUpdate], None] = send_status_update
 
-    @property
-    def _is_initialized(self) -> bool:
-        return self._path_on_disk is not None
-
     @property
     def uuid(self) -> str:
         return self._lfs_mapping['uuid']
 
     @property
     def docker_mount(self) -> docker.types.Mount:
-        if not self._is_initialized:
+        if not self._path_on_disk:
             raise LargeFileSystemError('LargeFileSystem not initialized')
 
         return docker.types.Mount(
@@ -68,19 +73,32 @@ class LargeFileSystem:
         )
 
     def initialize(self) -> None:
-        if self._is_initialized:
+        if self._path_on_disk:
+            logger_no_user_data.debug(f'LFS {self.uuid} is already initialized')
             return
 
         lfs_size_bytes = self._lfs_mapping['size_bytes']
+        logger_no_user_data.debug(f'Initializing LFS {self.uuid} of size {lfs_size_bytes} bytes...')
+
+        readonly_cache_state = LfsCacheState().get_read_only_dict_without_lock()
+        if readonly_cache_state:
+            readonly_lfs: Optional[LargeFileSystemCache] = readonly_cache_state['large_file_systems'].get(self.uuid)
+            if readonly_lfs and readonly_lfs['state'] == 'ready':
+                logger_no_user_data.debug(f'LFS {self.uuid} found to be ready in cache')
+                storage_partition = readonly_cache_state['storage_partitions'][readonly_lfs['storage_partition_uuid']]
+                self._path_on_disk = f"{storage_partition['path']}/lfs/{self.uuid}/data"
+                self._path_on_disk_for_write = f'{LfsCacheState().storage_path_for_write}/lfs/{self.uuid}/data'
+                return
 
         lfs_is_already_downloading = False
-
         with LfsCacheState() as cache_state:
-            lfs_cache = cache_state['large_file_systems'].get(self.uuid)
+            lfs_cache: Optional[LargeFileSystemCache] = cache_state['large_file_systems'].get(self.uuid)
 
             if lfs_cache is None:
-
+                logger_no_user_data.debug(f'LFS {self.uuid} was not found in cache')
                 storage_partition_to_use: Optional[StoragePartition] = None
+                logger_no_user_data.debug(f"Storage partitions to check: {cache_state['storage_partitions'].values()}")
+
                 for storage_partition in cache_state['storage_partitions'].values():
                     free_space_bytes = storage_partition['total_size_bytes'] - storage_partition['allocated_size_bytes']
                     if lfs_size_bytes < free_space_bytes:
@@ -93,7 +111,6 @@ class LargeFileSystem:
                     storage_partition_to_use['allocated_size_bytes'] += lfs_size_bytes
 
                 cache_state['large_file_systems'][self.uuid] = LargeFileSystemCache(
-                    active_jobs=[self._job_id],
                     last_used_at=LfsCacheState.get_timestamp_now(),
                     size_bytes=lfs_size_bytes,
                     state='downloading',
@@ -102,11 +119,15 @@ class LargeFileSystem:
                 )
 
                 self._path_on_disk = f"{storage_partition_to_use['path']}/lfs/{self.uuid}/data"
+                self._path_on_disk_for_write = f'{LfsCacheState().storage_path_for_write}/lfs/{self.uuid}/data'
+                logger_no_user_data.debug(f'Using path {self._path_on_disk} for LFS')
 
             else:
-                lfs_cache['active_jobs'].append(self._job_id)
+                logger_no_user_data.debug(f"LFS {self.uuid} found in cache with state {lfs_cache['state']}")
+                lfs_cache['last_used_at'] = LfsCacheState.get_timestamp_now()
                 storage_partition = cache_state['storage_partitions'][lfs_cache['storage_partition_uuid']]
                 self._path_on_disk = f"{storage_partition['path']}/lfs/{self.uuid}/data"
+                self._path_on_disk_for_write = f'{LfsCacheState().storage_path_for_write}/lfs/{self.uuid}/data'
 
                 if lfs_cache['state'] == 'ready':
                     return
@@ -126,24 +147,36 @@ class LargeFileSystem:
                 progress=30,
                 log_message=f'Downloading Large File System "{self.uuid}"...',
             ))
-            self._download_and_unzip()
+
+            try:
+                self._download_and_unzip()
+            except Exception as error:
+                logger_no_user_data.error(
+                    f'Failed to download LFS {self.uuid} got error: {error}. Cleaning up LFS cache state...'
+                )
+                self._remove_from_state()
+                raise error
+
             self._send_status_update(StatusUpdate(
                 progress=33,
                 log_message=f'Large File System "{self.uuid}" downloaded.',
             ))
             with LfsCacheState() as cache_state:
                 cache_state['large_file_systems'][self.uuid]['state'] = 'ready'
+                logger_no_user_data.debug(f'LFS cache state: {cache_state}')
 
-    def detach(self) -> None:
-        if not self._is_initialized:
-            return
+        logger_no_user_data.debug(f'LFS {self.uuid} is initialized')
 
+    def _remove_from_state(self) -> None:
         with LfsCacheState() as cache_state:
-            lfs_cache = cache_state['large_file_systems'][self.uuid]
-            lfs_cache['last_used_at'] = LfsCacheState.get_timestamp_now()
-            lfs_cache['active_jobs'] = [job_id for job_id in lfs_cache['active_jobs'] if job_id != self._job_id]
+            lfs = cache_state['large_file_systems'][self.uuid]
+            storage_partition = cache_state['storage_partitions'][lfs['storage_partition_uuid']]
+            storage_partition['allocated_size_bytes'] -= lfs['size_bytes']
+
+            cache_state['large_file_systems'].pop(self.uuid)
+            logger_no_user_data.debug(f'LFS cache state: {cache_state}')
 
-        self._path_on_disk = None
+        logger_no_user_data.debug('Cleaned up LFS cache state')
 
     def _wait_for_lfs_to_be_ready(self) -> None:
         # Timeout after 15 min
@@ -156,34 +189,50 @@ class LargeFileSystem:
         raise LargeFileSystemError(f'Waiting for Large File System "{self.uuid}" downloading timed out')
 
     def _download_and_unzip(self) -> None:
-        s3_lfs_bucket_name = os.environ.get('BIOLIB_S3_LFS_BUCKET_NAME')
-
+        logger_no_user_data.debug(f'Starting download and unzip of LFS {self.uuid}')
         lfs_size_bytes = self._lfs_mapping['size_bytes']
-
         tmp_storage_dir: Optional[str] = None
-        for path in LfsCacheState.get_tmp_storage_paths():
+
+        for path in LfsCacheState().tmp_storage_paths:
             disk_usage = shutil.disk_usage(path)
+            logger_no_user_data.debug(f'Path {path} has disk usage: {disk_usage}')
             if lfs_size_bytes < disk_usage.free:
                 tmp_storage_dir = path
 
         if tmp_storage_dir is None:
             raise LargeFileSystemError('No temporary storage available for downloading Large File System')
 
-        s3_data_zip_uri = f's3://{s3_lfs_bucket_name}/lfs/versions/{self.uuid}/data.zip'
         tmp_data_zip_path = f'{tmp_storage_dir}/lfs-{self.uuid}-data.zip'
+        logger_no_user_data.debug(f'Downloading LFS zip to path {tmp_data_zip_path}...')
 
-        download_result = subprocess.run(
-            ['s5cmd', 'cp', s3_data_zip_uri, tmp_data_zip_path],
-            check=False,
-            capture_output=True,
-        )
-
-        if download_result.returncode != 0:
-            raise LargeFileSystemError(
-                f'Failed to download Large File System data.zip: {download_result.stderr.decode()}'
+        try:
+            download_presigned_s3_url(
+                presigned_url=self._lfs_mapping['presigned_download_url'],
+                output_file_path=tmp_data_zip_path,
             )
+        except Exception as error:
+            logger_no_user_data.error(
+                f'Failed to download Large File System data.zip got error: {error}. Removing tmp_data_zip_path...'
+            )
+            if os.path.exists(tmp_data_zip_path):
+                os.remove(tmp_data_zip_path)
+                logger_no_user_data.debug(f'Removed {tmp_data_zip_path}')
+
+            raise LargeFileSystemError(f'Failed to download Large File System data.zip got error: {error}') from error
+
+        try:
+            logger_no_user_data.debug(f'Extracting {tmp_data_zip_path} to {self._path_on_disk_for_write} ...')
+            with zipfile.ZipFile(tmp_data_zip_path, 'r') as zip_ref:
+                zip_ref.extractall(self._path_on_disk_for_write)
+        except Exception as error:
+            logger_no_user_data.error(
+                f'Failed to unzip {tmp_data_zip_path} got error: {error}. '
+                f'Removing {self._path_on_disk_for_write}...'
+            )
+            if self._path_on_disk_for_write is not None and os.path.exists(self._path_on_disk_for_write):
+                shutil.rmtree(self._path_on_disk_for_write)
 
-        with zipfile.ZipFile(tmp_data_zip_path, 'r') as zip_ref:
-            zip_ref.extractall(self._path_on_disk)
-
-        os.remove(tmp_data_zip_path)
+            raise error
+        finally:
+            os.remove(tmp_data_zip_path)
+            logger_no_user_data.debug(f'Removed {tmp_data_zip_path}')

biolib/compute_node/job_worker/network_alloc.py

@@ -0,0 +1,99 @@
+import hashlib
+import ipaddress
+import uuid
+from typing import Dict, Optional, cast
+
+from docker.errors import APIError
+from docker.models.networks import Network
+from docker.types import IPAMConfig, IPAMPool
+
+from biolib.biolib_errors import BioLibError
+from biolib.biolib_logging import logger_no_user_data
+from biolib.compute_node.remote_host_proxy import get_static_ip_from_network
+
+
+def _iter_network_subnets(existing_network):
+    ipam_config = existing_network.attrs.get('IPAM', {}).get('Config', [])
+    for cfg in ipam_config:
+        subnet_str = cfg.get('Subnet')
+        if not subnet_str:
+            continue
+        try:
+            yield ipaddress.ip_network(subnet_str, strict=False)
+        except ValueError:
+            continue
+
+
+def _find_overlap(candidate_network, existing_networks):
+    for existing in existing_networks:
+        for subnet in _iter_network_subnets(existing):
+            if candidate_network.overlaps(subnet):
+                return existing, str(subnet)
+    return None
+
+
+def _allocate_network_with_retries(
+    name_prefix: str,
+    docker_client,
+    internal: bool = True,
+    driver: str = 'bridge',
+    max_attempts: int = 10,
+    labels: Optional[Dict[str, str]] = None,
+) -> Network:
+    base_network = ipaddress.ip_network('172.28.0.0/16', strict=False)
+
+    suffix = uuid.uuid4().hex
+    full_name = f'{name_prefix}{suffix}'
+    name_hash = int(hashlib.sha256(full_name.encode()).hexdigest(), 16)
+    starting_offset = name_hash % 256
+
+    for attempt in range(max_attempts):
+        offset = (starting_offset + attempt) % 256
+
+        if base_network.prefixlen == 16:
+            third_octet = offset
+            candidate_subnet = f'{base_network.network_address.exploded.rsplit(".", 2)[0]}.{third_octet}.0/24'
+        else:
+            candidate_subnet = f'{base_network.network_address.exploded.rsplit(".", 1)[0]}.{offset}.0/24'
+
+        candidate_network = ipaddress.ip_network(candidate_subnet, strict=False)
+
+        existing_networks = docker_client.networks.list()
+        overlap = _find_overlap(candidate_network, existing_networks)
+        if overlap:
+            existing_network, existing_subnet = overlap
+            logger_no_user_data.debug(
+                f'Subnet {candidate_subnet} conflicts with existing network '
+                f'{existing_network.name} ({existing_subnet}), trying next candidate'
+            )
+            continue
+
+        ipam_pool = IPAMPool(subnet=candidate_subnet)
+        computed_ipam_config = IPAMConfig(pool_configs=[ipam_pool])
+
+        try:
+            network = cast(
+                Network,
+                docker_client.networks.create(
+                    name=full_name,
+                    internal=internal,
+                    driver=driver,
+                    ipam=computed_ipam_config,
+                    labels=labels or {},
+                ),
+            )
+            static_ip = get_static_ip_from_network(network, offset=2)
+            logger_no_user_data.debug(
+                f'Created network {full_name} with subnet {candidate_subnet} and static IP {static_ip}'
+            )
+            return network
+        except APIError as api_error:
+            logger_no_user_data.debug(
+                f'Network creation failed with Docker API error for subnet {candidate_subnet}: {api_error}, '
+                f'trying next candidate (attempt {attempt + 1}/{max_attempts})'
+            )
+            continue
+
+    raise BioLibError(
+        f'Failed to allocate and create network {full_name} after {max_attempts} attempts. ' f'Base CIDR: 172.28.0.0/16'
+    )

biolib/compute_node/job_worker/network_buffer.py

@@ -0,0 +1,240 @@
+import contextlib
+import json
+import os
+import socket
+import time
+from typing import List, Optional
+
+from docker.errors import NotFound
+from docker.models.networks import Network
+
+from biolib import utils
+from biolib.biolib_docker_client import BiolibDockerClient
+from biolib.biolib_logging import logger_no_user_data
+from biolib.compute_node.job_worker.network_alloc import _allocate_network_with_retries
+
+
+class NetworkBuffer:
+    BUFFER_SIZE = 25
+    NETWORK_NAME_PREFIX = 'biolib-remote-host-network-'
+    NETWORK_LABEL = 'biolib-role=remote-host-network'
+
+    _BIOLIB_DIR = '/biolib' if utils.IS_RUNNING_IN_CLOUD else '/tmp/biolib'
+    _NETWORKS_FILE = os.path.join(_BIOLIB_DIR, 'remote-host-networks.json')
+    _LOCK_FILE = os.path.join(_BIOLIB_DIR, 'remote-host-networks.lock')
+    _LOCK_TIMEOUT_SECONDS = 60
+    _STALE_LOCK_THRESHOLD_SECONDS = 600
+
+    _instance: Optional['NetworkBuffer'] = None
+
+    def __init__(self):
+        os.makedirs(self._BIOLIB_DIR, exist_ok=True)
+        self._docker = BiolibDockerClient.get_docker_client()
+
+    @classmethod
+    def get_instance(cls) -> 'NetworkBuffer':
+        if cls._instance is None:
+            cls._instance = cls()
+        return cls._instance
+
+    def _acquire_lock(self) -> None:
+        start_time = time.time()
+        retry_count = 0
+
+        while time.time() - start_time < self._LOCK_TIMEOUT_SECONDS:
+            try:
+                with open(self._LOCK_FILE, 'x') as lock_file:
+                    lock_info = {
+                        'pid': os.getpid(),
+                        'hostname': socket.gethostname(),
+                        'started_at': time.time(),
+                    }
+                    json.dump(lock_info, lock_file)
+                return
+            except FileExistsError:
+                if retry_count == 0:
+                    self._check_and_remove_stale_lock()
+
+                time.sleep(0.5)
+                retry_count += 1
+
+        raise RuntimeError(
+            f'Failed to acquire network buffer lock after {self._LOCK_TIMEOUT_SECONDS}s: {self._LOCK_FILE}'
+        )
+
+    def _check_and_remove_stale_lock(self) -> None:
+        try:
+            if not os.path.exists(self._LOCK_FILE):
+                return
+
+            lock_mtime = os.path.getmtime(self._LOCK_FILE)
+            lock_age = time.time() - lock_mtime
+
+            if lock_age > self._STALE_LOCK_THRESHOLD_SECONDS:
+                try:
+                    with open(self._LOCK_FILE) as f:
+                        lock_info = json.load(f)
+                        lock_pid = lock_info.get('pid')
+
+                        if lock_pid:
+                            try:
+                                os.kill(lock_pid, 0)
+                                logger_no_user_data.warning(
+                                    f'Lock file is old ({lock_age:.0f}s) but process {lock_pid} is still alive'
+                                )
+                                return
+                            except (OSError, ProcessLookupError):
+                                pass
+
+                except (json.JSONDecodeError, KeyError, ValueError):
+                    pass
+
+                logger_no_user_data.warning(
+                    f'Removing stale lock file (age: {lock_age:.0f}s, threshold: {self._STALE_LOCK_THRESHOLD_SECONDS}s)'
+                )
+                os.remove(self._LOCK_FILE)
+
+        except Exception as error:
+            logger_no_user_data.debug(f'Error checking stale lock: {error}')
+
+    def _release_lock(self) -> None:
+        with contextlib.suppress(FileNotFoundError):
+            os.remove(self._LOCK_FILE)
+
+    def _read_available_networks(self) -> List[str]:
+        if not os.path.exists(self._NETWORKS_FILE):
+            return []
+
+        try:
+            with open(self._NETWORKS_FILE) as f:
+                network_ids = json.load(f)
+                if not isinstance(network_ids, list):
+                    logger_no_user_data.error(
+                        f'Invalid network buffer file format (expected list, got {type(network_ids).__name__})'
+                    )
+                    self._backup_corrupted_file()
+                    return []
+                return network_ids
+        except json.JSONDecodeError as error:
+            logger_no_user_data.error(f'Corrupted network buffer file: {error}')
+            self._backup_corrupted_file()
+            return []
+        except Exception as error:
+            logger_no_user_data.error(f'Failed to read network buffer file: {error}')
+            return []
+
+    def _write_available_networks(self, network_ids: List[str]) -> None:
+        temp_file = f'{self._NETWORKS_FILE}.tmp'
+        try:
+            with open(temp_file, 'w') as f:
+                json.dump(network_ids, f, indent=2)
+                f.flush()
+                os.fsync(f.fileno())
+
+            os.replace(temp_file, self._NETWORKS_FILE)
+        except Exception as error:
+            logger_no_user_data.error(f'Failed to write network buffer file: {error}')
+            with contextlib.suppress(FileNotFoundError):
+                os.remove(temp_file)
+            raise
+
+    def _backup_corrupted_file(self) -> None:
+        try:
+            timestamp = int(time.time())
+            backup_path = f'{self._NETWORKS_FILE}.corrupt-{timestamp}'
+            os.rename(self._NETWORKS_FILE, backup_path)
+            logger_no_user_data.error(f'Backed up corrupted file to {backup_path}')
+        except Exception as error:
+            logger_no_user_data.error(f'Failed to backup corrupted file: {error}')
+
+    def allocate_networks(self, job_id: str, count: int) -> List[Network]:
+        try:
+            self._acquire_lock()
+
+            available_ids = self._read_available_networks()
+            allocated: List[Network] = []
+
+            for _ in range(count):
+                network = None
+
+                while available_ids and network is None:
+                    net_id = available_ids.pop(0)
+                    try:
+                        network = self._docker.networks.get(net_id)
+                        logger_no_user_data.debug(
+                            f'Allocated network {network.id} ({network.name}) from buffer for job {job_id}'
+                        )
+                    except NotFound:
+                        logger_no_user_data.warning(
+                            f'Network {net_id} in buffer file no longer exists in Docker, skipping'
+                        )
+                        network = None
+
+                if network is None:
+                    logger_no_user_data.debug(f'Buffer exhausted, creating network on-the-fly for job {job_id}')
+                    network = self._create_network()
+
+                allocated.append(network)
+
+            self._write_available_networks(available_ids)
+            return allocated
+
+        except RuntimeError as error:
+            logger_no_user_data.warning(f'Lock acquisition failed: {error}. Creating networks on-the-fly.')
+            allocated = []
+            for _ in range(count):
+                network = self._create_network()
+                allocated.append(network)
+            return allocated
+
+        finally:
+            self._release_lock()
+
+    def fill_buffer(self) -> int:
+        try:
+            self._acquire_lock()
+
+            available_ids = self._read_available_networks()
+            current_count = len(available_ids)
+            needed = self.BUFFER_SIZE - current_count
+
+            if needed <= 0:
+                logger_no_user_data.debug(
+                    f'Buffer already has {current_count} available networks (target: {self.BUFFER_SIZE})'
+                )
+                return 0
+
+            logger_no_user_data.debug(
+                f'Filling buffer: current={current_count}, target={self.BUFFER_SIZE}, creating={needed}'
+            )
+
+            created_count = 0
+            for _ in range(needed):
+                try:
+                    network = self._create_network()
+                    if network.id:
+                        available_ids.append(network.id)
+                        created_count += 1
+                        logger_no_user_data.debug(f'Created buffer network {network.id} ({created_count}/{needed})')
+                    else:
+                        logger_no_user_data.error('Created network has no ID, skipping')
+                except Exception as error:
+                    logger_no_user_data.error(f'Failed to create buffer network: {error}')
+                    continue
+
+            self._write_available_networks(available_ids)
+            logger_no_user_data.debug(f'Buffer fill complete: created {created_count} networks')
+            return created_count
+
+        finally:
+            self._release_lock()
+
+    def _create_network(self) -> Network:
+        network = _allocate_network_with_retries(
+            name_prefix=self.NETWORK_NAME_PREFIX,
+            docker_client=self._docker,
+            internal=True,
+            driver='bridge',
+            labels={'biolib-role': 'remote-host-network'},
+        )
+        return network