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

biolib/compute_node/webserver/compute_node_results_proxy.py

@@ -0,0 +1,188 @@
+import io
+import os
+import tarfile
+import threading
+import time
+
+from docker.models.containers import Container
+
+from biolib.biolib_docker_client import BiolibDockerClient
+from biolib.biolib_errors import BioLibError
+from biolib.biolib_logging import logger, logger_no_user_data
+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 Optional
+
+
+class LogStreamingThread(threading.Thread):
+    def __init__(self, container, container_name: str):
+        super().__init__(daemon=True)
+        self._container = container
+        self._container_name = container_name
+        self._stop_event = threading.Event()
+
+    def run(self) -> None:
+        try:
+            logger_no_user_data.debug(f'Starting log streaming for container "{self._container_name}"')
+            log_stream = self._container.logs(follow=True, stream=True)
+            for log_line in log_stream:
+                if self._stop_event.is_set():
+                    break
+                if log_line:
+                    logger.debug(f'ComputeNodeResultsProxy | {log_line.decode("utf-8", errors="replace").rstrip()}')
+        except Exception as error:
+            logger_no_user_data.debug(f'Log streaming for container "{self._container_name}" ended: {error}')
+
+    def stop(self) -> None:
+        self._stop_event.set()
+
+
+class ComputeNodeResultsProxy:
+    _instance: Optional['ComputeNodeResultsProxy'] = None
+
+    def __init__(self, tls_pem_certificate_path: str, tls_pem_key_path: str):
+        assert tls_pem_certificate_path, 'tls_pem_certificate_path is required'
+        assert tls_pem_key_path, 'tls_pem_key_path is required'
+        self._name = 'biolib-compute-node-results-proxy'
+        self._container: Optional[Container] = None
+        self._docker = BiolibDockerClient().get_docker_client()
+        self._tls_pem_certificate_path = tls_pem_certificate_path
+        self._tls_pem_key_path = tls_pem_key_path
+        self._log_streaming_thread: Optional[LogStreamingThread] = None
+
+    @staticmethod
+    def start_proxy(tls_pem_certificate_path: str, tls_pem_key_path: str) -> None:
+        abs_cert_path = os.path.abspath(tls_pem_certificate_path)
+        abs_key_path = os.path.abspath(tls_pem_key_path)
+        assert os.path.exists(abs_cert_path), f'TLS certificate file does not exist: {abs_cert_path}'
+        assert os.path.exists(abs_key_path), f'TLS key file does not exist: {abs_key_path}'
+
+        if ComputeNodeResultsProxy._instance is None:
+            logger_no_user_data.debug(
+                f'Creating ComputeNodeResultsProxy instance with cert: {abs_cert_path}, key: {abs_key_path}'
+            )
+            ComputeNodeResultsProxy._instance = ComputeNodeResultsProxy(abs_cert_path, abs_key_path)
+            ComputeNodeResultsProxy._instance._start()  # pylint: disable=protected-access
+
+    @staticmethod
+    def stop_proxy() -> None:
+        if ComputeNodeResultsProxy._instance is not None:
+            ComputeNodeResultsProxy._instance._terminate()  # pylint: disable=protected-access
+            ComputeNodeResultsProxy._instance = None
+
+    def _start(self) -> None:
+        docker = BiolibDockerClient.get_docker_client()
+
+        for index in range(3):
+            logger_no_user_data.debug(
+                f'Attempt {index} at creating ComputeNodeResultsProxy container "{self._name}"...'
+            )
+            try:
+                self._container = docker.containers.create(
+                    detach=True,
+                    image=get_biolib_nginx_proxy_image(),
+                    name=self._name,
+                    network=BIOLIB_PROXY_NETWORK_NAME,
+                    ports={'443/tcp': 20443},
+                    volumes={
+                        self._tls_pem_certificate_path: {'bind': '/etc/ssl/certs/certificate.pem', 'mode': 'ro'},
+                        self._tls_pem_key_path: {'bind': '/etc/ssl/private/private_key.pem', 'mode': 'ro'},
+                    },
+                )
+                break
+            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 or not self._container.id:
+            raise BioLibError(f'Exceeded re-try limit for creating container {self._name}')
+
+        self._write_nginx_config_to_container()
+        self._container.start()
+
+        logger_no_user_data.debug(f'Waiting for container "{self._name}" to be ready...')
+        proxy_is_ready = False
+        for retry_count in range(1, 5):
+            time.sleep(0.5 * retry_count)
+            container_logs = self._container.logs()
+            if b'ready for start up\n' in container_logs or b'start worker process ' in container_logs:
+                proxy_is_ready = True
+                break
+
+        if not proxy_is_ready:
+            logger_no_user_data.error('ComputeNodeResultsProxy did not start properly.')
+            self._terminate()
+            raise Exception('ComputeNodeResultsProxy did not start properly')
+
+        self._container.reload()
+        logger.debug(f'ComputeNodeResultsProxy container "{self._name}" started with ID: {self._container.id}')
+
+        self._log_streaming_thread = LogStreamingThread(self._container, self._name)
+        self._log_streaming_thread.start()
+        logger_no_user_data.debug(f'Started log streaming for container "{self._name}"')
+
+    def _terminate(self):
+        logger_no_user_data.debug(f'Terminating ComputeNodeResultsProxy container "{self._name}"')
+        if self._log_streaming_thread:
+            self._log_streaming_thread.stop()
+            self._log_streaming_thread = None
+
+        logger.debug(f'Docker container removal temporarily disabled for debugging purposes (container "{self._name}")')
+        # TODO: Figure if we need to remove the container or keep it for debugging purposes
+        # if self._container:
+        #     self._container.remove(force=True)
+
+    def _write_nginx_config_to_container(self) -> None:
+        if not self._container:
+            raise Exception('ComputeNodeResultsProxy container not defined when attempting to write NGINX config')
+
+        docker = BiolibDockerClient.get_docker_client()
+
+        nginx_config = """
+events {
+  worker_connections  1024;
+}
+
+http {
+    client_max_body_size 0;
+    resolver 127.0.0.11 ipv6=off valid=30s;
+
+    server {
+        listen 443 ssl http2 default_server;
+
+        ssl_certificate /etc/ssl/certs/certificate.pem;
+        ssl_certificate_key /etc/ssl/private/private_key.pem;
+        ssl_protocols TLSv1.2 TLSv1.3;
+
+        location / {
+            if ($http_biolib_result_uuid = "") {
+                return 400 "Missing biolib-result-uuid header";
+            }
+
+            if ($http_biolib_result_port = "") {
+                return 400 "Missing biolib-result-port header";
+            }
+
+            set $target_host "biolib-app-caller-proxy-$http_biolib_result_uuid";
+            proxy_pass http://$target_host:1080$request_uri;
+            proxy_set_header biolib-result-uuid $http_biolib_result_uuid;
+            proxy_set_header biolib-result-port $http_biolib_result_port;
+            proxy_pass_request_headers on;
+        }
+    }
+}
+"""
+
+        nginx_config_bytes = nginx_config.encode()
+        tarfile_in_memory = io.BytesIO()
+        with tarfile.open(fileobj=tarfile_in_memory, mode='w:gz') as tar:
+            info = tarfile.TarInfo('/nginx.conf')
+            info.size = len(nginx_config_bytes)
+            tar.addfile(info, io.BytesIO(nginx_config_bytes))
+
+        tarfile_bytes = tarfile_in_memory.getvalue()
+        tarfile_in_memory.close()
+        logger_no_user_data.debug('Writing NGINX configuration to ComputeNodeResultsProxy container')
+        docker.api.put_archive(self._container.id, '/etc/nginx', tarfile_bytes)

