skypilot-nightly 1.0.0.dev20250509__py3-none-any.whl → 1.0.0.dev20251107__py3-none-any.whl

sky/skylet/skylet.py

@@ -1,11 +1,20 @@
 """skylet: a daemon running on the head node of a cluster."""
 
+import concurrent.futures
+import os
 import time
 
+import grpc
+
 import sky
 from sky import sky_logging
+from sky.schemas.generated import autostopv1_pb2_grpc
+from sky.schemas.generated import jobsv1_pb2_grpc
+from sky.schemas.generated import managed_jobsv1_pb2_grpc
+from sky.schemas.generated import servev1_pb2_grpc
 from sky.skylet import constants
 from sky.skylet import events
+from sky.skylet import services
 
 # Use the explicit logger name so that the logger is under the
 # `sky.skylet.skylet` namespace when executed directly, so as
@@ -24,12 +33,62 @@ EVENTS = [
     # This is for monitoring controller job status. If it becomes
     # unhealthy, this event will correctly update the controller
     # status to CONTROLLER_FAILED.
-    events.ServiceUpdateEvent(),
+    events.ServiceUpdateEvent(pool=False),
+    # Status refresh for pool.
+    events.ServiceUpdateEvent(pool=True),
     # Report usage heartbeat every 10 minutes.
     events.UsageHeartbeatReportEvent(),
 ]
 
-while True:
-    time.sleep(events.EVENT_CHECKING_INTERVAL_SECONDS)
+
+def start_grpc_server(port: int = constants.SKYLET_GRPC_PORT) -> grpc.Server:
+    """Start the gRPC server."""
+    # This is the default value in Python 3.8 - 3.12,
+    # putting it here for visibility.
+    # TODO(kevin): Determine the optimal max number of threads.
+    max_workers = min(32, (os.cpu_count() or 1) + 4)
+    server = grpc.server(
+        concurrent.futures.ThreadPoolExecutor(max_workers=max_workers))
+
+    autostopv1_pb2_grpc.add_AutostopServiceServicer_to_server(
+        services.AutostopServiceImpl(), server)
+    jobsv1_pb2_grpc.add_JobsServiceServicer_to_server(
+        services.JobsServiceImpl(), server)
+    servev1_pb2_grpc.add_ServeServiceServicer_to_server(
+        services.ServeServiceImpl(), server)
+    managed_jobsv1_pb2_grpc.add_ManagedJobsServiceServicer_to_server(
+        services.ManagedJobsServiceImpl(), server)
+
+    listen_addr = f'127.0.0.1:{port}'
+    server.add_insecure_port(listen_addr)
+
+    server.start()
+    logger.info(f'gRPC server started on {listen_addr}')
+
+    return server
+
+
+def run_event_loop():
+    """Run the existing event loop."""
+
     for event in EVENTS:
-        event.run()
+        event.start()
+
+    while True:
+        time.sleep(events.EVENT_CHECKING_INTERVAL_SECONDS)
+        for event in EVENTS:
+            event.run()
+
+
+def main():
+    grpc_server = start_grpc_server()
+    try:
+        run_event_loop()
+    except KeyboardInterrupt:
+        logger.info('Shutting down skylet...')
+    finally:
+        grpc_server.stop(grace=5)
+
+
+if __name__ == '__main__':
+    main()

sky/skylet/subprocess_daemon.py

@@ -4,11 +4,16 @@ processes of proc_pid.
 """
 import argparse
 import os
+import signal
 import sys
 import time
+from typing import List, Optional
 
 import psutil
 
+# Environment variable to enable kill_pg in subprocess daemon.
+USE_KILL_PG_ENV_VAR = 'SKYPILOT_SUBPROCESS_DAEMON_KILL_PG'
+
 
 def daemonize():
     """Detaches the process from its parent process with double-forking.
@@ -38,8 +43,74 @@ def daemonize():
     # This process is now fully detached from the original parent and terminal
 
 
