skypilot-nightly 1.0.0.dev20250723__py3-none-any.whl → 1.0.0.dev20250725__py3-none-any.whl

sky/logs/aws.py

@@ -0,0 +1,276 @@
+"""AWS CloudWatch logging agent."""
+
+from typing import Any, Dict, Optional
+
+import pydantic
+
+from sky.logs.agent import FluentbitAgent
+from sky.skylet import constants
+from sky.utils import common_utils
+from sky.utils import resources_utils
+
+
+class _CloudwatchLoggingConfig(pydantic.BaseModel):
+    """Configuration for AWS CloudWatch logging agent."""
+    region: Optional[str] = None
+    credentials_file: Optional[str] = None
+    log_group_name: str = 'skypilot-logs'
+    log_stream_prefix: str = 'skypilot-'
+    auto_create_group: bool = True
+    additional_tags: Optional[Dict[str, str]] = None
+
+
+class _CloudWatchOutputConfig(pydantic.BaseModel):
+    """Auxiliary model for building CloudWatch output config in YAML.
+
+    Ref: https://docs.fluentbit.io/manual/pipeline/outputs/cloudwatch
+    """
+    name: str = 'cloudwatch_logs'
+    match: str = '*'
+    region: Optional[str] = None
+    log_group_name: Optional[str] = None
+    log_stream_prefix: Optional[str] = None
+    auto_create_group: bool = True
+    additional_tags: Optional[Dict[str, str]] = None
+
+    def to_dict(self) -> Dict[str, Any]:
+        config = self.model_dump(exclude_none=True)
+        if 'auto_create_group' in config:
+            config['auto_create_group'] = 'true' if config[
+                'auto_create_group'] else 'false'
+        return config
+
+
+class CloudwatchLoggingAgent(FluentbitAgent):
+    """AWS CloudWatch logging agent.
+
+    This agent forwards logs from SkyPilot clusters to AWS CloudWatch using
+    Fluent Bit. It supports authentication via IAM roles (preferred), AWS
+    credentials file, or environment variables.
+
+    Example configuration:
+    ```yaml
+    logs:
+      store: aws
+      aws:
+        region: us-west-2
+        log_group_name: skypilot-logs
+        log_stream_prefix: my-cluster-
+        auto_create_group: true
+    ```
+    """
+
+    def __init__(self, config: Dict[str, Any]):
+        """Initialize the CloudWatch logging agent.
+
+        Args:
+            config: The configuration for the CloudWatch logging agent.
+                   See the class docstring for the expected format.
+        """
+        self.config = _CloudwatchLoggingConfig(**config)
+        super().__init__()
+
+    def get_setup_command(self,
+                          cluster_name: resources_utils.ClusterName) -> str:
+        """Get the command to set up the CloudWatch logging agent.
+
+        Args:
+            cluster_name: The name of the cluster.
+
+        Returns:
+            The command to set up the CloudWatch logging agent.
+        """
+
+        if self.config.credentials_file:
+            credential_path = self.config.credentials_file
+
+        # Set AWS credentials and check whether credentials are valid.
+        # CloudWatch plugin supports IAM roles, credentials file, and
+        # environment variables. We prefer IAM roles when available
+        # (on EC2 instances). If credentials file is provided, we use
+        # it. Otherwise, we check if credentials are available in
+        # the environment.
+        pre_cmd = ''
+        if self.config.credentials_file:
+            pre_cmd = (
+                f'export AWS_SHARED_CREDENTIALS_FILE={credential_path}; '
+                f'if [ ! -f {credential_path} ]; then '
+                f'echo "ERROR: AWS credentials file {credential_path} '
+                f'not found. Please check if the file exists and is '
+                f'accessible." && exit 1; '
+                f'fi; '
+                f'if ! grep -q "\\[.*\\]" {credential_path} || '
+                f'! grep -q "aws_access_key_id" {credential_path}; then '
+                f'echo "ERROR: AWS credentials file {credential_path} is '
+                f'invalid. It should contain a profile section '
+                f'[profile_name] and aws_access_key_id." && exit 1; '
+                f'fi;')
+        else:
+            # Check if we're running on EC2 with an IAM role or if
+            # AWS credentials are available in the environment
+            pre_cmd = (
+                'if ! curl -s -m 1 http://169.254.169.254'
+                '/latest/meta-data/iam/security-credentials/ > /dev/null; '
+                'then '
+                # failed EC2 check, look for env vars
+                'if [ -z "$AWS_ACCESS_KEY_ID" ] || '
+                '[ -z "$AWS_SECRET_ACCESS_KEY" ]; then '
+                'echo "ERROR: AWS CloudWatch logging configuration error. '
+                'Not running on EC2 with IAM role and AWS credentials not '
+                'found in environment. Please do one of the following: '
+                '1. Run on an EC2 instance with an IAM role that has '
+                'CloudWatch permissions, 2. Set AWS_ACCESS_KEY_ID and '
+                'AWS_SECRET_ACCESS_KEY environment variables, or '
+                '3. Provide a credentials file via logs.aws.credentials_file '
+                'in SkyPilot config." && exit 1; '
+                'fi; '
+                'fi;')
+
+        # If region is specified, set it in the environment
+        if self.config.region:
+            pre_cmd += f' export AWS_REGION={self.config.region};'
+        else:
+            # If region is not specified, check if it's available in
+            # the environment or credentials file
+            pre_cmd += (
+                ' if [ -z "$AWS_REGION" ] && '
+                '[ -z "$AWS_DEFAULT_REGION" ]; then '
+                'echo "WARNING: AWS region not specified in configuration or '
+                'environment. CloudWatch logging may fail if the region '
+                'cannot be determined. Consider setting logs.aws.region in '
+                'SkyPilot config."; '
+                'fi; ')
+
+        # Add a test command to verify AWS credentials work with CloudWatch
+        pre_cmd += (
+            ' echo "Verifying AWS CloudWatch access..."; '
+            'if command -v aws > /dev/null; then '
+            'aws cloudwatch list-metrics --namespace AWS/Logs --max-items 1 '
+            '> /dev/null 2>&1 || '
+            '{ echo "ERROR: Failed to access AWS CloudWatch. Please check '
+            'your credentials and permissions."; '
+            'echo "The IAM role or user must have cloudwatch:ListMetrics '
+            'and logs:* permissions."; '
+            'exit 1; }; '
+            'else echo "AWS CLI not installed, skipping CloudWatch access '
+            'verification."; '
+            'fi; ')
+
+        return pre_cmd + ' ' + super().get_setup_command(cluster_name)
+
+    def fluentbit_config(self,
+                         cluster_name: resources_utils.ClusterName) -> str:
+        """Get the Fluent Bit configuration for CloudWatch.
+
+        This overrides the base method to add a fallback output for local file
+        logging in case CloudWatch logging fails.
+
+        Args:
+            cluster_name: The name of the cluster.
+
+        Returns:
+            The Fluent Bit configuration as a YAML string.
+        """
+        display_name = cluster_name.display_name
+        unique_name = cluster_name.name_on_cloud
+        # Build tags for the log stream
+        tags = {
+            'skypilot.cluster_name': display_name,
+            'skypilot.cluster_id': unique_name,
+        }
+
+        # Add additional tags if provided
+        if self.config.additional_tags:
+            tags.update(self.config.additional_tags)
+
+        log_processors = []
+        for key, value in tags.items():
+            log_processors.append({
+                'name': 'content_modifier',
+                'action': 'upsert',
+                'key': key,
+                'value': value
+            })
+
+        cfg_dict = {
+            'pipeline': {
+                'inputs': [{
+                    'name': 'tail',
+                    'path': f'{constants.SKY_LOGS_DIRECTORY}/*/*.log',
+                    'path_key': 'log_path',
+                    # Shorten the refresh interval from 60s to 1s since every
+                    # job creates a new log file and we must be responsive
+                    # for this: the VM might be autodown within a minute
+                    # right after the job completion.
+                    'refresh_interval': 1,
+                    'processors': {
+                        'logs': log_processors,
+                    }
+                }],
+                'outputs': [self.fluentbit_output_config(cluster_name)],
+            }
+        }
+
+        # Add fallback outputs for graceful failure handling
+        cfg_dict = self.add_fallback_outputs(cfg_dict)
+
+        return common_utils.dump_yaml_str(cfg_dict)
+
+    def add_fallback_outputs(self, cfg_dict: Dict[str, Any]) -> Dict[str, Any]:
+        """Add fallback outputs to the Fluent Bit configuration.
+
+        This adds a local file output as a fallback in case
+        CloudWatch logging fails.
+
+        Args:
+            cfg_dict: The Fluent Bit configuration dictionary.
+
+        Returns:
+            The updated configuration dictionary.
+        """
+        # Add a local file output as a fallback
+        fallback_output = {
+            'name': 'file',
+            'match': '*',
+            'path': '/tmp/skypilot_logs_fallback.log',
+            'format': 'out_file',
+        }
+
+        # Add the fallback output to the configuration
+        cfg_dict['pipeline']['outputs'].append(fallback_output)
+
+        return cfg_dict
+
+    def fluentbit_output_config(
+            self, cluster_name: resources_utils.ClusterName) -> Dict[str, Any]:
+        """Get the Fluent Bit output configuration for CloudWatch.
+
+        Args:
+            cluster_name: The name of the cluster.
+
+        Returns:
+            The Fluent Bit output configuration for CloudWatch.
+        """
+        unique_name = cluster_name.name_on_cloud
+
+        # Format the log stream name to include cluster information
+        # This helps with identifying logs in CloudWatch
+        log_stream_prefix = f'{self.config.log_stream_prefix}{unique_name}-'
+
+        # Create the CloudWatch output configuration with error handling options
+        return _CloudWatchOutputConfig(
+            region=self.config.region,
+            log_group_name=self.config.log_group_name,
+            log_stream_prefix=log_stream_prefix,
+            auto_create_group=self.config.auto_create_group,
+        ).to_dict()
+
+    def get_credential_file_mounts(self) -> Dict[str, str]:
+        """Get the credential file mounts for the CloudWatch logging agent.
+
+        Returns:
+            A dictionary mapping local credential file paths to remote paths.
+        """
+        if self.config.credentials_file:
+            return {self.config.credentials_file: self.config.credentials_file}
+        return {}

