pybiolib 1.2.911__py3-none-any.whl → 1.2.1642__py3-none-any.whl

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

biolib/compute_node/job_worker/utilization_reporter_thread.py

@@ -1,7 +1,7 @@
 import threading
 import time
 import subprocess
-from datetime import datetime
+from datetime import datetime, timezone
 
 from docker.models.containers import Container  # type: ignore
 
@@ -173,7 +173,7 @@ class UtilizationReporterThread(threading.Thread):
             gpu_max_usage_in_percent=gpu_max_usage_in_percent,
             memory_average_usage_in_percent=memory_average_usage_in_percent,
             memory_max_usage_in_percent=memory_max_usage_in_percent,
-            recorded_at=datetime.utcnow().isoformat(),
+            recorded_at=datetime.now(timezone.utc).isoformat(),
             sampling_period_in_milliseconds=self._sampling_period_in_milliseconds * self._samples_between_writes,
         )
 

biolib/compute_node/remote_host_proxy.py

@@ -1,14 +1,13 @@
 import base64
 import io
-import subprocess
+import ipaddress
 import tarfile
 import time
 from urllib.parse import urlparse
 
-from docker.errors import ImageNotFound  # type: ignore
-from docker.models.containers import Container  # type: ignore
-from docker.models.images import Image  # type: ignore
-from docker.models.networks import Network  # type: ignore
+from docker.models.containers import Container
+from docker.models.networks import Network
+from docker.types import EndpointConfig
 
 from biolib import utils
 from biolib.biolib_api_client import BiolibApiClient, RemoteHost
@@ -16,7 +15,22 @@ from biolib.biolib_docker_client import BiolibDockerClient
 from biolib.biolib_errors import BioLibError
 from biolib.biolib_logging import logger_no_user_data
 from biolib.compute_node.cloud_utils import CloudUtils
-from biolib.typing_utils import List, Optional
+from biolib.compute_node.utils import BIOLIB_PROXY_NETWORK_NAME
+from biolib.compute_node.webserver.proxy_utils import get_biolib_nginx_proxy_image
+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
@@ -24,78 +38,98 @@ 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,
-        public_network: Network,
-        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._public_network: Network = public_network
-        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._enclave_traffic_forwarder_processes: List[subprocess.Popen] = []
         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]] = (
+            None
+            if not self.is_app_caller_proxy
+            else {
+                BIOLIB_PROXY_NETWORK_NAME: docker.api.create_endpoint_config(
+                    aliases=[f'biolib-app-caller-proxy-{self._job_uuid}']
+                )
+            }
+        )
+
         for index in range(3):
             logger_no_user_data.debug(f'Attempt {index} at creating RemoteHostProxy container "{self._name}"...')
             try:
                 self._container = docker.containers.create(
                     detach=True,
-                    image=self._get_biolib_remote_host_proxy_image(),
+                    image=get_biolib_nginx_proxy_image(),
                     name=self._name,
-                    network=self._public_network.name,
+                    network=BIOLIB_PROXY_NETWORK_NAME,
+                    networking_config=networking_config,
                 )
                 break
-            except Exception as error:  # pylint: disable=broad-except
+            except Exception as error:
                 logger_no_user_data.exception(f'Failed to create container "{self._name}" hit error: {error}')
 
             logger_no_user_data.debug('Sleeping before re-trying container creation...')
             time.sleep(3)
 
-        if not self._container:
+        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()
 
@@ -121,28 +155,7 @@ class RemoteHostProxy:
         if self._container:
             self._container.remove(force=True)
 
-        for process in self._enclave_traffic_forwarder_processes:
-            process.terminate()
-
-    def _get_biolib_remote_host_proxy_image(self) -> Image:
-        if utils.IS_RUNNING_IN_CLOUD:
-            try:
-                logger_no_user_data.debug('Getting local Docker image for remote host proxy')
-                return self._docker.images.get('biolib-remote-host-proxy:latest')
-            except ImageNotFound:
-                logger_no_user_data.debug(
-                    'Local Docker image for remote host proxy not available. Falling back to public image...'
-                )
-
-        public_image_uri = 'public.ecr.aws/h5y4b3l1/biolib-remote-host-proxy:latest'
-        try:
-            logger_no_user_data.debug('Getting public Docker image for remote host proxy')
-            return self._docker.images.get(public_image_uri)
-        except ImageNotFound:
-            logger_no_user_data.debug('Pulling public Docker image for remote host proxy')
-            return self._docker.images.pull(public_image_uri)
-
-    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')
 
@@ -273,6 +286,16 @@ http {{
             proxy_ssl_server_name    on;
         }}
 
+        location ~ "^/api/auth/oauth-token-exchange/$" {{
+            # Note: Using $1 here as URI part from regex must be used for proxy_pass
+            proxy_pass               https://$upstream_hostname/api/auth/oauth-token-exchange/$1;
+            proxy_set_header         authorization "";
+            proxy_set_header         compute-node-auth-token "{compute_node_auth_token}";
+            proxy_set_header         job-uuid "{self._job_uuid}";
+            proxy_set_header         cookie "";
+            proxy_ssl_server_name    on;
+        }}
+
         location /api/lfs/ {{
             proxy_pass               https://$upstream_hostname$request_uri;
             proxy_set_header         authorization "";
@@ -318,6 +341,22 @@ http {{
             proxy_ssl_server_name    on;
         }}
 
+        location /api/proxy/index/ {{
+            proxy_pass               https://$upstream_hostname$request_uri;
+            proxy_set_header         authorization "Basic {biolib_index_basic_auth_base64}";
+            proxy_set_header         cookie "";
+            proxy_ssl_server_name    on;
+        }}
+
+        location ~* "^/api/accounts/(?<account_id>[a-z0-9-]{{36}})/metrics/jobs/$" {{
+            proxy_pass               https://$upstream_hostname/api/accounts/$account_id/metrics/jobs/$is_args$args;
+            proxy_set_header         authorization "";
+            proxy_set_header         compute-node-auth-token "{compute_node_auth_token}";
+            proxy_set_header         job-uuid "{self._job_uuid}";
+            proxy_set_header         cookie "";
+            proxy_ssl_server_name    on;
+        }}
+
         location /api/ {{
             proxy_pass               https://$upstream_hostname$request_uri;
             proxy_set_header         authorization "";
@@ -357,28 +396,57 @@ http {{
             return 404 "Not found";
         }}
     }}
+
+    server {{
+        listen       1080;
+        resolver 127.0.0.11 ipv6=off valid=30s;
+
+        if ($http_biolib_result_uuid != "{self._job_uuid}") {{
+            return 403 "Invalid or missing biolib-result-uuid header";
+        }}
+
+        if ($http_biolib_result_port = "") {{
+            return 400 "Missing biolib-result-port header";
+        }}
+
+        location / {{
+            proxy_pass               http://main:$http_biolib_result_port$request_uri;
+            proxy_set_header         biolib-result-uuid "";
+            proxy_set_header         biolib-result-port "";
+            proxy_pass_request_headers on;
+        }}
+    }}
 }}
 """
         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 += """

biolib/compute_node/utils.py

@@ -4,6 +4,8 @@ from enum import Enum
 
 from biolib.biolib_logging import logger
 
+BIOLIB_PROXY_NETWORK_NAME = 'biolib-proxy-network'
+
 
 def get_package_type(package):
     package_type = int.from_bytes(package[1:2], 'big')