skypilot-nightly 1.0.0.dev20250905__py3-none-any.whl → 1.0.0.dev20251210__py3-none-any.whl

sky/utils/kubernetes/ssh_jump_lifecycle_manager.py

@@ -1,191 +0,0 @@
-"""Manages lifecycle of ssh jump pod.
-
-This script runs inside ssh jump pod as the main process (PID 1).
-
-It terminates itself (by removing ssh jump service and pod via a call to
-kubeapi) if it does not see ray pods in the duration of 10 minutes. If the
-user re-launches a task before the duration is over, then ssh jump pod is being
-reused and will terminate itself when it sees that no ray clusters exist in
-that duration.
-
-To allow multiple users to the share the same SSH jump pod,
-this script also reloads SSH keys from the mounted secret volume on an
-interval and updates `~/.ssh/authorized_keys`.
-"""
-import datetime
-import os
-import subprocess
-import sys
-import threading
-import time
-
-from kubernetes import client
-from kubernetes import config
-
-# Load kube config
-config.load_incluster_config()
-
-v1 = client.CoreV1Api()
-
-current_name = os.getenv('MY_POD_NAME')
-current_namespace = os.getenv('MY_POD_NAMESPACE')
-
-# The amount of time in seconds where no Ray pods exist in which after that time
-# ssh jump pod terminates itself
-alert_threshold = int(os.getenv('ALERT_THRESHOLD', '600'))
-# The amount of time in seconds to wait between Ray pods existence checks
-retry_interval = int(os.getenv('RETRY_INTERVAL', '60'))
-# The amount of time in seconds to wait between SSH key reloads
-reload_interval = int(os.getenv('RELOAD_INTERVAL', '5'))
-
-# Ray pods are labeled with this value i.e., ssh jump name which is unique per
-# user (based on user hash)
-label_selector = f'skypilot-ssh-jump={current_name}'
-
-
-def poll(interval, leading=True):
-    """Decorator factory for polling function. To stop polling, return True.
-
-    Args:
-        interval (int): The amount of time to wait between function calls.
-        leading (bool): Whether to wait before (rather than after) calls.
-    """
-
-    def decorator(func):
-
-        def wrapper(*args, **kwargs):
-            while True:
-                if leading:
-                    time.sleep(interval)
-                done = func(*args, **kwargs)
-                if done:
-                    return
-                if not leading:
-                    time.sleep(interval)
-
-        return wrapper
-
-    return decorator
-
-
-# Flag to terminate the reload keys thread when the lifecycle thread
-# terminates.
-terminated = False
-
-
-@poll(interval=reload_interval, leading=False)
-def reload_keys():
-    """Reloads SSH keys from mounted secret volume."""
-
-    if terminated:
-        sys.stdout.write('[SSH Key Reloader] Terminated.\n')
-        return True
-
-    # Reload SSH keys from mounted secret volume if changed.
-    tmpfile = '/tmp/sky-ssh-keys'
-    try:
-        subprocess.check_output(
-            f'cat /etc/secret-volume/ssh-publickey* > {tmpfile}', shell=True)
-        try:
-            subprocess.check_output(f'diff {tmpfile} ~/.ssh/authorized_keys',
-                                    shell=True)
-            sys.stdout.write(
-                '[SSH Key Reloader] No keys changed, continuing.\n')
-        except subprocess.CalledProcessError as e:
-            if e.returncode == 1:
-                sys.stdout.write(
-                    '[SSH Key Reloader] Changes detected, reloading.\n')
-                subprocess.check_output(f'mv {tmpfile} ~/.ssh/authorized_keys',
-                                        shell=True)
-            else:
-                raise
-    except Exception as e:
-        sys.stdout.write(
-            f'[SSH Key Reloader][ERROR] Failed to reload SSH keys: {e}\n')
-        raise
-
-
-alert_delta = datetime.timedelta(seconds=alert_threshold)
-retry_interval_delta = datetime.timedelta(seconds=retry_interval)
-# Accumulated time of where no SkyPilot cluster exists. Compared
-# against alert_threshold.
-nocluster_delta = datetime.timedelta()
-
-
-@poll(interval=retry_interval)
-def manage_lifecycle():
-    """Manages lifecycle of ssh jump pod."""
-
-    global terminated, nocluster_delta
-
-    try:
-        ret = v1.list_namespaced_pod(current_namespace,
-                                     label_selector=label_selector)
-    except Exception as e:
-        sys.stdout.write('[Lifecycle] [ERROR] listing pods failed with '
-                         f'error: {e}\n')
-        raise
-
-    if not ret.items:
-        sys.stdout.write(
-            f'[Lifecycle] Did not find pods with label '
-            f'"{label_selector}" in namespace {current_namespace}\n')
-        nocluster_delta = nocluster_delta + retry_interval_delta
-        sys.stdout.write(
-            f'[Lifecycle] Time since no pods found: {nocluster_delta}, alert '
-            f'threshold: {alert_delta}\n')
-    else:
-        sys.stdout.write(
-            f'[Lifecycle] Found pods with label "{label_selector}" in '
-            f'namespace {current_namespace}\n')
-        # reset ..
-        nocluster_delta = datetime.timedelta()
-        sys.stdout.write(
-            f'[Lifecycle] nocluster_delta is reset: {nocluster_delta}\n')
-
-    if nocluster_delta >= alert_delta:
-        sys.stdout.write(
-            f'[Lifecycle] nocluster_delta: {nocluster_delta} crossed alert '
-            f'threshold: {alert_delta}. Time to terminate myself and my '
-            'service.\n')
-        try:
-            # ssh jump resources created under same name
-            v1.delete_namespaced_service(current_name, current_namespace)
-            v1.delete_namespaced_pod(current_name, current_namespace)
-        except Exception as e:
-            sys.stdout.write('[Lifecycle][ERROR] Deletion failed. Exiting '
-                             f'poll() with error: {e}\n')
-            raise
-
-        terminated = True
-        return True
-
-
-def main():
-    sys.stdout.write('SkyPilot SSH Jump Pod Lifecycle Manager\n')
-    sys.stdout.write(f'current_name: {current_name}\n')
-    sys.stdout.write(f'current_namespace: {current_namespace}\n')
-    sys.stdout.write(f'alert_threshold time: {alert_threshold}\n')
-    sys.stdout.write(f'retry_interval time: {retry_interval}\n')
-    sys.stdout.write(f'reload_interval time: {reload_interval}\n')
-    sys.stdout.write(f'label_selector: {label_selector}\n')
-
-    if not current_name or not current_namespace:
-        # Raise Exception with message to terminate pod
-        raise Exception('Missing environment variables MY_POD_NAME or '
-                        'MY_POD_NAMESPACE')
-
-    threads = [
-        threading.Thread(target=manage_lifecycle),
-        threading.Thread(target=reload_keys)
-    ]
-    sys.stdout.write(f'Polling with {len(threads)} threads.\n')
-    for t in threads:
-        t.start()
-    for t in threads:
-        t.join()
-    sys.stdout.write('Done.\n')
-
-
-if __name__ == '__main__':
-    main()