biolib/compute_node/webserver/proxy_utils.py

@@ -0,0 +1,28 @@
+from docker.errors import ImageNotFound
+from docker.models.images import Image
+
+from biolib import utils
+from biolib.biolib_docker_client import BiolibDockerClient
+from biolib.biolib_logging import logger_no_user_data
+from biolib.typing_utils import cast
+
+
+def get_biolib_nginx_proxy_image() -> Image:
+    docker = BiolibDockerClient().get_docker_client()
+
+    if utils.IS_RUNNING_IN_CLOUD:
+        try:
+            logger_no_user_data.debug('Getting local Docker image for nginx proxy')
+            return cast(Image, docker.images.get('biolib-remote-host-proxy:latest'))
+        except ImageNotFound:
+            logger_no_user_data.debug(
+                'Local Docker image for nginx 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 nginx proxy')
+        return cast(Image, docker.images.get(public_image_uri))
+    except ImageNotFound:
+        logger_no_user_data.debug('Pulling public Docker image for nginx proxy')
+        return cast(Image, docker.images.pull(public_image_uri))

biolib/compute_node/webserver/webserver.py

@@ -1,19 +1,24 @@
 # pylint: disable=unsubscriptable-object
 
 import json
+import logging
 import os
-import time
 import tempfile
-import logging
-from flask import Flask, request, Response, jsonify
+import time
+
+from docker.models.networks import Network  # type: ignore
+from flask import Flask, Response, jsonify, request
 
 from biolib import utils
 from biolib.biolib_api_client import BiolibApiClient
 from biolib.biolib_binary_format import SavedJob
