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

sky/jobs/file_content_utils.py

@@ -0,0 +1,128 @@
+"""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 import skypilot_config
+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
+
+
+def restore_job_config_file(job_id: int) -> None:
+    """Restore config file from database if SKYPILOT_CONFIG is set.
+
+    This reads the config file content from the database and writes it to the
+    path specified in the SKYPILOT_CONFIG environment variable. This ensures
+    that jobs can run on any controller, even if the original config file
+    doesn't exist on disk.
+
+    For backward compatibility with jobs submitted before config persistence was
+    implemented, we fall back to using the file if it already exists on disk.
+
+    Args:
+        job_id: The job ID
+    """
+    config_path = os.environ.get(skypilot_config.ENV_VAR_SKYPILOT_CONFIG)
+    if not config_path:
+        # No config file for this job
+        return
+
+    file_info = managed_job_state.get_job_file_contents(job_id)
+    config_content = file_info['config_file_content']
+
+    # Expand ~ in config path
+    config_path_expanded = os.path.expanduser(config_path)
+
+    if config_content is not None:
+        # Config content is in database - restore it
+        # Ensure the directory exists
+        os.makedirs(os.path.dirname(config_path_expanded), exist_ok=True)
+        # Write the config file
+        with open(config_path_expanded, 'w', encoding='utf-8') as f:
+            f.write(config_content)
+        logger.info(f'Restored config file for job {job_id} to '
+                    f'{config_path_expanded} ({len(config_content)} bytes)')
+    elif os.path.exists(config_path_expanded):
+        # Backward compatibility: config not in DB but file exists on disk
+        # This can happen for jobs submitted before config persistence
+        logger.debug(f'Config file for job {job_id} not in database, but '
+                     f'found on disk at {config_path_expanded}')
+    else:
+        # Config should exist but doesn't - warn about it
+        logger.warning(
+            f'SKYPILOT_CONFIG is set to {config_path} but config content not '
+            f'found in database or on disk for job {job_id}. The job may fail '
+            f'if it relies on custom config settings.')

sky/jobs/log_gc.py

@@ -0,0 +1,193 @@
+"""Log garbage collection for managed jobs."""
+
+from datetime import datetime
+import os
+import pathlib
+import shutil
+import threading
+import time
+
+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
+
+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)
+
+
+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 = _clean_controller_logs_with_retention(
+                        controller_logs_retention)
+            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')
+        time.sleep(interval)
+
+
+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 = _clean_task_logs_with_retention(
+                        task_logs_retention)
+            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')
+        time.sleep(_next_gc_interval(task_logs_retention))
+
+
+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 = managed_job_state.get_controller_logs_to_clean(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 os.path.exists(log_file):
+            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.
+            with open(log_file, 'w', encoding='utf-8') as f:
+                f.write(msg + '\n')
+    # Batch the update, the timestamp will be not accurate but it's okay.
+    managed_job_state.set_controller_logs_cleaned(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
+
+
+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 = managed_job_state.get_task_logs_to_clean(retention_seconds,
+                                                     batch_size=batch_size)
+    tasks_to_update = []
+    for task in tasks:
+        local_log_file = pathlib.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')
+        local_log_file.unlink(missing_ok=True)
+        shutil.rmtree(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']))
+    managed_job_state.set_task_logs_cleaned(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
+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
+    pathlib.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))
+    tasks = []
+    tasks.append(
+        threading.Thread(target=gc_controller_logs_for_job, daemon=True))
+    tasks.append(threading.Thread(target=gc_task_logs_for_job, daemon=True))
+    for task in tasks:
+        task.start()
+    for task in tasks:
+        task.join()
+
+
+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):
+        run_log_gc()