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

biolib/compute_node/job_worker/job_worker.py

@@ -1,50 +1,63 @@
+import hashlib
 import io
 import json
-import socket
+import multiprocessing
+import os
 import shlex
+import signal
+import socket
 import sys
 import tempfile
 import zipfile
-from time import time
 from queue import Queue
-import multiprocessing
-import os
-import signal
+from time import time
 from types import FrameType
 
 from docker.models.networks import Network  # type: ignore
+from docker.types import IPAMConfig, IPAMPool  # type: ignore
 
-from biolib._internal.http_client import HttpClient
-from biolib.biolib_binary_format.stdout_and_stderr import StdoutAndStderr
-from biolib.compute_node.job_worker.job_legacy_input_wait_timeout_thread import JobLegacyInputWaitTimeout
-from biolib.compute_node.job_worker.job_storage import JobStorage
-from biolib.compute_node.job_worker.large_file_system import LargeFileSystem
-from biolib.biolib_errors import DockerContainerNotFoundDuringExecutionException, BioLibError, \
-    StorageDownloadFailed
-from biolib.compute_node.job_worker.job_max_runtime_timer_thread import JobMaxRuntimeTimerThread
-from biolib.compute_node.remote_host_proxy import RemoteHostProxy
-from biolib.typing_utils import Optional, List, Dict
 from biolib import utils
-from biolib.biolib_api_client import ModuleEnvironment, CreatedJobDict, JobWrapper, Module, AppVersionOnJob, \
-    BiolibApiClient, RemoteHost
+from biolib._internal.http_client import HttpClient
+from biolib.biolib_api_client import (
+    AppVersionOnJob,
+    BiolibApiClient,
+    CreatedJobDict,
+    JobWrapper,
+    Module,
+    ModuleEnvironment,
+)
 from biolib.biolib_api_client.biolib_job_api import BiolibJobApi
+from biolib.biolib_binary_format import (
+    InMemoryIndexableBuffer,
+    ModuleInput,
+    ModuleOutputV2,
+    SavedJob,
+    SystemException,
+    SystemStatusUpdate,
+)
+from biolib.biolib_binary_format.stdout_and_stderr import StdoutAndStderr
 from biolib.biolib_docker_client import BiolibDockerClient
+from biolib.biolib_errors import BioLibError, DockerContainerNotFoundDuringExecutionException, StorageDownloadFailed
 from biolib.biolib_logging import logger, logger_no_user_data
 from biolib.compute_node.job_worker.executors import DockerExecutor
 from biolib.compute_node.job_worker.executors.types import LocalExecutorOptions, StatusUpdate
-from biolib.compute_node.socker_listener_thread import SocketListenerThread
-from biolib.compute_node.socket_sender_thread import SocketSenderThread
+from biolib.compute_node.job_worker.job_legacy_input_wait_timeout_thread import JobLegacyInputWaitTimeout
+from biolib.compute_node.job_worker.job_max_runtime_timer_thread import JobMaxRuntimeTimerThread
+from biolib.compute_node.job_worker.job_storage import JobStorage
+from biolib.compute_node.job_worker.large_file_system import LargeFileSystem
 from biolib.compute_node.job_worker.mappings import Mappings, path_without_first_folder
+from biolib.compute_node.job_worker.network_buffer import NetworkBuffer
 from biolib.compute_node.job_worker.utils import ComputeProcessException, log_disk_and_memory_usage_info
-from biolib.compute_node.utils import get_package_type, SystemExceptionCodes, SystemExceptionCodeMap
-from biolib.biolib_binary_format import SavedJob, SystemStatusUpdate, ModuleInput, SystemException, \
-    ModuleOutputV2, InMemoryIndexableBuffer
+from biolib.compute_node.remote_host_proxy import RemoteHostMapping, RemoteHostProxy, get_static_ip_from_network
+from biolib.compute_node.socker_listener_thread import SocketListenerThread
+from biolib.compute_node.socket_sender_thread import SocketSenderThread
+from biolib.compute_node.utils import SystemExceptionCodeMap, SystemExceptionCodes, get_package_type
+from biolib.typing_utils import Dict, List, Optional
 
 SOCKET_HOST = '127.0.0.1'
 
 
 class JobWorkerProcess(multiprocessing.Process):