-from biolib.compute_node.webserver import webserver_utils
+from biolib.biolib_docker_client import BiolibDockerClient
+from biolib.biolib_logging import TRACE, logger, logger_no_user_data
 from biolib.compute_node.cloud_utils.cloud_utils import CloudUtils
+from biolib.compute_node.utils import BIOLIB_PROXY_NETWORK_NAME
+from biolib.compute_node.webserver import webserver_utils
+from biolib.compute_node.webserver.compute_node_results_proxy import ComputeNodeResultsProxy
 from biolib.compute_node.webserver.gunicorn_flask_application import GunicornFlaskApplication
-from biolib.biolib_logging import logger, TRACE
 from biolib.compute_node.webserver.webserver_utils import get_job_compute_state_or_404
 from biolib.typing_utils import Optional
 
@@ -102,15 +107,14 @@ def status(job_id):
     response_data['is_completed'] = compute_state['is_completed']
 
     if current_status['stdout_and_stderr_packages_b64']:
-        compute_state['streamed_logs_packages_b64'] = compute_state['streamed_logs_packages_b64'] + \
-                                                      current_status['stdout_and_stderr_packages_b64']
+        compute_state['streamed_logs_packages_b64'] = (
+            compute_state['streamed_logs_packages_b64'] + current_status['stdout_and_stderr_packages_b64']
+        )
 
         compute_state['status']['stdout_and_stderr_packages_b64'] = []
 
     if current_status['status_updates']:
-        compute_state['previous_status_updates'].extend(
-            current_status['status_updates']
-        )
+        compute_state['previous_status_updates'].extend(current_status['status_updates'])
         compute_state['status']['status_updates'] = []
 
     if return_full_logs:
@@ -132,21 +136,37 @@ def send_package_to_compute_process(job_id, package_bytes):
 
 
 def start_webserver(
-        host: str,
-        port: int,
-        tls_pem_certificate_path: Optional[str] = None,
-        tls_pem_key_path: Optional[str] = None,
+    host: str,
+    port: int,
+    tls_pem_certificate_path: Optional[str] = None,
+    tls_pem_key_path: Optional[str] = None,
 ) -> None:
     def worker_exit(server, worker):  # pylint: disable=unused-argument
-        active_compute_states = list(
-            webserver_utils.JOB_ID_TO_COMPUTE_STATE_DICT.values()) + webserver_utils.UNASSIGNED_COMPUTE_PROCESSES
+        active_compute_states = (
+            list(webserver_utils.JOB_ID_TO_COMPUTE_STATE_DICT.values()) + webserver_utils.UNASSIGNED_COMPUTE_PROCESSES
+        )
         logger.debug(f'Sending terminate signal to {len(active_compute_states)} compute processes')
         if active_compute_states:
             for compute_state in active_compute_states:
                 if compute_state['worker_thread']:
                     compute_state['worker_thread'].terminate()
             time.sleep(2)
-        return
+
+        if utils.IS_RUNNING_IN_CLOUD:
+            try:
+                logger_no_user_data.debug('Stopping ComputeNodeResultsProxy...')
+                ComputeNodeResultsProxy.stop_proxy()
+            except BaseException:
+                logger_no_user_data.exception('Failed to stop ComputeNodeResultsProxy')
+
+            try:
+                logger_no_user_data.debug(f'Removing Docker network {BIOLIB_PROXY_NETWORK_NAME}')
+                docker_client = BiolibDockerClient.get_docker_client()
+                biolib_proxy_network: Network = docker_client.networks.get(BIOLIB_PROXY_NETWORK_NAME)
+                biolib_proxy_network.remove()
+                logger_no_user_data.debug(f'Successfully removed Docker network {BIOLIB_PROXY_NETWORK_NAME}')
+            except BaseException:
+                logger_no_user_data.exception(f'Failed to clean up network {BIOLIB_PROXY_NETWORK_NAME}')
 
     def post_fork(server, worker):  # pylint: disable=unused-argument
         logger.info('Started compute node')
