skypilot-nightly 1.0.0.dev20251027__py3-none-any.whl → 1.0.0.dev20251101__py3-none-any.whl

sky/jobs/controller.py

@@ -1,15 +1,17 @@
 """Controller: handles scheduling and the life cycle of a managed job.
 """
 import asyncio
+import io
 import os
 import pathlib
 import resource
 import shutil
 import sys
+import threading
 import time
 import traceback
 import typing
-from typing import Dict, Optional, Set, Tuple
+from typing import Dict, Optional, Set
 
 import dotenv
 
@@ -22,6 +24,8 @@ from sky.backends import backend_utils
 from sky.backends import cloud_vm_ray_backend
 from sky.data import data_utils
 from sky.jobs import constants as jobs_constants
+from sky.jobs import file_content_utils
+from sky.jobs import log_gc
 from sky.jobs import recovery_strategy
 from sky.jobs import scheduler
 from sky.jobs import state as managed_job_state
@@ -29,6 +33,7 @@ from sky.jobs import utils as managed_job_utils
 from sky.skylet import constants
 from sky.skylet import job_lib
 from sky.usage import usage_lib
+from sky.utils import annotations
 from sky.utils import common
 from sky.utils import common_utils
 from sky.utils import context
@@ -61,17 +66,26 @@ async def create_background_task(coro: typing.Coroutine) -> None:
         task.add_done_callback(_background_tasks.discard)
 
 
-def _get_dag_and_name(dag_yaml: str) -> Tuple['sky.Dag', str]:
-    dag = dag_utils.load_chain_dag_from_yaml(dag_yaml)
-    dag_name = dag.name
-    assert dag_name is not None, dag
-    return dag, dag_name
+# Make sure to limit the size as we don't want to cache too many DAGs in memory.
+@annotations.lru_cache(scope='global', maxsize=50)
+def _get_dag(job_id: int) -> 'sky.Dag':
+    dag_content = file_content_utils.get_job_dag_content(job_id)
+    if dag_content is None:
+        raise RuntimeError('Managed job DAG YAML content is unavailable for '
+                           f'job {job_id}. This can happen if the job was '
+                           'submitted before file migration completed or if '
+                           'the submission failed to persist the DAG. Please '
+                           're-submit the job.')
 
+    dag = dag_utils.load_chain_dag_from_yaml_str(dag_content)
+    assert dag.name is not None, dag
+    return dag
 
-class JobsController:
+
+class JobController:
     """Controls the lifecycle of a single managed job.
 
-    This controller executes a chain DAG defined in ``dag_yaml`` by:
+    This controller executes the chain DAG recorded for the job by:
     - Loading the DAG and preparing per-task environment variables so each task
       has a stable global job identifier across recoveries.
     - Launching the task on the configured backend (``CloudVmRayBackend``),
@@ -91,7 +105,8 @@ class JobsController:
 
     Key attributes:
     - ``_job_id``: Integer identifier of this managed job.
-    - ``_dag_yaml`` / ``_dag`` / ``_dag_name``: The job definition and metadata.
+    - ``_dag`` / ``_dag_name``: The job definition and metadata loaded from the
+      database-backed job YAML.
     - ``_backend``: Backend used to launch and manage clusters.
     - ``_pool``: Optional pool name if using a cluster pool.
     - ``starting`` / ``starting_lock`` / ``starting_signal``: Shared scheduler
@@ -104,7 +119,6 @@ class JobsController:
     def __init__(
         self,
         job_id: int,
-        dag_yaml: str,
         starting: Set[int],
         starting_lock: asyncio.Lock,
         starting_signal: asyncio.Condition,
@@ -114,7 +128,6 @@ class JobsController:
 
         Args:
             job_id: Integer ID of the managed job.
-            dag_yaml: Path to the YAML file containing the chain DAG to run.
             starting: Shared set of job IDs currently in the STARTING phase,
                 used to limit concurrent launches.
             starting_lock: ``asyncio.Lock`` guarding access to the shared
@@ -130,12 +143,11 @@ class JobsController:
         self.starting_lock = starting_lock
         self.starting_signal = starting_signal
 
-        logger.info(f'Initializing JobsController for job_id={job_id}, '
-                    f'dag_yaml={dag_yaml}')
+        logger.info('Initializing JobsController for job_id=%s', job_id)
 
         self._job_id = job_id
-        self._dag_yaml = dag_yaml
-        self._dag, self._dag_name = _get_dag_and_name(dag_yaml)
+        self._dag = _get_dag(job_id)
+        self._dag_name = self._dag.name
         logger.info(f'Loaded DAG: {self._dag}')
 
         self._backend = cloud_vm_ray_backend.CloudVmRayBackend()
@@ -779,8 +791,11 @@ class JobsController:
                 task=self._dag.tasks[task_id]))
 
 
-class Controller:
-    """Controller for managing jobs."""
+class ControllerManager:
+    """Main loop for a job controller process.
+
+    Many jobs will be handled by this, each by a single JobController.
+    """
 
     def __init__(self, controller_uuid: str) -> None:
         self._controller_uuid = controller_uuid
@@ -799,10 +814,7 @@ class Controller:
 
         self._pid = os.getpid()
 
-    async def _cleanup(self,
-                       job_id: int,
-                       dag_yaml: str,
-                       pool: Optional[str] = None):
+    async def _cleanup(self, job_id: int, pool: Optional[str] = None):
         """Clean up the cluster(s) and storages.
 
         (1) Clean up the succeeded task(s)' ephemeral storage. The storage has
