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

biolib/compute_node/webserver/compute_node_results_proxy.py

@@ -0,0 +1,189 @@
+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 Host $http_host;
+            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,23 +1,34 @@
 # pylint: disable=unsubscriptable-object
 
 import json
-import time
-import base64
 import logging
-from flask import Flask, request, Response, jsonify
-from flask_cors import CORS  # type: ignore
+import os
+import tempfile
+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.enclave.enclave_remote_hosts import start_enclave_remote_hosts
-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
 
 app = Flask(__name__)
-CORS(app)
+
+if utils.IS_RUNNING_IN_CLOUD:
+    _BIOLIB_TMP_DIR = '/biolib/tmp'
+    os.makedirs(_BIOLIB_TMP_DIR, exist_ok=True)
+else:
+    _BIOLIB_TMP_DIR = tempfile.TemporaryDirectory().name
 
 
 @app.route('/hello/')
@@ -25,6 +36,11 @@ def hello():
     return 'Hello'
 
 
+@app.route('/health/')
+def health():
+    return 'biolib-compute-node is running', 200
+
+
 @app.route('/v1/job/', methods=['POST'])
 def save_job():
     saved_job = json.loads(request.data.decode())
@@ -34,79 +50,123 @@ def save_job():
         return jsonify({'job': 'Invalid job'}), 400
 
     job_id = saved_job['job']['public_id']
+    job_temporary_dir = os.path.join(_BIOLIB_TMP_DIR, job_id)
+    os.makedirs(job_temporary_dir)
     saved_job['BASE_URL'] = BiolibApiClient.get().base_url
+    saved_job['job_temporary_dir'] = job_temporary_dir
 
     compute_state = webserver_utils.get_compute_state(webserver_utils.UNASSIGNED_COMPUTE_PROCESSES)
     compute_state['job_id'] = job_id
+    compute_state['job'] = saved_job['job']
+    compute_state['job_temporary_dir'] = job_temporary_dir
+
     webserver_utils.JOB_ID_TO_COMPUTE_STATE_DICT[job_id] = compute_state
 
     if utils.IS_RUNNING_IN_CLOUD:
         config = CloudUtils.get_webserver_config()
         saved_job['compute_node_info'] = config['compute_node_info']
+        compute_state['cloud_job_id'] = saved_job['cloud_job']['public_id']
+        compute_state['cloud_job'] = saved_job['cloud_job']
+
+        webserver_utils.update_auto_shutdown_time()
 
     saved_job_bbf_package = SavedJob().serialize(json.dumps(saved_job))
     send_package_to_compute_process(job_id, saved_job_bbf_package)
 
-    if utils.BIOLIB_IS_RUNNING_IN_ENCLAVE:
-        return Response(base64.b64encode(compute_state['attestation_document']), status=201)
-    else:
-        return '', 201
+    return '', 201
 
 
 @app.route('/v1/job/<job_id>/start/', methods=['POST'])
 def start_compute(job_id):
     module_input_package = request.data
+
+    if 'AES-Key-String' in request.headers:
+        compute_state = webserver_utils.JOB_ID_TO_COMPUTE_STATE_DICT[job_id]
+        compute_state['aes_key_string_b64'] = request.headers['AES-Key-String']
+
     send_package_to_compute_process(job_id, module_input_package)
     return '', 201
 
 
+@app.route('/v1/job/<job_id>/', methods=['DELETE'])
+def terminate_job(job_id: str) -> Response:
+    compute_state = get_job_compute_state_or_404(job_id)
+    # TODO: Consider BBF package
+    compute_state['received_messages_queue'].put(b'CANCEL_JOB')
+    return Response()
+
+
 @app.route('/v1/job/<job_id>/status/')
 def status(job_id):
     # TODO Implement auth token
-    current_status = webserver_utils.JOB_ID_TO_COMPUTE_STATE_DICT[job_id]['status'].copy()
-    response = jsonify(current_status)
+    return_full_logs = request.args.get('logs') == 'full'
 
-    if current_status['status_updates']:
-        webserver_utils.JOB_ID_TO_COMPUTE_STATE_DICT[job_id]['status']['status_updates'] = []
+    compute_state = get_job_compute_state_or_404(job_id)
+    current_status = compute_state['status'].copy()
+    response_data = current_status
+    response_data['is_completed'] = compute_state['is_completed']
 
