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

sky/skylet/events.py

@@ -7,12 +7,12 @@ import time
 import traceback
 
 import psutil
-import yaml
 
 from sky import clouds
 from sky import sky_logging
 from sky.backends import cloud_vm_ray_backend
-from sky.jobs import scheduler as managed_job_scheduler
+from sky.jobs import constants as managed_job_constants
+from sky.jobs import scheduler
 from sky.jobs import state as managed_job_state
 from sky.jobs import utils as managed_job_utils
 from sky.serve import serve_utils
@@ -21,9 +21,10 @@ from sky.skylet import constants
 from sky.skylet import job_lib
 from sky.usage import usage_lib
 from sky.utils import cluster_utils
-from sky.utils import common_utils
 from sky.utils import registry
+from sky.utils import subprocess_utils
 from sky.utils import ux_utils
+from sky.utils import yaml_utils
 
 # Seconds of sleep between the processing of skylet events.
 EVENT_CHECKING_INTERVAL_SECONDS = 20
@@ -46,6 +47,9 @@ class SkyletEvent:
                       EVENT_CHECKING_INTERVAL_SECONDS))
         self._n = 0
 
+    def start(self):
+        pass
+
     def run(self):
         self._n = (self._n + 1) % self._event_interval
         if self._n % self._event_interval == 0:
@@ -74,9 +78,64 @@ class ManagedJobEvent(SkyletEvent):
     """Skylet event for updating and scheduling managed jobs."""
     EVENT_INTERVAL_SECONDS = 300
 
+    def start(self):
+        cpus_env_var = os.environ.get('SKYPILOT_POD_CPU_CORE_LIMIT')
+        if cpus_env_var is not None:
+            with open(os.path.expanduser(constants.CONTROLLER_K8S_CPU_FILE),
+                      'w',
+                      encoding='utf-8') as f:
+                f.write(cpus_env_var)
+        memory_env_var = os.environ.get('SKYPILOT_POD_MEMORY_GB_LIMIT')
+        if memory_env_var is not None:
+            with open(os.path.expanduser(constants.CONTROLLER_K8S_MEMORY_FILE),
+                      'w',
+                      encoding='utf-8') as f:
+                f.write(memory_env_var)
+
     def _run(self):
+        if not os.path.exists(
+                os.path.expanduser(
+                    managed_job_constants.JOB_CONTROLLER_INDICATOR_FILE)):
+            # Note: since the skylet is started before the user setup (in
+            # jobs-controller.yaml.j2) runs, it's possible that we hit this
+            # before the indicator file is written. However, since we will wait
+            # EVENT_INTERVAL_SECONDS before the first run, this should be very
+            # unlikely.
+            logger.info('No jobs controller indicator file found.')
+            all_job_ids = managed_job_state.get_all_job_ids_by_name(None)
+            if not all_job_ids:
+                logger.info('No jobs running. Stopping controllers.')
+                # TODO(cooperc): Move this to a shared function also called by
+                # sdk.api_stop(). (#7229)
+                try:
+                    with open(os.path.expanduser(
+                            scheduler.JOB_CONTROLLER_PID_PATH),
+                              'r',
+                              encoding='utf-8') as f:
+                        pids = f.read().split('\n')[:-1]
+                        for pid in pids:
+                            if subprocess_utils.is_process_alive(
+                                    int(pid.strip())):
+                                subprocess_utils.kill_children_processes(
+                                    parent_pids=[int(pid.strip())], force=True)
+                    os.remove(
+                        os.path.expanduser(scheduler.JOB_CONTROLLER_PID_PATH))
+                except FileNotFoundError:
+                    # its fine we will create it
+                    pass
+                except Exception as e:  # pylint: disable=broad-except
+                    # in case we get perm issues or something is messed up, just
+                    # ignore it and assume the process is dead
+                    logger.error(
+                        f'Error looking at job controller pid file: {e}')
+                    pass
+            logger.info(f'{len(all_job_ids)} jobs running. Assuming the '
+                        'indicator file hasn\'t been written yet.')
+            return
+
+        logger.info('=== Updating managed job status ===')
         managed_job_utils.update_managed_jobs_statuses()
-        managed_job_scheduler.maybe_schedule_next_jobs()
+        scheduler.maybe_start_controllers()
 
 
 class ServiceUpdateEvent(SkyletEvent):