@@ -157,8 +177,33 @@ def start_webserver(
             utils.IS_DEV = config['is_dev']
             BiolibApiClient.initialize(config['base_url'])
 
-            CloudUtils.initialize()
-            webserver_utils.start_auto_shutdown_timer()
+            biolib_proxy_network: Optional[Network] = None
+            try:
+                logger_no_user_data.debug(f'Creating Docker network {BIOLIB_PROXY_NETWORK_NAME}')
+                docker_client = BiolibDockerClient.get_docker_client()
+                biolib_proxy_network = docker_client.networks.create(
+                    name=BIOLIB_PROXY_NETWORK_NAME,
+                    internal=False,
+                    driver='bridge',
+                )
+                logger_no_user_data.debug(f'Successfully created Docker network {BIOLIB_PROXY_NETWORK_NAME}')
+            except BaseException:
+                logger_no_user_data.exception(f'Failed to create Docker network {BIOLIB_PROXY_NETWORK_NAME}')
+
+            if biolib_proxy_network:
+                try:
+                    logger_no_user_data.debug('Starting ComputeNodeResultsProxy...')
+                    ComputeNodeResultsProxy.start_proxy(tls_pem_certificate_path, tls_pem_key_path)
+                except BaseException:
+                    logger_no_user_data.exception('Failed to start ComputeNodeResultsProxy')
+
+                CloudUtils.initialize()
+                webserver_utils.start_auto_shutdown_timer()
+            else:
+                logger_no_user_data.error(
+                    f'Docker network {BIOLIB_PROXY_NETWORK_NAME} was not created, shutting down...'
+                )
+                CloudUtils.deregister_and_shutdown()
 
     if logger.level == TRACE:
         gunicorn_log_level_name = 'DEBUG'

biolib/experiments/experiment.py

@@ -1,13 +1,15 @@
 import time
 from collections import OrderedDict
+from pathlib import Path
 
 from biolib import api
-from biolib._internal.types.experiment import DeprecatedExperimentDict, ExperimentDict
-from biolib._internal.types.resource import ResourceDetailedDict
 from biolib._internal.utils import open_browser_window_from_notebook
+from biolib._shared.types import DeprecatedExperimentDict, ExperimentDict, ResourceDetailedDict
+from biolib.api.client import ApiClient
 from biolib.biolib_api_client import BiolibApiClient
 from biolib.biolib_errors import BioLibError
 from biolib.jobs.job import Job
+from biolib.jobs.job_result import PathFilter
 from biolib.jobs.types import JobsPaginatedResponse
 from biolib.tables import BioLibTable
 from biolib.typing_utils import Dict, List, Optional, Union
@@ -26,7 +28,13 @@ class Experiment:
         }
     )
 
-    def __init__(self, uri: str, _resource_dict: Optional[ResourceDetailedDict] = None):
+    def __init__(
+        self,
+        uri: str,
+        _resource_dict: Optional[ResourceDetailedDict] = None,
+        _api_client: Optional[ApiClient] = None,
+    ):
+        self._api_client = _api_client or api.client
         self._resource_dict: ResourceDetailedDict = _resource_dict or self._get_or_create_resource_dict(uri)
 
     def __enter__(self):
@@ -111,7 +119,7 @@ class Experiment:
             job_id = job
         elif job is None and job_id is None:
             raise BioLibError('A job ID or job object must be provided to add job')