-    # Check if any error occurred
-    if 'error_code' in current_status:
-        response.call_on_close(lambda: webserver_utils.finalize_and_clean_up_compute_job(job_id))
+    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']
+        )
 
-    return response
+        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['status']['status_updates'] = []
 
-@app.route('/v1/job/<job_id>/result/')
-def result(job_id):
-    if webserver_utils.JOB_ID_TO_COMPUTE_STATE_DICT[job_id]['result']:
-        result_data = webserver_utils.JOB_ID_TO_COMPUTE_STATE_DICT[job_id]['result']
-        # remove result from state dict, so we know the user has started the download
-        webserver_utils.JOB_ID_TO_COMPUTE_STATE_DICT[job_id]['result'] = None
-        response = Response(result_data)
-        response.call_on_close(lambda: webserver_utils.finalize_and_clean_up_compute_job(job_id))
+    if return_full_logs:
+        response_data['streamed_logs_packages_b64'] = compute_state['streamed_logs_packages_b64']
+        response_data['previous_status_updates'] = compute_state['previous_status_updates']
 
-        return response
+    return jsonify(response_data)
 
-    else:
-        return '', 404
+
+@app.route('/v1/job/<job_id>/result/')
+def result():
+    return 410, 'This endpoint is no longer available. Please update pybiolib by running `pip3 install -U pybiolib`'
 
 
 def send_package_to_compute_process(job_id, package_bytes):
-    message_queue = webserver_utils.JOB_ID_TO_COMPUTE_STATE_DICT[job_id]['messages_to_send_queue']
+    compute_state = get_job_compute_state_or_404(job_id)
+    message_queue = compute_state['messages_to_send_queue']
     message_queue.put(package_bytes)
 
 
