pybiolib 1.2.1298__tar.gz → 1.2.1304__tar.gz

{pybiolib-1.2.1298 → pybiolib-1.2.1304}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pybiolib
-Version: 1.2.1298
+Version: 1.2.1304
 Summary: BioLib Python Client
 License: MIT
 License-File: LICENSE

{pybiolib-1.2.1298 → pybiolib-1.2.1304}/biolib/compute_node/job_worker/executors/docker_executor.py

@@ -367,9 +367,11 @@ class DockerExecutor:
                 )
 
             logger_no_user_data.debug(f'Job "{job_uuid}" initializing Docker container. Getting IPs for proxies...')
+
+            networks_to_connect = []
             for proxy in self._options['remote_host_proxies']:
-                proxy_ip = proxy.get_ip_address_on_network(internal_network)
                 if proxy.is_app_caller_proxy:
+                    proxy_ip = proxy.get_ip_address_on_network(internal_network)
                     logger_no_user_data.debug('Found app caller proxy, setting both base URLs in compute container')
                     environment_vars.update(
                         {
@@ -383,7 +385,11 @@ class DockerExecutor:
                         }
                     )
                 else:
-                    extra_hosts[proxy.hostname] = proxy_ip
+                    extra_hosts.update(proxy.get_hostname_to_ip_mapping())
+
+                for network in proxy.get_remote_host_networks():
+                    if network != internal_network:
+                        networks_to_connect.append(network)
 
             logger_no_user_data.debug(f'Job "{job_uuid}" initializing Docker container. Constructing container args...')
             create_container_args = {
@@ -436,6 +442,17 @@ class DockerExecutor:
 
             logger_no_user_data.debug(f'Job "{job_uuid}" initializing Docker container. Creating container...')
             self._docker_container = docker_client.containers.create(**create_container_args)
+
+            if networks_to_connect:
+                network_connection_start = time.time()
+                for network in networks_to_connect:
+                    network.connect(self._docker_container.id)
+                    logger_no_user_data.debug(f'Connected app container to network {network.name}')
+                network_connection_time = time.time() - network_connection_start
+                logger_no_user_data.debug(
+                    f'Connected app container to {len(networks_to_connect)} networks in {network_connection_time:.2f}s'
+                )
+
             logger_no_user_data.debug(f'Job "{job_uuid}" finished initializing Docker container.')
         except Exception as exception:
             raise ComputeProcessException(

{pybiolib-1.2.1298 → pybiolib-1.2.1304}/biolib/compute_node/job_worker/job_worker.py

@@ -1,44 +1,56 @@
 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
 
-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.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'
 
@@ -199,6 +211,7 @@ 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()
@@ -210,11 +223,20 @@ 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
         logger_no_user_data.debug('Cleaned up networks...')
 
@@ -265,13 +287,6 @@ 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:
             self._internal_network = docker_client.networks.create(
@@ -288,7 +303,7 @@ class JobWorker:
             ) from exception
 
         if len(remote_hosts) > 0:
-            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}')
             try:
                 hostname_to_ports: Dict[str, List[int]] = {}
                 for remote_host in remote_hosts:
@@ -304,15 +319,44 @@ class JobWorker:
                     else:
                         hostname_to_ports[hostname] = [port]
 
+                remote_host_mappings: List[RemoteHostMapping] = []
                 for hostname, ports in hostname_to_ports.items():
+                    network_index = len(remote_host_mappings)
+                    subnet = f'172.28.{network_index}.0/24'
+
+                    ipam_pool = IPAMPool(subnet=subnet)
+                    ipam_config = IPAMConfig(pool_configs=[ipam_pool])
+
+                    network = docker_client.networks.create(
+                        name=f'biolib-remote-host-network-{job_id}-{network_index}',
+                        internal=True,
+                        driver='bridge',
+                        ipam=ipam_config,
+                    )
+
+                    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)
+                    logger_no_user_data.debug(
+                        f'Created network {network.name} with static IP {static_ip} for host {hostname}'
+                    )
+
+                if remote_host_mappings:
                     remote_host_proxy = RemoteHostProxy(
-                        remote_host=RemoteHost(hostname=hostname),
-                        internal_network=self._internal_network,
+                        remote_host_mappings=remote_host_mappings,
                         job_id=job_id,
-                        ports=ports,
+                        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:
                 raise ComputeProcessException(
@@ -322,7 +366,23 @@ class JobWorker:
                     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_id=job_id,
+                    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,

{pybiolib-1.2.1298 → pybiolib-1.2.1304}/biolib/compute_node/remote_host_proxy.py

@@ -1,5 +1,6 @@
 import base64
 import io
+import ipaddress
 import tarfile
 import time
 from urllib.parse import urlparse
@@ -19,52 +20,73 @@ from biolib.compute_node.webserver.proxy_utils import get_biolib_nginx_proxy_ima
 from biolib.typing_utils import Dict, List, Optional
 
 
+def get_static_ip_from_network(network: Network, offset: int = 2) -> str:
+    ipam_config = network.attrs['IPAM']['Config']
+    if not ipam_config:
+        raise BioLibError(f'Network {network.name} has no IPAM configuration')
+
+    subnet_str = ipam_config[0]['Subnet']
+    subnet = ipaddress.ip_network(subnet_str, strict=False)
+
+    static_ip = str(subnet.network_address + offset)
+
+    return static_ip
+
+
 # Prepare for remote hosts with specified port
 class RemoteHostExtended(RemoteHost):
     ports: List[int]
 
 
+class RemoteHostMapping:
+    def __init__(self, hostname: str, ports: List[int], network: Network, static_ip: str):
+        self.hostname = hostname
+        self.ports = ports
+        self.network = network
+        self.static_ip = static_ip
+
+
 class RemoteHostProxy:
     def __init__(
         self,
-        remote_host: RemoteHost,
-        internal_network: Optional[Network],
+        remote_host_mappings: List[RemoteHostMapping],
         job_id: str,
-        ports: List[int],
+        app_caller_network: Optional[Network] = None,
     ):
-        self.is_app_caller_proxy = remote_host['hostname'] == 'AppCallerProxy'
-        self._remote_host: RemoteHostExtended = RemoteHostExtended(hostname=remote_host['hostname'], ports=ports)
-        self._internal_network: Optional[Network] = internal_network
+        self._remote_host_mappings = remote_host_mappings
+        self._app_caller_network = app_caller_network
+        self.is_app_caller_proxy = app_caller_network is not None
 
         if not job_id:
             raise Exception('RemoteHostProxy missing argument "job_id"')
 
-        self._name = f'biolib-remote-host-proxy-{job_id}-{self.hostname}'
+        suffix = '-AppCallerProxy' if app_caller_network else ''
+        self._name = f'biolib-remote-host-proxy-{job_id}{suffix}'
         self._job_uuid = job_id
         self._container: Optional[Container] = None
         self._docker = BiolibDockerClient().get_docker_client()
 
-    @property
-    def hostname(self) -> str:
-        return self._remote_host['hostname']
+    def get_hostname_to_ip_mapping(self) -> Dict[str, str]:
+        return {mapping.hostname: mapping.static_ip for mapping in self._remote_host_mappings}
+
+    def get_remote_host_networks(self) -> List[Network]:
+        networks = [mapping.network for mapping in self._remote_host_mappings]
+        return networks
 
     def get_ip_address_on_network(self, network: Network) -> str:
         if not self._container:
-            raise Exception('RemoteHostProxy not yet started')
+            raise BioLibError('RemoteHostProxy not yet started')
 
         container_networks = self._container.attrs['NetworkSettings']['Networks']
         if network.name in container_networks:
             ip_address: str = container_networks[network.name]['IPAddress']
+            if not ip_address:
+                raise BioLibError(f'No IP address found for network {network.name}')
             return ip_address
 
-        raise Exception(f'RemoteHostProxy not connected to network {network.name}')
+        raise BioLibError(f'RemoteHostProxy not connected to network {network.name}')
 
     def start(self) -> None:
-        # TODO: Implement nice error handling in this method
-
-        upstream_server_name = self._remote_host['hostname']
-        upstream_server_ports = self._remote_host['ports']
-
         docker = BiolibDockerClient.get_docker_client()
 
         networking_config: Optional[Dict[str, EndpointConfig]] = (
@@ -97,13 +119,19 @@ class RemoteHostProxy:
         if not self._container or not self._container.id:
             raise BioLibError(f'Exceeded re-try limit for creating container {self._name}')
 
-        self._write_nginx_config_to_container(
-            upstream_server_name,
-            upstream_server_ports,
-        )
+        for mapping in self._remote_host_mappings:
+            mapping.network.connect(self._container.id, ipv4_address=mapping.static_ip)
+            logger_no_user_data.debug(
+                f'Connected proxy to network {mapping.network.name} with static IP {mapping.static_ip}'
+            )
+
+        if self._app_caller_network:
+            self._app_caller_network.connect(self._container.id)
+            logger_no_user_data.debug(
+                f'Connected app caller proxy to network {self._app_caller_network.name}'
+            )
 
-        if self._internal_network:
-            self._internal_network.connect(self._container.id)
+        self._write_nginx_config_to_container()
 
         self._container.start()
 
@@ -129,8 +157,7 @@ class RemoteHostProxy:
         if self._container:
             self._container.remove(force=True)
 
-
-    def _write_nginx_config_to_container(self, upstream_server_name: str, upstream_server_ports: List[int]) -> None:
+    def _write_nginx_config_to_container(self) -> None:
         if not self._container:
             raise Exception('RemoteHostProxy container not defined when attempting to write NGINX config')
 
@@ -378,25 +405,34 @@ http {{
 }}
 """
         else:
+            port_to_mappings: Dict[int, List[RemoteHostMapping]] = {}
+            for mapping in self._remote_host_mappings:
+                for port in mapping.ports:
+                    if port not in port_to_mappings:
+                        port_to_mappings[port] = []
+                    port_to_mappings[port].append(mapping)
+
             nginx_config = """
 events {}
 error_log /dev/stdout info;
 stream {
     resolver 127.0.0.11 valid=30s;"""
-            for idx, upstream_server_port in enumerate(upstream_server_ports):
+
+            for port, mappings in port_to_mappings.items():
                 nginx_config += f"""
-    map "" $upstream_{idx} {{
-        default {upstream_server_name}:{upstream_server_port};
-    }}
+    map $server_addr $backend_{port} {{"""
+                for mapping in mappings:
+                    nginx_config += f'\n        {mapping.static_ip} {mapping.hostname}:{port};'
 
+                nginx_config += f"""
+    }}
     server {{
-        listen          {self._remote_host['ports'][idx]};
-        proxy_pass      $upstream_{idx};
+        listen 0.0.0.0:{port};
+        proxy_pass $backend_{port};
     }}
-
     server {{
-        listen          {self._remote_host['ports'][idx]} udp;
-        proxy_pass      $upstream_{idx};
+        listen 0.0.0.0:{port} udp;
+        proxy_pass $backend_{port};
     }}"""
 
             nginx_config += """

{pybiolib-1.2.1298 → pybiolib-1.2.1304}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pybiolib"
-version = "1.2.1298"
+version = "1.2.1304"
 description = "BioLib Python Client"
 readme = "PYPI_README.md"
 license = "MIT"