@@ -892,7 +904,7 @@ class Controller:
             if error is not None:
                 raise error
 
-        dag, _ = _get_dag_and_name(dag_yaml)
+        dag = _get_dag(job_id)
         error = None
         for task in dag.tasks:
             # most things in this function are blocking
@@ -911,57 +923,45 @@ class Controller:
     @context.contextual_async
     async def run_job_loop(self,
                            job_id: int,
-                           dag_yaml: str,
                            log_file: str,
-                           env_file_path: Optional[str] = None,
                            pool: Optional[str] = None):
         """Background task that runs the job loop."""
         ctx = context.get()
         assert ctx is not None, 'Context is not initialized'
         ctx.redirect_log(pathlib.Path(log_file))
 
-        logger.info(f'Starting job loop for {job_id}')
-        logger.info(f'  dag_yaml={dag_yaml}')
-        logger.info(f'  log_file={log_file}')
-        logger.info(f'  env_file_path={env_file_path}')
-        logger.info(f'  pool={pool}')
+        logger.info('Starting job loop for %s', job_id)
+        logger.info('  log_file=%s', log_file)
+        logger.info('  pool=%s', pool)
         logger.info(f'From controller {self._controller_uuid}')
         logger.info(f'  pid={self._pid}')
 
-        # Load and apply environment variables from the job's environment file
-        if env_file_path and os.path.exists(env_file_path):
+        env_content = file_content_utils.get_job_env_content(job_id)
+        if env_content:
             try:
-                # Load environment variables from the file
-                env_vars = dotenv.dotenv_values(env_file_path)
-                logger.info(f'Loading environment from {env_file_path}: '
-                            f'{list(env_vars.keys())}')
-
-                # Apply environment variables to the job's context
+                env_vars = dotenv.dotenv_values(stream=io.StringIO(env_content))
+                logger.info('Loading %d environment variables for job %s',
+                            len(env_vars), job_id)
                 if ctx is not None:
                     for key, value in env_vars.items():
                         if value is not None:
                             ctx.override_envs({key: value})
-                            logger.debug(
-                                f'Set environment variable: {key}={value}')
-                    # Reload the skypilot config for this context to make sure
-                    # the latest config is used.
+                            logger.debug('Set environment variable: %s=%s', key,
+                                         value)
                     skypilot_config.reload_config()
-                else:
-                    logger.error(
-                        'Context is None, cannot set environment variables')
+                else:  # pragma: no cover - defensive
+                    logger.error('Context is None, cannot set environment '
+                                 'variables')
             except Exception as e:  # pylint: disable=broad-except
                 logger.error(
-                    f'Failed to load environment file {env_file_path}: {e}')
-        elif env_file_path:
-            logger.error(f'Environment file not found: {env_file_path}')
+                    'Failed to load environment variables for job %s: '
+                    '%s', job_id, e)
 
         cancelling = False
         try:
-            logger.info(f'Starting job loop for {job_id}')
-
-            controller = JobsController(job_id, dag_yaml, self.starting,
-                                        self._job_tasks_lock,
-                                        self._starting_signal, pool)
+            controller = JobController(job_id, self.starting,
+                                       self._job_tasks_lock,
+                                       self._starting_signal, pool)
 
             async with self._job_tasks_lock:
                 if job_id in self.job_tasks:
@@ -976,7 +976,7 @@ class Controller:
             await task
         except asyncio.CancelledError:
             logger.info(f'Job {job_id} was cancelled')
-            dag, _ = _get_dag_and_name(dag_yaml)
+            dag = _get_dag(job_id)
             task_id, _ = await (
                 managed_job_state.get_latest_task_id_status_async(job_id))
             assert task_id is not None, job_id
@@ -994,7 +994,7 @@ class Controller:
             raise
         finally:
             try:
-                await self._cleanup(job_id, dag_yaml=dag_yaml, pool=pool)
+                await self._cleanup(job_id, pool=pool)
                 logger.info(
                     f'Cluster of managed job {job_id} has been cleaned up.')
             except Exception as e:  # pylint: disable=broad-except
@@ -1056,29 +1056,23 @@ class Controller:
     async def start_job(
         self,
         job_id: int,
-        dag_yaml: str,
-        env_file_path: Optional[str] = None,
         pool: Optional[str] = None,
     ):
         """Start a new job.
 
         Args:
             job_id: The ID of the job to start.
-            dag_yaml: Path to the YAML file containing the DAG definition.
-            env_file_path: Optional path to environment file for the job.
         """
         # Create log file path for job output redirection
         log_dir = os.path.expanduser(jobs_constants.JOBS_CONTROLLER_LOGS_DIR)
         os.makedirs(log_dir, exist_ok=True)
         log_file = os.path.join(log_dir, f'{job_id}.log')
 
-        logger.info(f'Starting job {job_id} with dag_yaml={dag_yaml}, '
-                    f'env_file_path={env_file_path}, and log_file={log_file}')
+        logger.info(f'Starting job {job_id} with log_file={log_file}')
 
         async with self._job_tasks_lock:
             self.starting.add(job_id)
-        await create_background_task(
-            self.run_job_loop(job_id, dag_yaml, log_file, env_file_path, pool))
+        await create_background_task(self.run_job_loop(job_id, log_file, pool))
 
         logger.info(f'Job {job_id} started successfully')
 
@@ -1151,8 +1145,6 @@ class Controller:
 
             logger.info(f'Claiming job {waiting_job["job_id"]}')
             job_id = waiting_job['job_id']
-            dag_yaml_path = waiting_job['dag_yaml_path']
-            env_file_path = waiting_job.get('env_file_path')
             pool = waiting_job.get('pool', None)
 
             cancels = os.listdir(jobs_constants.CONSOLIDATED_SIGNAL_PATH)
@@ -1172,7 +1164,7 @@ class Controller:
                             job_id=job_id, task_id=None, task=None))
                     continue
 
-            await self.start_job(job_id, dag_yaml_path, env_file_path, pool)
+            await self.start_job(job_id, pool)
 
 
 async def main(controller_uuid: str):
@@ -1180,7 +1172,7 @@ async def main(controller_uuid: str):
 
     context_utils.hijack_sys_attrs()
 
-    controller = Controller(controller_uuid)
+    controller = ControllerManager(controller_uuid)
 
     # Will happen multiple times, who cares though
     os.makedirs(jobs_constants.CONSOLIDATED_SIGNAL_PATH, exist_ok=True)
@@ -1199,7 +1191,10 @@ async def main(controller_uuid: str):
     # Will loop forever, do it in the background
     cancel_job_task = asyncio.create_task(controller.cancel_job())
     monitor_loop_task = asyncio.create_task(controller.monitor_loop())
-
+    # Run the garbage collector in a dedicated daemon thread to avoid affecting
+    # the main event loop.
+    gc_thread = threading.Thread(target=log_gc.elect_for_log_gc, daemon=True)
+    gc_thread.start()
     try:
         await asyncio.gather(cancel_job_task, monitor_loop_task)
     except Exception as e:  # pylint: disable=broad-except

sky/jobs/file_content_utils.py