-def start_webserver(port, host):
+def start_webserver(
+    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')
@@ -117,10 +177,33 @@ def start_webserver(port, host):
             utils.IS_DEV = config['is_dev']
             BiolibApiClient.initialize(config['base_url'])
 
-            if utils.BIOLIB_IS_RUNNING_IN_ENCLAVE:
-                start_enclave_remote_hosts(config)
-
-            CloudUtils.initialize()
+            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'
@@ -133,12 +216,15 @@ def start_webserver(port, host):
 
     options = {
         'bind': f'{host}:{port}',
-        'workers': 1,
-        'post_fork': post_fork,
-        'worker_exit': worker_exit,
-        'timeout': '300',
+        'certfile': tls_pem_certificate_path,
         'graceful_timeout': 4,
+        'keyfile': tls_pem_key_path,
         'loglevel': gunicorn_log_level_name,
+        'post_fork': post_fork,
+        'ssl_version': 'TLSv1_2',
+        'timeout': '7200',  # Reduce to 300 when frontend no longer downloads from webserver
+        'worker_exit': worker_exit,
+        'workers': 1,
     }
 
     GunicornFlaskApplication(app, options).run()

biolib/compute_node/webserver/webserver_types.py

@@ -3,21 +3,17 @@ from biolib.typing_utils import TypedDict
 
 class ComputeNodeInfo(TypedDict):
     auth_token: str
-    public_id: str
     ip_address: str
+    public_id: str
+    pybiolib_version: str
 
 
 class ShutdownTimes(TypedDict):
     auto_shutdown_time_in_seconds: int
-    job_max_runtime_shutdown_time_in_seconds: int
 
 
 class WebserverConfig(TypedDict):
     base_url: str
-    ecr_region_name: str
-    max_docker_image_cache_size_bytes: int
-    s3_general_storage_bucket_name: str
-    s3_lfs_bucket_name: str
     compute_node_info: ComputeNodeInfo
     is_dev: bool
     shutdown_times: ShutdownTimes

biolib/compute_node/webserver/webserver_utils.py

@@ -1,6 +1,11 @@
 import time
+from datetime import datetime, timedelta
+from multiprocessing import Process
+
+from werkzeug.exceptions import NotFound
 
 from biolib import utils
+from biolib.biolib_logging import logger_no_user_data
 from biolib.compute_node.cloud_utils import CloudUtils
 from biolib.typing_utils import Dict, List
 from biolib.compute_node.webserver.worker_thread import WorkerThread
@@ -10,10 +15,12 @@ JOB_ID_TO_COMPUTE_STATE_DICT: Dict = {}
 UNASSIGNED_COMPUTE_PROCESSES: List = []
 
 
-def finalize_and_clean_up_compute_job(job_id: str):
-    JOB_ID_TO_COMPUTE_STATE_DICT.pop(job_id)
-    if utils.IS_RUNNING_IN_CLOUD:
-        CloudUtils.finish_job(job_id)
+def get_job_compute_state_or_404(job_id: str):
+    compute_state = JOB_ID_TO_COMPUTE_STATE_DICT.get(job_id)
+    if compute_state:
+        return compute_state
+
+    raise NotFound('Job not found')
 
 
 def get_compute_state(unassigned_compute_processes):
@@ -26,6 +33,12 @@ def get_compute_state(unassigned_compute_processes):
 def start_compute_process(unassigned_compute_processes):
     compute_state = {
         'job_id': None,
+        'job': None,
+        'cloud_job_id': None,
+        'cloud_job': None,
+        'streamed_logs_packages_b64': [],
+        'previous_status_updates': [],
+        'is_completed': False,
         'status': {
             'status_updates': [
                 {
@@ -33,21 +46,16 @@ def start_compute_process(unassigned_compute_processes):
                     'log_message': 'Initializing'
                 }
             ],
+            'stdout_and_stderr_packages_b64': []
         },
-        'result': None,
-        'attestation_document': b'',
+        'progress': 0,
+        'job_temporary_dir': '',
         'received_messages_queue': None,
         'messages_to_send_queue': None,
         'worker_process': None
     }
 
     WorkerThread(compute_state).start()
-
-    while True:
-        if compute_state['attestation_document']:
-            break
-        time.sleep(0.25)
-
     unassigned_compute_processes.append(compute_state)
 
 
@@ -62,3 +70,60 @@ def validate_saved_job(saved_job):
         return False
 
     return True
+
+
+def start_auto_shutdown_timer() -> None:
+    if not utils.IS_RUNNING_IN_CLOUD:
+        logger_no_user_data.error('Not running in cloud so skipping auto shutdown time start')
+        return
+
+    update_auto_shutdown_time()
+
+    timer = CloudAutoShutdownTimer()
+    timer.start()
+    logger_no_user_data.debug(f"Started auto shutdown timer on pid: {timer.pid}")
+
+
+def update_auto_shutdown_time() -> None:
+    if not utils.IS_RUNNING_IN_CLOUD:
+        logger_no_user_data.error('Not running in cloud so skipping auto shutdown time update')
+        return
+
+    webserver_config = CloudUtils.get_webserver_config()
+    auto_shutdown_buffer = webserver_config['shutdown_times']['auto_shutdown_time_in_seconds']  # pylint: disable=unsubscriptable-object
+
+    if JOB_ID_TO_COMPUTE_STATE_DICT:
+        highest_max_job_compute_state = max(
+            JOB_ID_TO_COMPUTE_STATE_DICT.values(),
+            key=lambda compute_state: compute_state['cloud_job']['max_runtime_in_seconds']  # type: ignore
+        )
+        highest_max_job_runtime = highest_max_job_compute_state['cloud_job']['max_runtime_in_seconds']
+    else:
+        highest_max_job_runtime = 0
+
+    auto_shutdown_time = datetime.now() + timedelta(seconds=highest_max_job_runtime) + timedelta(
+        seconds=auto_shutdown_buffer
+    )
+    auto_shutdown_time_isoformat = datetime.isoformat(auto_shutdown_time)
+
+    with open(CloudAutoShutdownTimer.SHUTDOWN_TIMESTAMP_FILE_PATH, 'w') as auto_shutdown_time_file:
+        auto_shutdown_time_file.write(auto_shutdown_time_isoformat)
+
+    logger_no_user_data.debug(f'Extending auto shutdown timer to: {auto_shutdown_time_isoformat}')
+
+
+class CloudAutoShutdownTimer(Process):
+    SHUTDOWN_TIMESTAMP_FILE_PATH = '/tmp/auto_shutdown_time.timestamp'
+
+    def _get_auto_shutdown_time(self) -> datetime:
+        with open(self.SHUTDOWN_TIMESTAMP_FILE_PATH, mode='r') as file:
+            return datetime.fromisoformat(file.read())
+
+    def run(self) -> None:
+        auto_shutdown_time = self._get_auto_shutdown_time()
+        while datetime.now() < auto_shutdown_time:
+            time.sleep(60)
+            auto_shutdown_time = self._get_auto_shutdown_time()
+
+        logger_no_user_data.debug(f'Hit auto shutdown timer since {datetime.now()} > {auto_shutdown_time}')
+        CloudUtils.deregister_and_shutdown()