-
     # note: this method is run in the parent process
     def __init__(self, socket_port: int, log_level: int):
         super().__init__()
@@ -69,6 +82,13 @@ class JobWorker:
             # handle termination signal from parent
             signal.signal(signal.SIGTERM, self._handle_exit_gracefully)
 
+            try:
+                docker_client = BiolibDockerClient.get_docker_client()
+                networks = docker_client.networks.list()
+                logger_no_user_data.debug(f'Docker networks at JobWorker init: {[net.name for net in networks]}')
+            except Exception as error:
+                logger_no_user_data.debug(f'Failed to list docker networks at init: {error}')
+
             self._socket_port = socket_port
             self._received_messages_queue: Queue = Queue()
             self._messages_to_send_queue: Queue = Queue()
@@ -80,9 +100,9 @@ class JobWorker:
 
             self._remote_host_proxies: List[RemoteHostProxy] = []
             self._internal_network: Optional[Network] = None
-            self._public_network: Optional[Network] = None
             self._executors: List[DockerExecutor] = []
             self.is_cleaning_up: bool = False
+            self._network_buffer = NetworkBuffer.get_instance()
 
             self.job_temporary_dir: Optional[str] = None
 
@@ -91,18 +111,18 @@ class JobWorker:
                 exception,
                 SystemExceptionCodes.FAILED_TO_INIT_COMPUTE_PROCESS_VARIABLES.value,
                 self.send_system_exception,
-                may_contain_user_data=False
+                may_contain_user_data=False,
             ) from exception
 
         if socket_port:
             self._connect_to_parent()
 
     def _handle_exit_gracefully(
-            self,
-            signum: int,
-            frame: Optional[FrameType],  # pylint: disable=unused-argument
+        self,
+        signum: int,
+        frame: Optional[FrameType],  # pylint: disable=unused-argument
     ) -> None:
-        job_id = self._root_job_wrapper["job"]["public_id"] if self._root_job_wrapper else None
+        job_id = self._root_job_wrapper['job']['public_id'] if self._root_job_wrapper else None
         logger_no_user_data.debug(
             f'_JobWorker ({job_id}) got exit signal {signal.Signals(signum).name}'  # pylint: disable=no-member
         )
@@ -187,9 +207,7 @@ class JobWorker:
 
             except Exception as exception:
                 raise ComputeProcessException(
-                    exception,
-                    SystemExceptionCodes.UNKNOWN_COMPUTE_PROCESS_ERROR.value,
-                    self.send_system_exception
+                    exception, SystemExceptionCodes.UNKNOWN_COMPUTE_PROCESS_ERROR.value, self.send_system_exception
                 ) from exception
 
     def _cleanup(self) -> None:
@@ -200,6 +218,8 @@ class JobWorker:
             executor.cleanup()
 
         proxy_count = len(self._remote_host_proxies)
+        cleaned_networks = set()
+
         if proxy_count > 0:
             logger_no_user_data.debug('Cleaning up proxies...')
             proxy_cleanup_start_time = time()
@@ -211,21 +231,37 @@ class JobWorker:
                     logger_no_user_data.error('Failed to clean up remote host proxy')
                     logger.error(exception)
 
+                for network in proxy.get_remote_host_networks():
+                    try:
+                        self._cleanup_network(network)
+                        cleaned_networks.add(network.id)
+                    except Exception as exception:  # pylint: disable=broad-except
+                        logger_no_user_data.error(f'Failed to clean up network {network.name}')
+                        logger.error(exception)
+
             self._remote_host_proxies = []
             logger_no_user_data.debug(f'Cleaned up {proxy_count} proxies in {time() - proxy_cleanup_start_time}')
 
         logger_no_user_data.debug('Cleaning up networks...')
-        self._cleanup_network(self._internal_network)
+        if self._internal_network and self._internal_network.id not in cleaned_networks:
+            self._cleanup_network(self._internal_network)
         self._internal_network = None