sky/server/common.py

@@ -13,12 +13,14 @@ import shutil
 import subprocess
 import sys
 import tempfile
+import threading
 import time
 import typing
 from typing import Any, Dict, Literal, Optional, Tuple, Union
 from urllib import parse
 import uuid
 
+import cachetools
 import colorama
 import filelock
 
@@ -132,6 +134,8 @@ def get_api_cookie_jar() -> requests.cookies.RequestsCookieJar:
 def set_api_cookie_jar(cookie_jar: CookieJar,
                        create_if_not_exists: bool = True) -> None:
     """Updates the file cookie jar with the given cookie jar."""
+    if len(cookie_jar) == 0:
+        return
     cookie_path = get_api_cookie_jar_path()
     if not cookie_path.exists() and not create_if_not_exists:
         # if the file doesn't exist and we don't want to create it, do nothing
@@ -274,6 +278,10 @@ def _handle_non_200_server_status(
     return ApiServerInfo(status=ApiServerStatus.UNHEALTHY)
 
 
+@cachetools.cached(cache=cachetools.TTLCache(maxsize=10,
+                                             ttl=5.0,
+                                             timer=time.time),
+                   lock=threading.RLock())
 def get_api_server_status(endpoint: Optional[str] = None) -> ApiServerInfo:
     """Retrieve the status of the API server.
 
@@ -351,7 +359,9 @@ def get_api_server_status(endpoint: Optional[str] = None) -> ApiServerInfo:
                                      error=version_info.error)
 
             cookies = get_cookies_from_response(response)
-            set_api_cookie_jar(cookies, create_if_not_exists=False)
+            # Save or refresh the cookie jar in case of session affinity and
+            # OAuth.
+            set_api_cookie_jar(cookies, create_if_not_exists=True)
             return server_info
         except (json.JSONDecodeError, AttributeError) as e:
             # Try to check if we got redirected to a login page.
@@ -409,6 +419,7 @@ def _start_api_server(deploy: bool = False,
     server_url = get_server_url(host)
     assert server_url in AVAILABLE_LOCAL_API_SERVER_URLS, (
         f'server url {server_url} is not a local url')
+
     with rich_utils.client_status('Starting SkyPilot API server, '
                                   f'view logs at {constants.API_SERVER_LOGS}'):
         logger.info(f'{colorama.Style.DIM}Failed to connect to '
@@ -484,6 +495,8 @@ def _start_api_server(deploy: bool = False,
                         'SkyPilot API server process exited unexpectedly.\n'
                         f'View logs at: {constants.API_SERVER_LOGS}')
             try:
+                # Clear the cache to ensure fresh checks during startup
+                get_api_server_status.cache_clear()  # type: ignore
                 check_server_healthy()
             except exceptions.APIVersionMismatchError:
                 raise

sky/server/requests/payloads.py

@@ -203,17 +203,33 @@ class DagRequestBody(RequestBody):
         return kwargs
 
 
-class ValidateBody(DagRequestBody):
+class DagRequestBodyWithRequestOptions(DagRequestBody):
+    """Request body base class for endpoints with a dag and request options."""
+    request_options: Optional[admin_policy.RequestOptions]
+
+    def get_request_options(self) -> Optional[admin_policy.RequestOptions]:
+        """Get the request options."""
+        if self.request_options is None:
+            return None
+        if isinstance(self.request_options, dict):
+            return admin_policy.RequestOptions(**self.request_options)
+        return self.request_options
+
+    def to_kwargs(self) -> Dict[str, Any]:
+        kwargs = super().to_kwargs()
+        kwargs['request_options'] = self.get_request_options()
+        return kwargs
+
+
+class ValidateBody(DagRequestBodyWithRequestOptions):
     """The request body for the validate endpoint."""
     dag: str
-    request_options: Optional[admin_policy.RequestOptions]
 
 
-class OptimizeBody(DagRequestBody):
+class OptimizeBody(DagRequestBodyWithRequestOptions):
     """The request body for the optimize endpoint."""
     dag: str
     minimize: common_lib.OptimizeTarget = common_lib.OptimizeTarget.COST
-    request_options: Optional[admin_policy.RequestOptions]
 
 
 class LaunchBody(RequestBody):

sky/server/rest.py

@@ -89,6 +89,12 @@ def retry_transient_errors(max_retries: int = 3,
                 for retry_cnt in range(max_retries):
                     try:
                         return func(*args, **kwargs)
+                    # Occurs when the server proactively interrupts the request
+                    # during rolling update, we can retry immediately on the
+                    # new replica.
+                    except exceptions.RequestInterruptedError:
+                        logger.debug('Request interrupted. Retry immediately.')
+                        continue
                     except Exception as e:  # pylint: disable=broad-except
                         if retry_cnt >= max_retries - 1:
                             # Retries exhausted.

sky/server/server.py

@@ -827,7 +827,8 @@ async def validate(validate_body: payloads.ValidateBody) -> None:
         # added RTTs. For now, we stick to doing the validation inline in the
         # server thread.
         with admin_policy_utils.apply_and_use_config_in_current_request(
-                dag, request_options=validate_body.request_options) as dag:
+                dag,
+                request_options=validate_body.get_request_options()) as dag:
             # Skip validating workdir and file_mounts, as those need to be
             # validated after the files are uploaded to the SkyPilot API server
             # with `upload_mounts_to_api_server`.
@@ -1763,6 +1764,9 @@ if __name__ == '__main__':
 
     from sky.server import uvicorn as skyuvicorn
 
+    # Initialize global user state db
+    global_user_state.initialize_and_get_db()
+    # Initialize request db
     requests_lib.reset_db_and_logs()
 
     parser = argparse.ArgumentParser()

sky/templates/aws-ray.yml.j2

@@ -19,7 +19,7 @@ docker:
     username: |-
       {{docker_login_config.username}}
     password: |-
-      {{docker_login_config.password}}
+      {{docker_login_config.password | indent(6) }}
     server: |-
       {{docker_login_config.server}}
   {%- endif %}
@@ -131,6 +131,12 @@ available_node_types:
           - systemctl disable apt-daily.timer apt-daily-upgrade.timer unattended-upgrades.service
           - systemctl mask apt-daily.service apt-daily-upgrade.service unattended-upgrades.service
           - systemctl daemon-reload
+        {%- if runcmd %}
+        runcmd:
+        {%- for cmd in runcmd %}
+          - {{cmd}}
+        {%- endfor %}
+        {%- endif %}
       TagSpecifications:
         - ResourceType: instance
           Tags:

sky/templates/azure-ray.yml.j2

@@ -19,7 +19,7 @@ docker:
     username: |-
       {{docker_login_config.username}}
     password: |-
-      {{docker_login_config.password}}
+      {{docker_login_config.password | indent(6) }}
     server: |-
       {{docker_login_config.server}}
   {%- endif %}

sky/templates/do-ray.yml.j2

@@ -19,7 +19,7 @@ docker:
     username: |-
       {{docker_login_config.username}}
     password: |-
-      {{docker_login_config.password}}
+      {{docker_login_config.password | indent(6) }}
     server: |-
       {{docker_login_config.server}}
   {%- endif %}

sky/templates/lambda-ray.yml.j2

@@ -19,7 +19,7 @@ docker:
     username: |-
       {{docker_login_config.username}}
     password: |-
-      {{docker_login_config.password}}
+      {{docker_login_config.password | indent(6) }}
     server: |-
       {{docker_login_config.server}}
   {%- endif %}

sky/templates/nebius-ray.yml.j2

@@ -25,7 +25,7 @@ docker:
     username: |-
       {{docker_login_config.username}}
     password: |-
-      {{docker_login_config.password}}
+      {{docker_login_config.password | indent(6) }}
     server: |-
       {{docker_login_config.server}}
   {%- endif %}

sky/templates/paperspace-ray.yml.j2

@@ -19,7 +19,7 @@ docker:
     username: |-
       {{docker_login_config.username}}
     password: |-
-      {{docker_login_config.password}}
+      {{docker_login_config.password | indent(6) }}
     server: |-
       {{docker_login_config.server}}
   {%- endif %}

sky/templates/runpod-ray.yml.j2

@@ -20,7 +20,7 @@ provider:
     username: |-
       {{docker_login_config.username}}
     password: |-
-      {{docker_login_config.password}}
+      {{docker_login_config.password | indent(6) }}
     server: |-
       {{docker_login_config.server}}
   {%- endif %}

sky/utils/config_utils.py

@@ -6,6 +6,8 @@ from sky import sky_logging
 
 logger = sky_logging.init_logger(__name__)
 
+_REGION_CONFIG_CLOUDS = ['nebius', 'oci']
+
 
 class Config(Dict[str, Any]):
     """SkyPilot config that supports setting/getting values with nested keys."""
@@ -248,7 +250,7 @@ def get_cloud_config_value_from_dict(
     region_key = None
     if cloud == 'kubernetes':
         region_key = 'context_configs'
-    if cloud == 'nebius':
+    elif cloud in _REGION_CONFIG_CLOUDS:
         region_key = 'region_configs'
 
     per_context_config = None
@@ -257,7 +259,7 @@ def get_cloud_config_value_from_dict(
             keys=(cloud, region_key, region) + keys,
             default_value=None,
             override_configs=override_configs)
-        if not per_context_config and cloud == 'nebius':
+        if not per_context_config and cloud in _REGION_CONFIG_CLOUDS:
             # TODO (kyuds): Backward compatibility, remove after 0.11.0.
             per_context_config = input_config.get_nested(
                 keys=(cloud, region) + keys,
@@ -265,9 +267,9 @@ def get_cloud_config_value_from_dict(
                 override_configs=override_configs)
             if per_context_config is not None:
                 logger.info(
-                    'Nebius configuration is using the legacy format. \n'
+                    f'{cloud} configuration is using the legacy format. \n'
                     'This format will be deprecated after 0.11.0, refer to '
-                    '`https://docs.skypilot.co/en/latest/reference/config.html#nebius` '  # pylint: disable=line-too-long
+                    '`https://docs.skypilot.co/en/latest/reference/config.html` '  # pylint: disable=line-too-long
                     'for the new format. Please use `region_configs` to specify region specific configuration.'
                 )
     # if no override found for specified region

sky/utils/db/migration_utils.py

@@ -3,6 +3,7 @@
 import contextlib
 import logging
 import os
+import pathlib
 
 from alembic import command as alembic_command
 from alembic.config import Config
@@ -10,6 +11,12 @@ from alembic.runtime import migration
 import filelock
 import sqlalchemy
 
+from sky import sky_logging
+from sky import skypilot_config
+from sky.skylet import constants
+
+logger = sky_logging.init_logger(__name__)
+
 DB_INIT_LOCK_TIMEOUT_SECONDS = 10
 
 GLOBAL_USER_STATE_DB_NAME = 'state_db'
@@ -21,6 +28,21 @@ SPOT_JOBS_VERSION = '001'
 SPOT_JOBS_LOCK_PATH = '~/.sky/locks/.spot_jobs_db.lock'
 
 
+def get_engine(db_name: str):
+    conn_string = None
+    if os.environ.get(constants.ENV_VAR_IS_SKYPILOT_SERVER) is not None:
+        conn_string = skypilot_config.get_nested(('db',), None)
+    if conn_string:
+        logger.debug(f'using db URI from {conn_string}')
+        engine = sqlalchemy.create_engine(conn_string,
+                                          poolclass=sqlalchemy.NullPool)
+    else:
+        db_path = os.path.expanduser(f'~/.sky/{db_name}.db')
+        pathlib.Path(db_path).parents[0].mkdir(parents=True, exist_ok=True)
+        engine = sqlalchemy.create_engine('sqlite:///' + db_path)
+    return engine
+
+
 @contextlib.contextmanager
 def db_lock(db_name: str):
     lock_path = os.path.expanduser(f'~/.sky/locks/.{db_name}.lock')
@@ -37,7 +59,6 @@ def db_lock(db_name: str):
 
 def get_alembic_config(engine: sqlalchemy.engine.Engine, section: str):
     """Get Alembic configuration for the given section"""
-    # Use the alembic.ini file from setup_files (included in wheel)
     # From sky/utils/db/migration_utils.py -> sky/setup_files/alembic.ini
     alembic_ini_path = os.path.join(
         os.path.dirname(os.path.dirname(os.path.dirname(__file__))),
@@ -47,31 +68,29 @@ def get_alembic_config(engine: sqlalchemy.engine.Engine, section: str):
     # Override the database URL to match SkyPilot's current connection
     # Use render_as_string to get the full URL with password
     url = engine.url.render_as_string(hide_password=False)
+    # Replace % with %% to escape the % character in the URL
+    # set_section_option uses variable interpolation, which treats % as a
+    # special character.
+    # any '%' symbol not used for interpolation needs to be escaped.
+    url = url.replace('%', '%%')
     alembic_cfg.set_section_option(section, 'sqlalchemy.url', url)
 
     return alembic_cfg
 
 
-def safe_alembic_upgrade(engine: sqlalchemy.engine.Engine,
-                         alembic_config: Config, target_revision: str):
-    """Only upgrade if current version is older than target.
-
-    This handles the case where a database was created with a newer version of
-    the code and we're now running older code. Since our migrations are purely
-    additive, it's safe to run a newer database with older code.
+def needs_upgrade(engine: sqlalchemy.engine.Engine, section: str,
+                  target_revision: str):
+    """Check if the database needs to be upgraded.
 
     Args:
         engine: SQLAlchemy engine for the database
-        alembic_config: Alembic configuration object
+        section: Alembic section to upgrade (e.g., 'state_db' or 'spot_jobs_db')
         target_revision: Target revision to upgrade to (e.g., '001')
     """
-    # set alembic logger to warning level
-    alembic_logger = logging.getLogger('alembic')
-    alembic_logger.setLevel(logging.WARNING)
-
     current_rev = None
 
-    # Get the current revision from the database
+    # get alembic config for the given section
+    alembic_config = get_alembic_config(engine, section)
     version_table = alembic_config.get_section_option(
         alembic_config.config_ini_section, 'version_table', 'alembic_version')
 
@@ -81,13 +100,35 @@ def safe_alembic_upgrade(engine: sqlalchemy.engine.Engine,
         current_rev = context.get_current_revision()
 
     if current_rev is None:
-        alembic_command.upgrade(alembic_config, target_revision)
-        return
+        return True
 
     # Compare revisions - assuming they are numeric strings like '001', '002'
     current_rev_num = int(current_rev)
     target_rev_num = int(target_revision)
 
-    # only upgrade if current revision is older than target revision
-    if current_rev_num < target_rev_num:
-        alembic_command.upgrade(alembic_config, target_revision)
+    return current_rev_num < target_rev_num
+
+
+def safe_alembic_upgrade(engine: sqlalchemy.engine.Engine, section: str,
+                         target_revision: str):
+    """Upgrade the database if needed. Uses a file lock to ensure
+    that only one process tries to upgrade the database at a time.
+
+    Args:
+        engine: SQLAlchemy engine for the database
+        section: Alembic section to upgrade (e.g., 'state_db' or 'spot_jobs_db')
+        target_revision: Target revision to upgrade to (e.g., '001')
+    """
+    # set alembic logger to warning level
+    alembic_logger = logging.getLogger('alembic')
+    alembic_logger.setLevel(logging.WARNING)
+
+    alembic_config = get_alembic_config(engine, section)
+
+    # only acquire lock if db needs upgrade
+    if needs_upgrade(engine, section, target_revision):
+        with db_lock(section):
+            # check again if db needs upgrade in case another
+            # process upgraded it while we were waiting for the lock
+            if needs_upgrade(engine, section, target_revision):
+                alembic_command.upgrade(alembic_config, target_revision)