@@ -0,0 +1,80 @@
+"""Utilities for managing managed job file content.
+
+The helpers in this module fetch job file content (DAG YAML/env files) from the
+database-first storage added for managed jobs, transparently falling back to
+legacy on-disk paths when needed. Consumers should prefer the string-based
+helpers so controllers never have to rely on local disk state.
+"""
+
+import os
+from typing import Optional
+
+from sky import sky_logging
+from sky.jobs import state as managed_job_state
+
+logger = sky_logging.init_logger(__name__)
+
+
+def get_job_dag_content(job_id: int) -> Optional[str]:
+    """Get DAG YAML content for a job from database or disk.
+
+    Args:
+        job_id: The job ID
+
+    Returns:
+        DAG YAML content as string, or None if not found
+    """
+    file_info = managed_job_state.get_job_file_contents(job_id)
+
+    # Prefer content stored in the database
+    if file_info['dag_yaml_content'] is not None:
+        return file_info['dag_yaml_content']
+
+    # Fallback to disk path for backward compatibility
+    dag_yaml_path = file_info.get('dag_yaml_path')
+    if dag_yaml_path and os.path.exists(dag_yaml_path):
+        try:
+            with open(dag_yaml_path, 'r', encoding='utf-8') as f:
+                content = f.read()
+                logger.debug('Loaded DAG YAML from disk for job %s: %s', job_id,
+                             dag_yaml_path)
+                return content
+        except (FileNotFoundError, IOError, OSError) as e:
+            logger.warning(
+                f'Failed to read DAG YAML from disk {dag_yaml_path}: {e}')
+
+    logger.warning(f'DAG YAML content not found for job {job_id}')
+    return None
+
+
+def get_job_env_content(job_id: int) -> Optional[str]:
+    """Get environment file content for a job from database or disk.
+
+    Args:
+        job_id: The job ID
+
+    Returns:
+        Environment file content as string, or None if not found
+    """
+    file_info = managed_job_state.get_job_file_contents(job_id)
+
+    # Prefer content stored in the database
+    if file_info['env_file_content'] is not None:
+        return file_info['env_file_content']
+
+    # Fallback to disk path for backward compatibility
+    env_file_path = file_info.get('env_file_path')
+    if env_file_path and os.path.exists(env_file_path):
+        try:
+            with open(env_file_path, 'r', encoding='utf-8') as f:
+                content = f.read()
+                logger.debug('Loaded environment file from disk for job %s: %s',
+                             job_id, env_file_path)
+                return content
+        except (FileNotFoundError, IOError, OSError) as e:
+            logger.warning(
+                f'Failed to read environment file from disk {env_file_path}: '
+                f'{e}')
+
+    # Environment file is optional, so don't warn if not found
+    return None

sky/jobs/log_gc.py