-        self._cleanup_network(self._public_network)
-        self._public_network = None
+
+        try:
+            logger_no_user_data.debug('Refilling network buffer...')
+            created = self._network_buffer.fill_buffer()
+            logger_no_user_data.debug(f'Refilled buffer with {created} new networks')
+        except Exception as exception:  # pylint: disable=broad-except
+            logger_no_user_data.error('Failed to refill network buffer')
+            logger.error(exception)
+
         logger_no_user_data.debug('Cleaned up networks...')
 
     @staticmethod
     def _cleanup_network(network: Optional[Network]) -> None:
         if network:
             network_cleanup_start_time = time()
-            network_name = network
+            network_name = network.name
             try:
                 network.remove()
             except Exception as exception:  # pylint: disable=broad-except
@@ -237,10 +273,7 @@ class JobWorker:
     def _handle_save_job_wrapper(self, package: bytes):
         job_wrapper_json_string = SavedJob(package).deserialize()
         job_wrapper: JobWrapper = json.loads(job_wrapper_json_string)
-        BiolibApiClient.initialize(
-            base_url=job_wrapper['BASE_URL'],
-            access_token=job_wrapper['access_token']
-        )
+        BiolibApiClient.initialize(base_url=job_wrapper['BASE_URL'], access_token=job_wrapper['access_token'])
         self._root_job_wrapper = job_wrapper
         if not utils.IS_RUNNING_IN_CLOUD:
             job_wrapper['cloud_job'] = None
@@ -250,6 +283,15 @@ class JobWorker:
         job = job_wrapper['job']
         self._jobs[job['public_id']] = job
 
+        app_version = job['app_version']
+        modules = app_version.get('modules', [])
+        for module in modules:
+            module_ports = module.get('ports', [])
+            if module_ports:
+                logger_no_user_data.debug(
+                    f"Job '{job['public_id']}' module '{module['name']}' has ports: {module_ports}"
+                )
+
         if job['app_version'].get('modules') is not None and BiolibDockerClient.is_docker_running():
             self._start_network_and_remote_host_proxies(job)
 
@@ -259,44 +301,33 @@ class JobWorker:
         app_version = job['app_version']
         job_id = job['public_id']
         remote_hosts = app_version['remote_hosts']
-        if utils.IS_RUNNING_IN_CLOUD:
-            remote_hosts.append(
-                {
-                    'hostname': 'AppCallerProxy',
-                },
-            )
-
         docker_client = BiolibDockerClient.get_docker_client()
         try:
+            name_hash = int(hashlib.sha256(job_id.encode()).hexdigest(), 16)
+            third_octet = name_hash % 256
+            internal_subnet = f'172.29.{third_octet}.0/24'
+
+            ipam_pool = IPAMPool(subnet=internal_subnet)
+            ipam_config = IPAMConfig(pool_configs=[ipam_pool])
+
             self._internal_network = docker_client.networks.create(
                 name=f'biolib-sandboxed-network-{job_id}',
                 internal=True,
                 driver='bridge',
+                ipam=ipam_config,
             )
+            logger_no_user_data.debug(f'Created internal network for job {job_id} with subnet {internal_subnet}')
         except Exception as exception:
             raise ComputeProcessException(
                 exception,
                 SystemExceptionCodes.FAILED_TO_CREATE_DOCKER_NETWORKS.value,
                 self.send_system_exception,
-                may_contain_user_data=False
+                may_contain_user_data=False,
             ) from exception
 
         if len(remote_hosts) > 0:
-            logger_no_user_data.debug(f'Job "{job_id}" creating networks for remote host proxies...')
-            try:
-                self._public_network = docker_client.networks.create(
-                    name=f'biolib-proxy-network-{job_id}',
-                    internal=False,
-                    driver='bridge',
-                )
-            except Exception as exception:
-                raise ComputeProcessException(
-                    exception,
-                    SystemExceptionCodes.FAILED_TO_CREATE_DOCKER_NETWORKS.value,
-                    self.send_system_exception,
-                    may_contain_user_data=False
-                ) from exception
-            logger_no_user_data.debug(f'Job "{job_id}" starting proxies for remote hosts: {remote_hosts}')
+            logger_no_user_data.debug(f'Job "{job_id}" starting proxy for remote hosts: {remote_hosts}')
+            created_networks: List[Network] = []
             try:
                 hostname_to_ports: Dict[str, List[int]] = {}
                 for remote_host in remote_hosts:
@@ -312,33 +343,67 @@ class JobWorker:
                     else:
                         hostname_to_ports[hostname] = [port]
 
-                for hostname, ports in hostname_to_ports.items():
+                remote_host_mappings: List[RemoteHostMapping] = []
+                networks = self._network_buffer.allocate_networks(job_id, len(hostname_to_ports))
+                created_networks.extend(networks)
+
+                for (hostname, ports), network in zip(hostname_to_ports.items(), networks):
+                    static_ip = get_static_ip_from_network(network, offset=2)
+
+                    mapping = RemoteHostMapping(
+                        hostname=hostname,
+                        ports=ports,
+                        network=network,
+                        static_ip=static_ip,
+                    )
+                    remote_host_mappings.append(mapping)
+
+                if remote_host_mappings:
                     remote_host_proxy = RemoteHostProxy(
-                        RemoteHost(hostname=hostname),
-                        self._public_network,
-                        self._internal_network,
-                        job_id,
-                        ports,
+                        remote_host_mappings=remote_host_mappings,
+                        job=job,
+                        app_caller_network=None,
                     )
                     remote_host_proxy.start()
                     self._remote_host_proxies.append(remote_host_proxy)
+                    num_hosts = len(remote_host_mappings)
+                    logger_no_user_data.debug(f'Started single proxy container for {num_hosts} remote hosts')
 
             except Exception as exception:
+                for network in created_networks:
+                    self._cleanup_network(network)
+
                 raise ComputeProcessException(
                     exception,
                     SystemExceptionCodes.FAILED_TO_START_REMOTE_HOST_PROXIES.value,
                     self.send_system_exception,
-                    may_contain_user_data=False
+                    may_contain_user_data=False,
                 ) from exception
 
-            logger_no_user_data.debug(f'Job "{job_id}" startup of remote host proxies completed')
+        if utils.IS_RUNNING_IN_CLOUD:
+            try:
+                app_caller_proxy = RemoteHostProxy(
+                    remote_host_mappings=[],
+                    job=job,
+                    app_caller_network=self._internal_network,
+                )
+                app_caller_proxy.start()
+                self._remote_host_proxies.append(app_caller_proxy)
+                logger_no_user_data.debug('Started app caller proxy')
+            except Exception as exception:
+                raise ComputeProcessException(
+                    exception,
+                    SystemExceptionCodes.FAILED_TO_START_REMOTE_HOST_PROXIES.value,
+                    self.send_system_exception,
+                    may_contain_user_data=False,
+                ) from exception
 
     def _run_app_version(
-            self,
-            app_version_id: str,
-            module_input_path: str,
-            caller_job: CreatedJobDict,
-            main_module_output_path: str,
+        self,
+        app_version_id: str,
+        module_input_path: str,
+        caller_job: CreatedJobDict,
+        main_module_output_path: str,
     ) -> None:
         job: CreatedJobDict = BiolibJobApi.create(app_version_id, caller_job=caller_job['public_id'])
         self._jobs[job['public_id']] = job
@@ -357,17 +422,17 @@ class JobWorker:
         root_job_id = root_job['public_id']
         if job.get('arguments_override_command') and not job['app_version']['app']['allow_client_side_execution']:
             raise ComputeProcessException(
-                Exception("Command override not allowed"),
+                Exception('Command override not allowed'),
                 SystemExceptionCodes.COMMAND_OVERRIDE_NOT_ALLOWED.value,
-                self.send_system_exception
+                self.send_system_exception,
             )
 
         modules = job['app_version'].get('modules')
         if not modules:
             raise ComputeProcessException(
-                Exception("No modules found on job"),
+                Exception('No modules found on job'),
                 SystemExceptionCodes.NO_MODULES_FOUND_ON_JOB.value,
-                self.send_system_exception
+                self.send_system_exception,
             )
 
         main_module = self._get_module_from_name(modules, module_name='main')