-if __name__ == '__main__':
-    daemonize()
+def get_pgid_if_leader(pid) -> Optional[int]:
+    """Get the process group ID of the target process if it is the leader."""
+    try:
+        pgid = os.getpgid(pid)
+        # Only use process group if the target process is the leader. This is
+        # to avoid killing the entire process group while the target process is
+        # just a subprocess in the group.
+        if pgid == pid:
+            print(f'Process group {pgid} is the leader.')
+            return pgid
+        return None
+    except Exception:  # pylint: disable=broad-except
+        # Process group is only available in UNIX.
+        return None
+
+
+def kill_process_group(pgid: int) -> bool:
+    """Kill the target process group."""
+    try:
+        print(f'Terminating process group {pgid}...')
+        os.killpg(pgid, signal.SIGTERM)
+    except Exception:  # pylint: disable=broad-except
+        return False
+
+    # Wait 30s for the process group to exit gracefully.
+    time.sleep(30)
+
+    try:
+        print(f'Force killing process group {pgid}...')
+        os.killpg(pgid, signal.SIGKILL)
+    except Exception:  # pylint: disable=broad-except
+        pass
+
+    return True
+
+
+def kill_process_tree(process: psutil.Process,
+                      children: List[psutil.Process]) -> bool:
+    """Kill the process tree of the target process."""
+    if process is not None:
+        # Kill the target process first to avoid having more children, or fail
+        # the process due to the children being defunct.
+        children = [process] + children
+
+    if not children:
+        sys.exit()
+
+    for child in children:
+        try:
+            child.terminate()
+        except psutil.NoSuchProcess:
+            continue
+
+    # Wait 30s for the processes to exit gracefully.
+    time.sleep(30)
+
+    # SIGKILL if they're still running.
+    for child in children:
+        try:
+            child.kill()
+        except psutil.NoSuchProcess:
+            continue
+
+    return True
+
+
+def main():
+    # daemonize()
     parser = argparse.ArgumentParser()
     parser.add_argument('--parent-pid', type=int, required=True)
     parser.add_argument('--proc-pid', type=int, required=True)
@@ -72,37 +143,40 @@ if __name__ == '__main__':
             except (psutil.NoSuchProcess, ValueError):
                 pass
 
+    pgid: Optional[int] = None
+    if os.environ.get(USE_KILL_PG_ENV_VAR) == '1':
+        # Use kill_pg on UNIX system if allowed to reduce the resource usage.
+        # Note that both implementations might leave subprocessed uncancelled:
+        # - kill_process_tree(default): a subprocess is able to detach itself
+        #   from the process tree use the same technique as daemonize(). Also,
+        #   since we refresh the process tree per second, if the subprocess is
+        #   launched between the [last_poll, parent_die] interval, the
+        #   subprocess will not be captured will not be killed.
+        # - kill_process_group: kill_pg will kill all the processed in the group
+        #   but if a subprocess calls setpgid(0, 0) to detach itself from the
+        #   process group (usually to daemonize itself), the subprocess will
+        #   not be killed.
+        pgid = get_pgid_if_leader(process.pid)
+
     if process is not None and parent_process is not None:
         # Wait for either parent or target process to exit
         while process.is_running() and parent_process.is_running():
-            try:
-                tmp_children = process.children(recursive=True)
-                if tmp_children:
-                    children = tmp_children
-            except psutil.NoSuchProcess:
-                pass
+            if pgid is None:
+                # Refresh process tree for cleanup if process group is not
+                # available.
+                try:
+                    tmp_children = process.children(recursive=True)
+                    if tmp_children:
+                        children = tmp_children
+                except psutil.NoSuchProcess:
+                    pass
             time.sleep(1)
 
-    if process is not None:
-        # Kill the target process first to avoid having more children, or fail
-        # the process due to the children being defunct.
-        children = [process] + children
+    if pgid is not None:
+        kill_process_group(pgid)
+    else:
+        kill_process_tree(process, children)
 
-    if not children:
-        sys.exit()
 
-    for child in children:
-        try:
-            child.terminate()
-        except psutil.NoSuchProcess:
-            continue
-
-    # Wait 30s for the processes to exit gracefully.
-    time.sleep(30)
-
-    # SIGKILL if they're still running.
-    for child in children:
-        try:
-            child.kill()
-        except psutil.NoSuchProcess:
-            continue
+if __name__ == '__main__':
+    main()