@@ -0,0 +1,201 @@
+"""Log garbage collection for managed jobs."""
+
+import asyncio
+from datetime import datetime
+import os
+import pathlib
+import shutil
+import time
+
+import anyio
+import filelock
+
+from sky import sky_logging
+from sky import skypilot_config
+from sky.jobs import constants as managed_job_constants
+from sky.jobs import state as managed_job_state
+from sky.jobs import utils as managed_job_utils
+from sky.utils import context
+from sky.utils import context_utils
+
+logger = sky_logging.init_logger(__name__)
+
+# Filelock for garbage collector leader election.
+_JOB_CONTROLLER_GC_LOCK_PATH = os.path.expanduser(
+    '~/.sky/locks/job_controller_gc.lock')
+
+_DEFAULT_TASK_LOGS_GC_RETENTION_HOURS = 24 * 7
+_DEFAULT_CONTROLLER_LOGS_GC_RETENTION_HOURS = 24 * 7
+
+_LEAST_FREQUENT_GC_INTERVAL_SECONDS = 3600
+_MOST_FREQUENT_GC_INTERVAL_SECONDS = 30
+
+
+def _next_gc_interval(retention_seconds: int) -> int:
+    """Get the next GC interval."""
+    # Run the GC at least per hour to ensure hourly accuracy and
+    # at most per 30 seconds (when retention_seconds is small) to
+    # avoid too frequent cleanup.
+    return max(min(retention_seconds, _LEAST_FREQUENT_GC_INTERVAL_SECONDS),
+               _MOST_FREQUENT_GC_INTERVAL_SECONDS)
+
+
+async def gc_controller_logs_for_job():
+    """Garbage collect job and controller logs."""
+    while True:
+        skypilot_config.reload_config()
+        controller_logs_retention = skypilot_config.get_nested(
+            ('jobs', 'controller', 'controller_logs_gc_retention_hours'),
+            _DEFAULT_CONTROLLER_LOGS_GC_RETENTION_HOURS) * 3600
+        # Negative value disables the GC
+        if controller_logs_retention >= 0:
+            logger.info(f'GC controller logs for job: retention '
+                        f'{controller_logs_retention} seconds')
+            try:
+                finished = False
+                while not finished:
+                    finished = await _clean_controller_logs_with_retention(
+                        controller_logs_retention)
+            except asyncio.CancelledError:
+                logger.info('Managed jobs logs GC task cancelled')
+                break
+            except Exception as e:  # pylint: disable=broad-except
+                logger.error(f'Error GC controller logs for job: {e}',
+                             exc_info=True)
+        else:
+            logger.info('Controller logs GC is disabled')
+
+        interval = _next_gc_interval(controller_logs_retention)
+        logger.info('Next controller logs GC is scheduled after '
+                    f'{interval} seconds')
+        await asyncio.sleep(interval)
+
+
+async def gc_task_logs_for_job():
+    """Garbage collect task logs for job."""
+    while True:
+        skypilot_config.reload_config()
+        task_logs_retention = skypilot_config.get_nested(
+            ('jobs', 'controller', 'task_logs_gc_retention_hours'),
+            _DEFAULT_TASK_LOGS_GC_RETENTION_HOURS) * 3600
+        # Negative value disables the GC
+        if task_logs_retention >= 0:
+            logger.info('GC task logs for job: '
+                        f'retention {task_logs_retention} seconds')
+            try:
+                finished = False
+                while not finished:
+                    finished = await _clean_task_logs_with_retention(
+                        task_logs_retention)
+            except asyncio.CancelledError:
+                logger.info('Task logs GC task cancelled')
+                break
+            except Exception as e:  # pylint: disable=broad-except
+                logger.error(f'Error GC task logs for job: {e}', exc_info=True)
+        else:
+            logger.info('Controller logs GC is disabled')
+
+        interval = _next_gc_interval(task_logs_retention)
+        logger.info(f'Next task logs GC is scheduled after {interval} seconds')
+        await asyncio.sleep(_next_gc_interval(task_logs_retention))
+
+
+async def _clean_controller_logs_with_retention(retention_seconds: int,
+                                                batch_size: int = 100):
+    """Clean controller logs with retention.
+
+    Returns:
+        Whether the GC of this round has finished, False means there might
+        still be more controller logs to clean.
+    """
+    assert batch_size > 0, 'Batch size must be positive'
+    jobs = await managed_job_state.get_controller_logs_to_clean_async(
+        retention_seconds, batch_size=batch_size)
+    job_ids_to_update = []
+    for job in jobs:
+        job_ids_to_update.append(job['job_id'])
+        log_file = managed_job_utils.controller_log_file_for_job(job['job_id'])
+        cleaned_at = time.time()
+        if await anyio.Path(log_file).exists():
+            ts_str = datetime.fromtimestamp(cleaned_at).strftime(
+                '%Y-%m-%d %H:%M:%S')
+            msg = f'Controller log has been cleaned at {ts_str}.'
+            # Sync down logs will reference to this file directly, so we
+            # keep the file and delete the content.
+            # TODO(aylei): refactor sync down logs if the inode usage
+            # becomes an issue.
+            async with await anyio.open_file(log_file, 'w',
+                                             encoding='utf-8') as f:
+                await f.write(msg + '\n')
+    # Batch the update, the timestamp will be not accurate but it's okay.
+    await managed_job_state.set_controller_logs_cleaned_async(
+        job_ids=job_ids_to_update, logs_cleaned_at=time.time())
+    complete = len(jobs) < batch_size
+    logger.info(f'Cleaned {len(jobs)} controller logs with retention '
+                f'{retention_seconds} seconds, complete: {complete}')
+    return complete
+
+
+async def _clean_task_logs_with_retention(retention_seconds: int,
+                                          batch_size: int = 100):
+    """Clean task logs with retention.
+
+    Returns:
+        Whether the GC of this round has finished, False means there might
+        still be more task logs to clean.
+    """
+    assert batch_size > 0, 'Batch size must be positive'
+    tasks = await managed_job_state.get_task_logs_to_clean_async(
+        retention_seconds, batch_size=batch_size)
+    tasks_to_update = []
+    for task in tasks:
+        local_log_file = anyio.Path(task['local_log_file'])
+        # We assume the log directory has the following layout:
+        # task-id/
+        #   - run.log
+        #   - tasks/
+        #     - run.log
+        # and also remove the tasks directory on cleanup.
+        task_log_dir = local_log_file.parent.joinpath('tasks')
+        await local_log_file.unlink(missing_ok=True)
+        await context_utils.to_thread(shutil.rmtree,
+                                      str(task_log_dir),
+                                      ignore_errors=True)
+        # We have at least once semantic guarantee for the cleanup here.
+        tasks_to_update.append((task['job_id'], task['task_id']))
+    await managed_job_state.set_task_logs_cleaned_async(
+        tasks=list(tasks_to_update), logs_cleaned_at=time.time())
+    complete = len(tasks) < batch_size
+    logger.info(f'Cleaned {len(tasks)} task logs with retention '
+                f'{retention_seconds} seconds, complete: {complete}')
+    return complete
+
+
+@context.contextual_async
+async def run_log_gc():
+    """Run the log garbage collector."""
+    log_dir = os.path.expanduser(managed_job_constants.JOBS_CONTROLLER_LOGS_DIR)
+    os.makedirs(log_dir, exist_ok=True)
+    log_path = os.path.join(log_dir, 'garbage_collector.log')
+    # Remove previous log file
+    await anyio.Path(log_path).unlink(missing_ok=True)
+    ctx = context.get()
+    assert ctx is not None, 'Context is not initialized'
+    ctx.redirect_log(pathlib.Path(log_path))
+    gc_controller_logs_for_job_task = asyncio.create_task(
+        gc_controller_logs_for_job())
+    gc_task_logs_for_job_task = asyncio.create_task(gc_task_logs_for_job())
+    await asyncio.gather(gc_controller_logs_for_job_task,
+                         gc_task_logs_for_job_task)
+
+
+def elect_for_log_gc():
+    """Use filelock to elect for the log garbage collector.
+
+    The log garbage collector runs in the controller process to avoid the
+    overhead of launching a new process and the lifecycle management, the
+    threads that does not elected as the log garbage collector just wait.
+    on the filelock and bring trivial overhead.
+    """
+    with filelock.FileLock(_JOB_CONTROLLER_GC_LOCK_PATH):
+        asyncio.run(run_log_gc())