@@ -87,8 +146,12 @@ class ServiceUpdateEvent(SkyletEvent):
     """
     EVENT_INTERVAL_SECONDS = 300
 
+    def __init__(self, pool: bool) -> None:
+        super().__init__()
+        self._pool = pool
+
     def _run(self):
-        serve_utils.update_service_status()
+        serve_utils.update_service_status(self._pool)
 
 
 class UsageHeartbeatReportEvent(SkyletEvent):
@@ -128,23 +191,37 @@ class AutostopEvent(SkyletEvent):
             logger.debug('autostop_config not set. Skipped.')
             return
 
-        if (job_lib.is_cluster_idle() and
-                not managed_job_state.get_num_alive_jobs()):
-            idle_minutes = (time.time() -
-                            autostop_lib.get_last_active_time()) // 60
+        ignore_idle_check = (
+            autostop_config.wait_for == autostop_lib.AutostopWaitFor.NONE)
+        is_idle = True
+        if not ignore_idle_check:
+            if not job_lib.is_cluster_idle(
+            ) or managed_job_state.get_num_alive_jobs() or (
+                    autostop_config.wait_for
+                    == autostop_lib.AutostopWaitFor.JOBS_AND_SSH and
+                    autostop_lib.has_active_ssh_sessions()):
+                is_idle = False
+
+        if ignore_idle_check or is_idle:
+            minutes_since_last_active = (
+                time.time() - autostop_lib.get_last_active_time()) // 60
             logger.debug(
-                f'Idle minutes: {idle_minutes}, '
-                f'AutoStop config: {autostop_config.autostop_idle_minutes}')
+                f'Minutes since last active: {minutes_since_last_active}, '
+                f'AutoStop idle minutes: '
+                f'{autostop_config.autostop_idle_minutes}, '
+                f'Wait for: {autostop_config.wait_for.value}')
         else:
             autostop_lib.set_last_active_time_to_now()
-            idle_minutes = -1
-            logger.debug(
-                'Not idle. Reset idle minutes.'
-                f'AutoStop config: {autostop_config.autostop_idle_minutes}')
-        if idle_minutes >= autostop_config.autostop_idle_minutes:
+            minutes_since_last_active = -1
+            logger.debug('Not idle. Reset idle minutes. '
+                         f'AutoStop idle minutes: '
+                         f'{autostop_config.autostop_idle_minutes}, '
+                         f'Wait for: {autostop_config.wait_for.value}')
+        if minutes_since_last_active >= autostop_config.autostop_idle_minutes:
             logger.info(
-                f'{idle_minutes} idle minutes reached; threshold: '
-                f'{autostop_config.autostop_idle_minutes} minutes. Stopping.')
+                f'{minutes_since_last_active} minute(s) since last active; '
+                f'threshold: {autostop_config.autostop_idle_minutes} minutes. '
+                f'Stopping.')
             self._stop_cluster(autostop_config)
 
     def _stop_cluster(self, autostop_config):
@@ -154,7 +231,7 @@ class AutostopEvent(SkyletEvent):
 
             config_path = os.path.abspath(
                 os.path.expanduser(cluster_utils.SKY_CLUSTER_YAML_REMOTE_PATH))
-            config = common_utils.read_yaml(config_path)
+            config = yaml_utils.read_yaml(config_path)
             provider_name = cluster_utils.get_provider_name(config)
             cloud = registry.CLOUD_REGISTRY.from_str(provider_name)
             assert cloud is not None, f'Unknown cloud: {provider_name}'
@@ -249,8 +326,15 @@ class AutostopEvent(SkyletEvent):
         cluster_name_on_cloud = cluster_config['cluster_name']
         is_cluster_multinode = cluster_config['max_workers'] > 0
 
+        # Clear AWS credentials from environment to force boto3 to use IAM
+        # role attached to the instance (lowest priority in credential chain).
+        # This allows the cluster to stop/terminate itself using its IAM role.
         os.environ.pop('AWS_ACCESS_KEY_ID', None)
         os.environ.pop('AWS_SECRET_ACCESS_KEY', None)
+        os.environ.pop('AWS_SESSION_TOKEN', None)
+        # Point boto3 to /dev/null to skip reading credentials from files.
+        os.environ['AWS_SHARED_CREDENTIALS_FILE'] = '/dev/null'
+        os.environ['AWS_CONFIG_FILE'] = '/dev/null'
 
         # Stop the ray autoscaler to avoid scaling up, during
         # stopping/terminating of the cluster.
@@ -282,7 +366,7 @@ class AutostopEvent(SkyletEvent):
         else:
             yaml_str = self._CATCH_NODES.sub(r'cache_stopped_nodes: true',
                                              yaml_str)
-        config = yaml.safe_load(yaml_str)
+        config = yaml_utils.safe_load(yaml_str)
         # Set the private key with the existed key on the remote instance.
         config['auth']['ssh_private_key'] = '~/ray_bootstrap_key.pem'
         # NOTE: We must do this, otherwise with ssh_proxy_command still under
@@ -299,5 +383,5 @@ class AutostopEvent(SkyletEvent):
         config['auth'].pop('ssh_proxy_command', None)
         # Empty the file_mounts.
         config['file_mounts'] = {}
-        common_utils.dump_yaml(yaml_path, config)
+        yaml_utils.dump_yaml(yaml_path, config)
         logger.debug('Replaced upscaling speed to 0.')