@@ -420,8 +485,8 @@ class JobWorker:
             log_disk_and_memory_usage_info()
 
     def _run_module(
-            self,
-            options: LocalExecutorOptions,
+        self,
+        options: LocalExecutorOptions,
     ) -> None:
         module = options['module']
         job_id = options['job']['public_id']
@@ -434,7 +499,7 @@ class JobWorker:
             if not self.job_temporary_dir:
                 raise BioLibError('Undefined job_temporary_dir')
             logger_no_user_data.debug(f'Job "{job_id}" starting child job...')
-            with open(module_input_path,'rb') as fp:
+            with open(module_input_path, 'rb') as fp:
                 module_input_serialized = fp.read()
             module_input = ModuleInput(module_input_serialized).deserialize()
             module_input_with_runtime_zip = self._add_runtime_zip_and_command_to_module_input(options, module_input)
@@ -443,7 +508,7 @@ class JobWorker:
                 arguments=module_input_with_runtime_zip['arguments'],
                 files=module_input_with_runtime_zip['files'],
             )
-            module_input_path_new = os.path.join(self.job_temporary_dir, "runtime." + JobStorage.module_input_file_name)
+            module_input_path_new = os.path.join(self.job_temporary_dir, 'runtime.' + JobStorage.module_input_file_name)
             open(module_input_path_new, 'wb').write(module_input_with_runtime_zip_serialized)
             return self._run_app_version(
                 module['image_uri'],
@@ -460,7 +525,7 @@ class JobWorker:
                     exception,
                     SystemExceptionCodes.FAILED_TO_INITIALIZE_DOCKER_EXECUTOR.value,
                     self.send_system_exception,
-                    may_contain_user_data=False
+                    may_contain_user_data=False,
                 ) from exception
         else:
             err_string = f'Job "{job_id}" hit unsupported module environment "{module["environment"]}"'
@@ -485,7 +550,7 @@ class JobWorker:
                 exception,
                 SystemExceptionCodes.FAILED_TO_CONNECT_TO_WORKER_THREAD_SOCKET.value,
                 self.send_system_exception,
-                may_contain_user_data=False
+                may_contain_user_data=False,
             ) from exception
 
         try:
@@ -496,7 +561,7 @@ class JobWorker:
                 exception,
                 SystemExceptionCodes.FAILED_TO_START_SENDER_THREAD_OR_RECEIVER_THREAD.value,
                 self.send_system_exception,
-                may_contain_user_data=False
+                may_contain_user_data=False,
             ) from exception
 
     # TODO: move this mapping logic to the ModuleInput class
@@ -524,7 +589,7 @@ class JobWorker:
                 exception,
                 SystemExceptionCodes.FAILED_TO_CREATE_NEW_JOB.value,
                 self.send_system_exception,
-                may_contain_user_data=False
+                may_contain_user_data=False,
             ) from exception
 
         return module_input
@@ -550,7 +615,7 @@ class JobWorker:
                     exception,
                     SystemExceptionCodes.FAILED_TO_DOWNLOAD_RUNTIME_ZIP.value,
                     self.send_system_exception,
-                    may_contain_user_data=False
+                    may_contain_user_data=False,
                 ) from exception
             finally:
                 download_time = time() - start_time
@@ -596,7 +661,7 @@ class JobWorker:
                 exception,
                 SystemExceptionCodes.FAILED_TO_SEND_STATUS_UPDATE.value,
                 self.send_system_exception,
-                may_contain_user_data=False
+                may_contain_user_data=False,
             ) from exception
 
     def _run_root_job(self, module_input_path: str) -> str:

biolib/compute_node/job_worker/large_file_system.py

@@ -56,17 +56,13 @@ class LargeFileSystem:
         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(
@@ -77,7 +73,7 @@ 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
 

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'
+    )