sky/jobs/scheduler.py

@@ -263,6 +263,7 @@ def maybe_start_controllers(from_scheduler: bool = False) -> None:
 
             if started > 0:
                 logger.info(f'Started {started} controllers')
+
     except filelock.Timeout:
         # If we can't get the lock, just exit. The process holding the lock
         # should launch any pending jobs.
@@ -289,8 +290,20 @@ def submit_job(job_id: int, dag_yaml_path: str, original_user_yaml_path: str,
             maybe_start_controllers(from_scheduler=True)
             return
 
-    state.scheduler_set_waiting(job_id, dag_yaml_path, original_user_yaml_path,
-                                env_file_path, priority)
+    with open(dag_yaml_path, 'r', encoding='utf-8') as dag_file:
+        dag_yaml_content = dag_file.read()
+    with open(original_user_yaml_path, 'r',
+              encoding='utf-8') as original_user_yaml_file:
+        original_user_yaml_content = original_user_yaml_file.read()
+    with open(env_file_path, 'r', encoding='utf-8') as env_file:
+        env_file_content = env_file.read()
+    logger.debug(f'Storing job {job_id} file contents in database '
+                 f'(DAG bytes={len(dag_yaml_content)}, '
+                 f'original user yaml bytes={len(original_user_yaml_content)}, '
+                 f'env bytes={len(env_file_content)}).')
+    state.scheduler_set_waiting(job_id, dag_yaml_content,
+                                original_user_yaml_content, env_file_content,
+                                priority)
     if state.get_ha_recovery_script(job_id) is None:
         # the run command is just the command that called scheduler
         run = (f'source {env_file_path} && '