-        api.client.post(
+        self._api_client.post(
             path=f'/experiments/{self.uuid}/jobs/',
             data={'job_uuid': job_id},
         )
@@ -124,7 +132,7 @@ class Experiment:
         else:
             raise BioLibError('A job ID or job object must be provided to remove job')
 
-        api.client.delete(path=f'/experiments/{self.uuid}/jobs/{job_id}/')
+        self._api_client.delete(path=f'/experiments/{self.uuid}/jobs/{job_id}/')
 
     def mount_files(self, mount_path: str) -> None:
         try:
@@ -173,7 +181,7 @@ class Experiment:
 
     # Prints a table listing info about the jobs in this experiment
     def show_jobs(self) -> None:
-        response: JobsPaginatedResponse = api.client.get(
+        response: JobsPaginatedResponse = self._api_client.get(
             path=f'/experiments/{self.uuid}/jobs/',
             params=dict(page_size=10),
         ).json()
@@ -195,11 +203,13 @@ class Experiment:
         if status:
             params['status'] = status
 
-        response: JobsPaginatedResponse = api.client.get(url, params=params).json()
+        response: JobsPaginatedResponse = self._api_client.get(url, params=params).json()
         jobs: List[Job] = [Job(job_dict) for job_dict in response['results']]
 
         for page_number in range(2, response['page_count'] + 1):
-            page_response: JobsPaginatedResponse = api.client.get(url, params=dict(**params, page=page_number)).json()
+            page_response: JobsPaginatedResponse = self._api_client.get(
+                url, params=dict(**params, page=page_number)
+            ).json()
             jobs.extend([Job(job_dict) for job_dict in page_response['results']])
 
         return jobs
@@ -223,22 +233,94 @@ class Experiment:
         """
         return self.get_jobs(status=status)
 
+    def save_completed_results(
+        self,
+        output_dir: Optional[str] = None,
+        path_filter: Optional[PathFilter] = None,
+        skip_file_if_exists: bool = False,
+        overwrite: bool = False,
+    ) -> None:
+        r"""Save all completed results in this experiment to local folders.
+
+        Creates a folder structure with the experiment name as the root directory,
+        containing a subfolder for each completed result. Only results with
+        'completed' status will be saved.
+
+        Args:
+            output_dir (str, optional): Base directory where the experiment folder
+                will be created. If None, uses the current working directory.
+            path_filter (PathFilter, optional): Filter to select which files in the results to save.
+                Can be a glob pattern string or a callable function.
+            skip_file_if_exists (bool, optional): Whether to skip files that already exist
+                locally instead of raising an error. Defaults to False.
+            overwrite (bool, optional): Whether to overwrite existing files.
+                Defaults to False.
+
+        Example::
+
+            >>> # Save all completed results to current directory
+            >>> experiment.save_completed_results()
+            >>> # This creates: ./experiment_name/result_1/, ./experiment_name/result_2/, etc.
+
+            >>> # Save to specific directory
+            >>> experiment.save_completed_results(output_dir="/path/to/save")
+            >>> # This creates: /path/to/save/experiment_name/result_1/, etc.
+        """
+        base_dir = Path(output_dir) if output_dir else Path.cwd()
+
+        if base_dir == Path('/'):
+            raise BioLibError("Cannot save experiment results to root directory '/'")
+
+        experiment_folder = base_dir / self.name
+        experiment_folder.mkdir(parents=True, exist_ok=True)
+
+        completed_results: List[Job] = []
+        failed_results = False
+        print('Getting experiment status...')
+        for result in self.get_results():
+            if result.get_status() == 'completed':
+                completed_results.append(result)
+            elif result.get_status() != 'in_progress':
+                failed_results = True
+
+        if failed_results:
+            print(
+                'WARNING: Found failed or cancelled results in the experiment. '
+                'Please verify you have all your results, and consider removing the failed ones.'
+            )
+        if not completed_results:
+            print(f"No completed results found in experiment '{self.name}'")
+            return
+
+        print(f"Saving {len(completed_results)} completed results from experiment '{self.name}' to {experiment_folder}")
+
+        for result in completed_results:
+            result_name = result.get_name()
+            result_folder = experiment_folder / result_name
+
+            result_folder.mkdir(parents=True, exist_ok=True)
+
+            result.save_files(
+                output_dir=str(result_folder),
+                path_filter=path_filter,
+                skip_file_if_exists=skip_file_if_exists,
+                overwrite=overwrite,
+            )
+
     def rename(self, destination: str) -> None:
-        api.client.patch(f'/resources/{self.uuid}/', data={'uri': destination})
+        self._api_client.patch(f'/resources/{self.uuid}/', data={'uri': destination})
         self._refetch()
 
-    @staticmethod
-    def _get_resource_dict_by_uuid(uuid: str) -> ResourceDetailedDict:
-        resource_dict: ResourceDetailedDict = api.client.get(f'/resources/{uuid}/').json()
+    def _get_resource_dict_by_uuid(self, uuid: str) -> ResourceDetailedDict:
+        resource_dict: ResourceDetailedDict = self._api_client.get(f'/resources/{uuid}/').json()
         if not resource_dict['experiment']:
             raise ValueError('Resource from URI is not an experiment')
 
         return resource_dict
 
-    @staticmethod
-    def _get_or_create_resource_dict(uri: str) -> ResourceDetailedDict:
-        response_dict = api.client.post(path='/experiments/', data={'uri' if '/' in uri else 'name': uri}).json()
-        return Experiment._get_resource_dict_by_uuid(uuid=response_dict['uuid'])
+    def _get_or_create_resource_dict(self, uri: str) -> ResourceDetailedDict:
+        response_dict = self._api_client.post(path='/experiments/', data={'uri' if '/' in uri else 'name': uri}).json()
+        return self._get_resource_dict_by_uuid(uuid=response_dict['uuid'])
 
     def _refetch(self) -> None:
         self._resource_dict = self._get_resource_dict_by_uuid(uuid=self._resource_dict['uuid'])