skypilot-nightly 1.0.0.dev20251004__py3-none-any.whl → 1.0.0.dev20251008__py3-none-any.whl

sky/server/config.py

@@ -111,7 +111,9 @@ def compute_server_config(deploy: bool,
     process after API server was introduced.
     """
     cpu_count = common_utils.get_cpu_count()
+    logger.debug(f'CPU count: {cpu_count}')
     mem_size_gb = common_utils.get_mem_size_gb()
+    logger.debug(f'Memory size: {mem_size_gb}GB')
     max_parallel_for_long = _max_long_worker_parallism(cpu_count,
                                                        mem_size_gb,
                                                        local=not deploy)

sky/server/metrics.py

@@ -24,8 +24,10 @@ logger = sky_logging.init_logger(__name__)
 metrics_app = fastapi.FastAPI()
 
 
+# Serve /metrics in dedicated thread to avoid blocking the event loop
+# of metrics server.
 @metrics_app.get('/metrics')
-async def metrics() -> fastapi.Response:
+def metrics() -> fastapi.Response:
     """Expose aggregated Prometheus metrics from all worker processes."""
     if os.environ.get('PROMETHEUS_MULTIPROC_DIR'):
         # In multiprocess mode, we need to collect metrics from all processes.

sky/server/requests/executor.py

@@ -349,27 +349,30 @@ def override_request_env_and_config(
         os.environ.update(original_env)
 
 
-def _redirect_output(file: TextIO) -> Tuple[int, int]:
-    """Redirect stdout and stderr to the log file."""
-    fd = file.fileno()  # Get the file descriptor from the file object
-    # Store copies of the original stdout and stderr file descriptors
-    original_stdout = os.dup(sys.stdout.fileno())
-    original_stderr = os.dup(sys.stderr.fileno())
+def _get_current_output() -> Tuple[int, int]:
+    """Get the current stdout and stderr file descriptors."""
+    return os.dup(sys.stdout.fileno()), os.dup(sys.stderr.fileno())
+
 
+def _redirect_output(file: TextIO) -> None:
+    """Redirect stdout and stderr to the log file."""
+    # Get the file descriptor from the file object
+    fd = file.fileno()
     # Copy this fd to stdout and stderr
     os.dup2(fd, sys.stdout.fileno())
     os.dup2(fd, sys.stderr.fileno())
-    return original_stdout, original_stderr
 
 
-def _restore_output(original_stdout: int, original_stderr: int) -> None:
+def _restore_output(original_stdout: Optional[int],
+                    original_stderr: Optional[int]) -> None:
     """Restore stdout and stderr to their original file descriptors."""
-    os.dup2(original_stdout, sys.stdout.fileno())
-    os.dup2(original_stderr, sys.stderr.fileno())
+    if original_stdout is not None:
+        os.dup2(original_stdout, sys.stdout.fileno())
+        os.close(original_stdout)
 
-    # Close the duplicate file descriptors
-    os.close(original_stdout)
-    os.close(original_stderr)
+    if original_stderr is not None:
+        os.dup2(original_stderr, sys.stderr.fileno())
+        os.close(original_stderr)
 
 
 def _sigterm_handler(signum: int, frame: Optional['types.FrameType']) -> None:
@@ -397,24 +400,38 @@ def _request_execution_wrapper(request_id: str,
     signal.signal(signal.SIGTERM, _sigterm_handler)
 
     logger.info(f'Running request {request_id} with pid {pid}')
-    with api_requests.update_request(request_id) as request_task:
-        assert request_task is not None, request_id
-        log_path = request_task.log_path
-        request_task.pid = pid
-        request_task.status = api_requests.RequestStatus.RUNNING
-        func = request_task.entrypoint
-        request_body = request_task.request_body
-        request_name = request_task.name
 
-    # Append to the log file instead of overwriting it since there might be
-    # logs from previous retries.
-    with log_path.open('a', encoding='utf-8') as f:
+    original_stdout = original_stderr = None
+    try:
+        # As soon as the request is updated with the executor PID, we can
+        # receive SIGTERM from cancellation. So, we update the request inside
+        # the try block to ensure we have the KeyboardInterrupt handling.
+        with api_requests.update_request(request_id) as request_task:
+            assert request_task is not None, request_id
+            if request_task.status != api_requests.RequestStatus.PENDING:
+                logger.debug(f'Request is already {request_task.status.value}, '
+                             f'skipping execution')
+                return
+            log_path = request_task.log_path
+            request_task.pid = pid
+            request_task.status = api_requests.RequestStatus.RUNNING
+            func = request_task.entrypoint
+            request_body = request_task.request_body
+            request_name = request_task.name
+
         # Store copies of the original stdout and stderr file descriptors
-        original_stdout, original_stderr = _redirect_output(f)
-        # Redirect the stdout/stderr before overriding the environment and
-        # config, as there can be some logs during override that needs to be
-        # captured in the log file.
-        try:
+        # We do this in two steps because we should make sure to restore the
+        # original values even if we are cancelled or fail during the redirect.
+        original_stdout, original_stderr = _get_current_output()
+
+        # Append to the log file instead of overwriting it since there might be
+        # logs from previous retries.
+        with log_path.open('a', encoding='utf-8') as f:
+            # Redirect the stdout/stderr before overriding the environment and
+            # config, as there can be some logs during override that needs to be
+            # captured in the log file.
+            _redirect_output(f)
+
             with sky_logging.add_debug_log_handler(request_id), \
                 override_request_env_and_config(
                     request_body, request_id, request_name), \
@@ -429,53 +446,59 @@ def _request_execution_wrapper(request_id: str,
                                            group='request_execution'):
                     return_value = func(**request_body.to_kwargs())
                 f.flush()
-        except KeyboardInterrupt:
-            logger.info(f'Request {request_id} cancelled by user')
-            # Kill all children processes related to this request.
-            # Each executor handles a single request, so we can safely kill all
-            # children processes related to this request.
-            # This is required as python does not pass the KeyboardInterrupt
-            # to the threads that are not main thread.
-            subprocess_utils.kill_children_processes()
-            _restore_output(original_stdout, original_stderr)
-            return
-        except exceptions.ExecutionRetryableError as e:
-            logger.error(e)
-            logger.info(e.hint)
-            with api_requests.update_request(request_id) as request_task:
-                assert request_task is not None, request_id
-                # Retried request will undergo rescheduling and a new execution,
-                # clear the pid of the request.
-                request_task.pid = None
-            # Yield control to the scheduler for uniform handling of retries.
-            _restore_output(original_stdout, original_stderr)
-            raise
-        except (Exception, SystemExit) as e:  # pylint: disable=broad-except
-            api_requests.set_request_failed(request_id, e)
-            _restore_output(original_stdout, original_stderr)
-            logger.info(f'Request {request_id} failed due to '
-                        f'{common_utils.format_exception(e)}')
-            return
-        else:
-            api_requests.set_request_succeeded(
-                request_id, return_value if not ignore_return_value else None)
-            _restore_output(original_stdout, original_stderr)
-            logger.info(f'Request {request_id} finished')
-        finally:
-            try:
-                # Capture the peak RSS before GC.
-                peak_rss = max(proc.memory_info().rss,
-                               metrics_lib.peak_rss_bytes)
-                # Clear request level cache to release all memory used by
-                # the request.
-                annotations.clear_request_level_cache()
-                with metrics_utils.time_it(name='release_memory',
-                                           group='internal'):
-                    common_utils.release_memory()
-                _record_memory_metrics(request_name, proc, rss_begin, peak_rss)
-            except Exception as e:  # pylint: disable=broad-except
-                logger.error(f'Failed to record memory metrics: '
-                             f'{common_utils.format_exception(e)}')
+    except KeyboardInterrupt:
+        logger.info(f'Request {request_id} cancelled by user')
+        # Kill all children processes related to this request.
+        # Each executor handles a single request, so we can safely kill all
+        # children processes related to this request.
+        # This is required as python does not pass the KeyboardInterrupt to the
+        # threads that are not main thread.
+        subprocess_utils.kill_children_processes()
+        return
+    except exceptions.ExecutionRetryableError as e:
+        logger.error(e)
+        logger.info(e.hint)
+        with api_requests.update_request(request_id) as request_task:
+            assert request_task is not None, request_id
+            # Retried request will undergo rescheduling and a new execution,
+            # clear the pid of the request.
+            request_task.pid = None
+        # Yield control to the scheduler for uniform handling of retries.
+        _restore_output(original_stdout, original_stderr)
+        raise
+    except (Exception, SystemExit) as e:  # pylint: disable=broad-except
+        api_requests.set_request_failed(request_id, e)
+        # Manually reset the original stdout and stderr file descriptors early
+        # so that the "Request xxxx failed due to ..." log message will be
+        # written to the original stdout and stderr file descriptors.
+        _restore_output(original_stdout, original_stderr)
+        original_stdout = original_stderr = None
+        logger.info(f'Request {request_id} failed due to '
+                    f'{common_utils.format_exception(e)}')
+        return
+    else:
+        api_requests.set_request_succeeded(
+            request_id, return_value if not ignore_return_value else None)
+        # Manually reset the original stdout and stderr file descriptors early
+        # so that the "Request xxxx failed due to ..." log message will be
+        # written to the original stdout and stderr file descriptors.
+        _restore_output(original_stdout, original_stderr)
+        original_stdout = original_stderr = None
+        logger.info(f'Request {request_id} finished')
+    finally:
+        _restore_output(original_stdout, original_stderr)
+        try:
+            # Capture the peak RSS before GC.
+            peak_rss = max(proc.memory_info().rss, metrics_lib.peak_rss_bytes)
+            # Clear request level cache to release all memory used by the
+            # request.
+            annotations.clear_request_level_cache()
+            with metrics_utils.time_it(name='release_memory', group='internal'):
+                common_utils.release_memory()
+            _record_memory_metrics(request_name, proc, rss_begin, peak_rss)
+        except Exception as e:  # pylint: disable=broad-except
+            logger.error(f'Failed to record memory metrics: '
+                         f'{common_utils.format_exception(e)}')
 
 
 _first_request = True
@@ -596,11 +619,14 @@ async def _execute_request_coroutine(request: api_requests.Request):
     except (Exception, KeyboardInterrupt, SystemExit) as e:
         # Handle any other error
         ctx.redirect_log(original_output)
-        ctx.cancel()
         api_requests.set_request_failed(request.request_id, e)
         logger.error(f'Request {request.request_id} interrupted due to '
                      f'unhandled exception: {common_utils.format_exception(e)}')
         raise
+    finally:
+        # Always cancel the context to kill potentially running background
+        # routine.
+        ctx.cancel()
 
 
 def prepare_request(

sky/server/requests/requests.py

@@ -449,9 +449,15 @@ def init_db_async(func):
 
 def reset_db_and_logs():
     """Create the database."""
+    logger.debug('clearing local API server database')
     server_common.clear_local_api_server_database()
+    logger.debug(
+        f'clearing local API server logs directory at {REQUEST_LOG_PATH_PREFIX}'
+    )
     shutil.rmtree(pathlib.Path(REQUEST_LOG_PATH_PREFIX).expanduser(),
                   ignore_errors=True)
+    logger.debug('clearing local API server client directory at '
+                 f'{server_common.API_SERVER_CLIENT_DIR.expanduser()}')
     shutil.rmtree(server_common.API_SERVER_CLIENT_DIR.expanduser(),
                   ignore_errors=True)
 
@@ -467,10 +473,13 @@ def request_lock_path(request_id: str) -> str:
 @metrics_lib.time_me
 def update_request(request_id: str) -> Generator[Optional[Request], None, None]:
     """Get and update a SkyPilot API request."""
-    request = _get_request_no_lock(request_id)
-    yield request
-    if request is not None:
-        _add_or_update_request_no_lock(request)
+    # Acquire the lock to avoid race conditions between multiple request
+    # operations, e.g. execute and cancel.
+    with filelock.FileLock(request_lock_path(request_id)):
+        request = _get_request_no_lock(request_id)
+        yield request
+        if request is not None:
+            _add_or_update_request_no_lock(request)
 
 
 @init_db
@@ -485,12 +494,15 @@ def update_request_async(
 
     @contextlib.asynccontextmanager
     async def _cm():
-        request = await _get_request_no_lock_async(request_id)
-        try:
-            yield request
-        finally:
-            if request is not None:
-                await _add_or_update_request_no_lock_async(request)
+        # Acquire the lock to avoid race conditions between multiple request
+        # operations, e.g. execute and cancel.
+        async with filelock.AsyncFileLock(request_lock_path(request_id)):
+            request = await _get_request_no_lock_async(request_id)
+            try:
+                yield request
+            finally:
+                if request is not None:
+                    await _add_or_update_request_no_lock_async(request)
 
     return _cm()
 
@@ -775,9 +787,12 @@ def set_request_succeeded(request_id: str, result: Optional[Any]) -> None:
 
 
 def set_request_cancelled(request_id: str) -> None:
-    """Set a request to cancelled."""
+    """Set a pending or running request to cancelled."""
     with update_request(request_id) as request_task:
         assert request_task is not None, request_id
+        # Already finished or cancelled.
+        if request_task.status > RequestStatus.RUNNING:
+            return
         request_task.finished_at = time.time()
         request_task.status = RequestStatus.CANCELLED
 

sky/server/server.py

@@ -1943,6 +1943,7 @@ if __name__ == '__main__':
 
     from sky.server import uvicorn as skyuvicorn
 
+    logger.info('Initializing SkyPilot API server')
     skyuvicorn.add_timestamp_prefix_for_server_logs()
 
     parser = argparse.ArgumentParser()
@@ -1954,20 +1955,35 @@ if __name__ == '__main__':
     parser.add_argument('--metrics-port', default=9090, type=int)
     cmd_args = parser.parse_args()
     if cmd_args.port == cmd_args.metrics_port:
+        logger.error('port and metrics-port cannot be the same, exiting.')
         raise ValueError('port and metrics-port cannot be the same')
 
+    # Fail fast if the port is not available to avoid corrupt the state
+    # of potential running server instance.
+    # We might reach here because the running server is currently not
+    # responding, thus the healthz check fails and `sky api start` think
+    # we should start a new server instance.
+    if not common_utils.is_port_available(cmd_args.port):
+        logger.error(f'Port {cmd_args.port} is not available, exiting.')
+        raise RuntimeError(f'Port {cmd_args.port} is not available')
+
     # Show the privacy policy if it is not already shown. We place it here so
     # that it is shown only when the API server is started.
     usage_lib.maybe_show_privacy_policy()
 
     # Initialize global user state db
     db_utils.set_max_connections(1)
+    logger.info('Initializing database engine')
     global_user_state.initialize_and_get_db()
+    logger.info('Database engine initialized')
     # Initialize request db
     requests_lib.reset_db_and_logs()
     # Restore the server user hash
+    logger.info('Initializing server user hash')
     _init_or_restore_server_user_hash()
+
     max_db_connections = global_user_state.get_max_db_connections()
+    logger.info(f'Max db connections: {max_db_connections}')
     config = server_config.compute_server_config(cmd_args.deploy,
                                                  max_db_connections)
 

sky/skylet/constants.py

@@ -100,7 +100,7 @@ TASK_ID_LIST_ENV_VAR = f'{SKYPILOT_ENV_VAR_PREFIX}TASK_IDS'
 # cluster yaml is updated.
 #
 # TODO(zongheng,zhanghao): make the upgrading of skylet automatic?
-SKYLET_VERSION = '21'
+SKYLET_VERSION = '22'
 # The version of the lib files that skylet/jobs use. Whenever there is an API
 # change for the job_lib or log_lib, we need to bump this version, so that the
 # user can be notified to update their SkyPilot version on the remote cluster.
@@ -331,6 +331,14 @@ FILE_MOUNTS_LOCAL_TMP_BASE_PATH = '~/.sky/tmp/'
 # controller_utils.translate_local_file_mounts_to_two_hop().
 FILE_MOUNTS_CONTROLLER_TMP_BASE_PATH = '~/.sky/tmp/controller'
 
+# For passing in CPU and memory limits to the controller pod when running
+# in k8s. Right now, we only use this for the jobs controller, but we may
+# use this for the serve controller as well in the future.
+# These files are written to disk by the skylet, who reads it from env vars
+# passed by the backend when starting the skylet (start_skylet_on_head_node).
+CONTROLLER_K8S_CPU_FILE = '~/.sky/_internal_k8s_pod_cpu'
+CONTROLLER_K8S_MEMORY_FILE = '~/.sky/_internal_k8s_pod_memory'
+
 # Used when an managed jobs are created and
 # files are synced up to the cloud.
 FILE_MOUNTS_WORKDIR_SUBPATH = 'job-{run_id}/workdir'

sky/skylet/events.py

@@ -47,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:
@@ -75,6 +78,20 @@ 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(

sky/skylet/skylet.py

@@ -71,6 +71,9 @@ def start_grpc_server(port: int = constants.SKYLET_GRPC_PORT) -> grpc.Server:
 def run_event_loop():
     """Run the existing event loop."""
 
+    for event in EVENTS:
+        event.start()
+
     while True:
         time.sleep(events.EVENT_CHECKING_INTERVAL_SECONDS)
         for event in EVENTS:

sky/skypilot_config.py

@@ -825,7 +825,8 @@ def _compose_cli_config(cli_config: Optional[List[str]]) -> config_utils.Config:
     except ValueError as e:
         raise ValueError(f'Invalid config override: {cli_config}. '
                          f'Check if config file exists or if the dotlist '
-                         f'is formatted as: key1=value1,key2=value2') from e
+                         f'is formatted as: key1=value1,key2=value2.\n'
+                         f'Details: {e}') from e
     logger.debug('CLI overrides config syntax check passed.')
 
     return parsed_config

sky/templates/kubernetes-ray.yml.j2

@@ -632,6 +632,9 @@ available_node_types:
           command: ["/bin/bash", "-c", "--"]
           args:
             - |
+              # Set -x to print the commands and their arguments as they are executed.
+              # Useful for debugging.
+              set -x
               # Helper function to conditionally use sudo
               # TODO(zhwu): consolidate the two prefix_cmd and sudo replacements
               prefix_cmd() { if [ $(id -u) -ne 0 ]; then echo "sudo"; else echo ""; fi; }
@@ -1086,6 +1089,8 @@ available_node_types:
 
               touch {{k8s_high_availability_deployment_volume_mount_path}}/k8s_container_ready
               {% endif %}
+              # Set +x to stop printing the commands and their arguments as they are executed.
+              set +x
 
               trap : TERM INT; log_tail || sleep infinity & wait
 

sky/utils/context_utils.py

@@ -130,7 +130,11 @@ def wait_process(ctx: context.Context,
             # Kill the process despite the caller's callback, the utility
             # function gracefully handles the case where the process is
             # already terminated.
-            subprocess_utils.kill_process_with_grace_period(proc)
+            # Bash script typically does not forward SIGTERM to childs, thus
+            # cannot be killed gracefully, shorten the grace period for faster
+            # termination.
+            subprocess_utils.kill_process_with_grace_period(proc,
+                                                            grace_period=1)
             raise asyncio.CancelledError()
         try:
             proc.wait(poll_interval)

sky/utils/controller_utils.py

@@ -506,6 +506,9 @@ def shared_controller_vars_to_fill(
         # before popping allowed_contexts. If it is not on Kubernetes,
         # we may be able to use allowed_contexts.
         local_user_config.pop('allowed_contexts', None)
+        # Remove api_server config so that the controller does not try to use
+        # a remote API server.
+        local_user_config.pop('api_server', None)
         with tempfile.NamedTemporaryFile(
                 delete=False,
                 suffix=_LOCAL_SKYPILOT_CONFIG_PATH_SUFFIX) as temp_file:
@@ -726,6 +729,17 @@ def get_controller_resources(
     return result
 
 
+def get_controller_mem_size_gb() -> float:
+    try:
+        with open(os.path.expanduser(constants.CONTROLLER_K8S_MEMORY_FILE),
+                  'r',
+                  encoding='utf-8') as f:
+            return float(f.read())
+    except FileNotFoundError:
+        pass
+    return common_utils.get_mem_size_gb()
+
+
 def _setup_proxy_command_on_controller(
         controller_launched_cloud: 'clouds.Cloud',
         user_config: Dict[str, Any]) -> config_utils.Config:

sky/utils/db/db_utils.py

@@ -410,6 +410,8 @@ def get_engine(
                 conn_string, poolclass=sqlalchemy.NullPool)
         with _db_creation_lock:
             if conn_string not in _postgres_engine_cache:
+                logger.debug('Creating a new postgres engine with '
+                             f'maximum {_max_connections} connections')
                 if _max_connections == 0:
                     _postgres_engine_cache[conn_string] = (
                         sqlalchemy.create_engine(

sky/utils/db/migration_utils.py

@@ -11,13 +11,14 @@ import filelock
 import sqlalchemy
 
 from sky import sky_logging
+from sky.skylet import constants
 
 logger = sky_logging.init_logger(__name__)
 
 DB_INIT_LOCK_TIMEOUT_SECONDS = 10
 
 GLOBAL_USER_STATE_DB_NAME = 'state_db'
-GLOBAL_USER_STATE_VERSION = '009'
+GLOBAL_USER_STATE_VERSION = '010'
 GLOBAL_USER_STATE_LOCK_PATH = '~/.sky/locks/.state_db.lock'
 
 SPOT_JOBS_DB_NAME = 'spot_jobs_db'
@@ -85,12 +86,20 @@ def needs_upgrade(engine: sqlalchemy.engine.Engine, section: str,
             connection, opts={'version_table': version_table})
         current_rev = context.get_current_revision()
 
+    target_rev_num = int(target_revision)
     if current_rev is None:
+        if os.environ.get(constants.ENV_VAR_IS_SKYPILOT_SERVER) is not None:
+            logger.debug(f'{section} database currently uninitialized, '
+                         f'targeting revision {target_rev_num}')
         return True
 
     # Compare revisions - assuming they are numeric strings like '001', '002'
     current_rev_num = int(current_rev)
-    target_rev_num = int(target_revision)
+    if (current_rev_num < target_rev_num and
+            os.environ.get(constants.ENV_VAR_IS_SKYPILOT_SERVER) is not None):
+        logger.debug(
+            f'{section} database currently at revision {current_rev_num}, '
+            f'targeting revision {target_rev_num}')
 
     return current_rev_num < target_rev_num
 

sky/volumes/server/server.py

@@ -24,11 +24,11 @@ async def volume_list(request: fastapi.Request) -> None:
     auth_user_env_vars_kwargs = {
         'env_vars': auth_user.to_env_vars()
     } if auth_user else {}
-    volume_list_body = payloads.VolumeListBody(**auth_user_env_vars_kwargs)
+    request_body = payloads.RequestBody(**auth_user_env_vars_kwargs)
     executor.schedule_request(
         request_id=request.state.request_id,
         request_name='volume_list',
-        request_body=volume_list_body,
+        request_body=request_body,
         func=core.volume_list,
         schedule_type=requests_lib.ScheduleType.SHORT,
     )