skypilot-nightly 1.0.0.dev20250616__py3-none-any.whl → 1.0.0.dev20250618__py3-none-any.whl

sky/logs/agent.py

@@ -0,0 +1,73 @@
+"""Base class for all logging agents."""
+import abc
+import os
+import shlex
+from typing import Any, Dict
+
+from sky.skylet import constants
+from sky.utils import common_utils
+from sky.utils import resources_utils
+
+
+class LoggingAgent(abc.ABC):
+    """Base class for all logging agents.
+
+    Each agent should implement the `get_setup_command` and
+    `get_credential_file_mounts` methods to return the setup command and
+    credential file mounts for the agent for provisioner to setup the agent
+    on each node.
+    """
+
+    @abc.abstractmethod
+    def get_setup_command(self,
+                          cluster_name: resources_utils.ClusterName) -> str:
+        pass
+
+    @abc.abstractmethod
+    def get_credential_file_mounts(self) -> Dict[str, str]:
+        pass
+
+
+class FluentbitAgent(LoggingAgent):
+    """Base class for logging store that use fluentbit as the agent."""
+
+    def get_setup_command(self,
+                          cluster_name: resources_utils.ClusterName) -> str:
+        install_cmd = (
+            'if ! command -v fluent-bit >/dev/null 2>&1; then '
+            'sudo apt-get install -y gnupg; '
+            # pylint: disable=line-too-long
+            'curl https://raw.githubusercontent.com/fluent/fluent-bit/master/install.sh | sh; '
+            'fi')
+        cfg = self.fluentbit_config(cluster_name)
+        cfg_path = os.path.join(constants.LOGGING_CONFIG_DIR, 'fluentbit.yaml')
+        config_cmd = (f'mkdir -p {constants.LOGGING_CONFIG_DIR} && '
+                      f'echo {shlex.quote(cfg)} > {cfg_path}')
+        start_cmd = ('nohup $(command -v fluent-bit || '
+                     'echo "/opt/fluent-bit/bin/fluent-bit") '
+                     f'-c {cfg_path} > /tmp/fluentbit.log 2>&1 &')
+        return f'set -e; {install_cmd}; {config_cmd}; {start_cmd}'
+
+    def fluentbit_config(self,
+                         cluster_name: resources_utils.ClusterName) -> str:
+        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,
+                }],
+                'outputs': [self.fluentbit_output_config(cluster_name)],
+            }
+        }
+        return common_utils.dump_yaml_str(cfg_dict)
+
+    @abc.abstractmethod
+    def fluentbit_output_config(
+            self, cluster_name: resources_utils.ClusterName) -> Dict[str, Any]:
+        pass

sky/logs/gcp.py

@@ -0,0 +1,91 @@
+"""GCP logging agent."""
+
+from typing import Any, Dict, Optional
+
+import pydantic
+
+from sky.clouds import gcp
+from sky.logs.agent import FluentbitAgent
+from sky.utils import resources_utils
+
+
+class _GCPLoggingConfig(pydantic.BaseModel):
+    """Configuration for GCP logging agent."""
+    project_id: Optional[str] = None
+    credentials_file: Optional[str] = None
+    additional_labels: Optional[Dict[str, str]] = None
+
+
+class _StackdriverOutputConfig(pydantic.BaseModel):
+    """Auxiliary model for building stackdriver output config in YAML.
+
+    Ref: https://docs.fluentbit.io/manual/1.7/pipeline/outputs/stackdriver
+    """
+    name: str = 'stackdriver'
+    match: str = '*'
+    export_to_project_id: Optional[str] = None
+    labels: Optional[Dict[str, str]] = None
+
+    def to_dict(self) -> Dict[str, Any]:
+        config = self.model_dump(exclude_none=True)
+        if self.labels:
+            # Replace the label format from `{k: v}` to `k=v`
+            label_str = ','.join([f'{k}={v}' for k, v in self.labels.items()])
+            config['labels'] = label_str
+        return config
+
+
+class GCPLoggingAgent(FluentbitAgent):
+    """GCP logging agent."""
+
+    def __init__(self, config: Dict[str, Any]):
+        self.config = _GCPLoggingConfig(**config)
+
+    def get_setup_command(self,
+                          cluster_name: resources_utils.ClusterName) -> str:
+        credential_path = gcp.DEFAULT_GCP_APPLICATION_CREDENTIAL_PATH
+        if self.config.credentials_file:
+            credential_path = self.config.credentials_file
+        # Set GOOGLE_APPLICATION_CREDENTIALS and check whether credentials
+        # is valid.
+        # Stackdriver only support service account credentials or credentials
+        # from metadata server (only available on GCE or GKE). If the default
+        # credentials uploaded by API server is NOT a service account key and
+        # there is NO metadata server available, the logging agent will fail to
+        # authenticate and we require the user to upload a service account key
+        # via logs.gcp.credentials_file in this case.
+        # Also note that we use env var instead of YAML config to specify the
+        # service account key file path in order to resolve the home directory
+        # more reliably.
+        # Ref: https://github.com/fluent/fluent-bit/issues/8804
+        # TODO(aylei): check whether the credentials config is valid before
+        # provision.
+        pre_cmd = (f'export GOOGLE_APPLICATION_CREDENTIALS={credential_path}; '
+                   f'cat {credential_path} | grep "service_account" || '
+                   f'(echo "Credentials file {credential_path} is not a '
+                   'service account key, check metadata server" && '
+                   'curl -s http://metadata.google.internal >/dev/null || '
+                   f'(echo "Neither service account key nor metadata server is '
+                   'available. Set logs.gcp.credentials_file to a service '
+                   'account key in server config and retry." && '
+                   'exit 1;))')
+        return pre_cmd + ' && ' + super().get_setup_command(cluster_name)
+
+    def fluentbit_output_config(
+            self, cluster_name: resources_utils.ClusterName) -> Dict[str, Any]:
+        display_name = cluster_name.display_name
+        unique_name = cluster_name.name_on_cloud
+
+        return _StackdriverOutputConfig(
+            export_to_project_id=self.config.project_id,
+            labels={
+                'skypilot_cluster_name': display_name,
+                'skypilot_cluster_id': unique_name,
+                **(self.config.additional_labels or {})
+            },
+        ).to_dict()
+
+    def get_credential_file_mounts(self) -> Dict[str, str]:
+        if self.config.credentials_file:
+            return {self.config.credentials_file: self.config.credentials_file}
+        return {}

sky/models.py

@@ -17,6 +17,7 @@ class User:
     id: str
     # Display name of the user
     name: Optional[str] = None
+    password: Optional[str] = None
 
     def to_dict(self) -> Dict[str, Any]:
         return {'id': self.id, 'name': self.name}

sky/provision/__init__.py

@@ -24,6 +24,7 @@ from sky.provision import lambda_cloud
 from sky.provision import nebius
 from sky.provision import oci
 from sky.provision import runpod
+from sky.provision import scp
 from sky.provision import ssh
 from sky.provision import vast
 from sky.provision import vsphere

sky/provision/instance_setup.py

@@ -8,6 +8,7 @@ import time
 from typing import Any, Callable, Dict, List, Optional, Tuple
 
 from sky import exceptions
+from sky import logs
 from sky import provision
 from sky import sky_logging
 from sky.provision import common
@@ -21,6 +22,7 @@ from sky.utils import accelerator_registry
 from sky.utils import command_runner
 from sky.utils import common_utils
 from sky.utils import env_options
+from sky.utils import resources_utils
 from sky.utils import subprocess_utils
 from sky.utils import timeline
 from sky.utils import ux_utils
@@ -557,3 +559,36 @@ def internal_file_mounts(cluster_name: str, common_file_mounts: Dict[str, str],
         ssh_credentials=ssh_credentials,
         max_workers=subprocess_utils.get_max_workers_for_file_mounts(
             common_file_mounts, cluster_info.provider_name))
+
+
+@common.log_function_start_end
+@timeline.event
+def setup_logging_on_cluster(logging_agent: logs.LoggingAgent,
+                             cluster_name: resources_utils.ClusterName,
+                             cluster_info: common.ClusterInfo,
+                             ssh_credentials: Dict[str, Any]) -> None:
+    """Setup logging agent (fluentbit) on all nodes after provisioning."""
+    _hint_worker_log_path(cluster_name.name_on_cloud, cluster_info,
+                          'logging_setup')
+
+    @_auto_retry()
+    def _setup_node(runner: command_runner.CommandRunner, log_path: str):
+        cmd = logging_agent.get_setup_command(cluster_name)
+        logger.info(f'Running command on node: {cmd}')
+        returncode, stdout, stderr = runner.run(cmd,
+                                                stream_logs=False,
+                                                require_outputs=True,
+                                                log_path=log_path,
+                                                source_bashrc=True)
+        if returncode:
+            raise RuntimeError(f'Failed to setup logging agent\n{cmd}\n'
+                               f'(exit code {returncode}). Error: '
+                               f'===== stdout ===== \n{stdout}\n'
+                               f'===== stderr ====={stderr}')
+
+    _parallel_ssh_with_cache(_setup_node,
+                             cluster_name.name_on_cloud,
+                             stage_name='logging_setup',
+                             digest=None,
+                             cluster_info=cluster_info,
+                             ssh_credentials=ssh_credentials)

sky/provision/provisioner.py

@@ -16,6 +16,7 @@ import sky
 from sky import clouds
 from sky import exceptions
 from sky import global_user_state
+from sky import logs
 from sky import provision
 from sky import sky_logging
 from sky import skypilot_config
@@ -648,6 +649,15 @@ def _post_provision_setup(
             logger.debug('Ray cluster is ready. Skip starting ray cluster on '
                          'worker nodes.')
 
+        logging_agent = logs.get_logging_agent()
+        if logging_agent:
+            status.update(
+                ux_utils.spinner_message('Setting up logging agent',
+                                         provision_logging.config.log_path))
+            instance_setup.setup_logging_on_cluster(logging_agent, cluster_name,
+                                                    cluster_info,
+                                                    ssh_credentials)
+
         instance_setup.start_skylet_on_head_node(cluster_name.name_on_cloud,
                                                  cluster_info, ssh_credentials)
 
@@ -672,6 +682,7 @@ def post_provision_runtime_setup(
        and other necessary files to the VM.
     3. Run setup commands to install dependencies.
     4. Start ray cluster and skylet.
+    5. (Optional) Setup logging agent.
 
     Raises:
         RuntimeError: If the setup process encounters any error.

sky/provision/scp/__init__.py

@@ -0,0 +1,15 @@
+"""SCP provisioner for SkyPilot."""
+
+from sky.provision.scp.config import bootstrap_instances
+from sky.provision.scp.instance import cleanup_ports
+from sky.provision.scp.instance import get_cluster_info
+from sky.provision.scp.instance import open_ports
+from sky.provision.scp.instance import query_instances
+from sky.provision.scp.instance import run_instances
+from sky.provision.scp.instance import stop_instances
+from sky.provision.scp.instance import terminate_instances
+from sky.provision.scp.instance import wait_instances
+
+__all__ = ('bootstrap_instances', 'cleanup_ports', 'get_cluster_info',
+           'open_ports', 'query_instances', 'run_instances', 'stop_instances',
+           'terminate_instances', 'wait_instances')

sky/provision/scp/config.py

@@ -0,0 +1,93 @@
+"""SCP configuration bootstrapping."""
+
+import subprocess
+
+from sky.clouds.utils import scp_utils
+from sky.provision import common
+
+
+def bootstrap_instances(
+        region: str, cluster_name: str,
+        config: common.ProvisionConfig) -> common.ProvisionConfig:
+    """Bootstraps instances for the given cluster."""
+    del cluster_name
+
+    node_cfg = config.node_config
+    zone_id = _get_zone_id(region)
+    node_cfg['zone_id'] = zone_id
+
+    docker_cfg = config.docker_config
+    docker_cfg['imageId'] = node_cfg['imageId']
+    docker_cfg['serviceZoneId'] = zone_id
+    docker_cfg['serverType'] = node_cfg['InstanceType']
+    docker_cfg['contractId'] = 'None'
+    ssh_public_key = node_cfg['AuthorizedKey']
+    docker_cfg['initialScript'] = _get_init_script(ssh_public_key)
+
+    key_pair_id = _get_key_pair_id()
+    miscellaneous = {
+        'deletionProtectionEnabled': False,
+        'keyPairId': key_pair_id,
+        'blockStorage': {
+            'blockStorageName': 'skystorage',
+            'diskSize': node_cfg['diskSize'],
+            'encryptEnabled': False,
+            'productId': 'PRODUCT-sRlJ34iBr9hOxN9J5PrQxo'
+        },
+        'nic': {
+            'natEnabled': True
+        },
+    }
+
+    docker_cfg.update(miscellaneous)
+
+    return config
+
+
+def _get_zone_id(region_name: str):
+    zone_contents = scp_utils.SCPClient().get_zones()
+    zone_dict = {
+        item['serviceZoneName']: item['serviceZoneId'] for item in zone_contents
+    }
+    return zone_dict[region_name]
+
+
+def _get_init_script(ssh_public_key: str):
+    init_script_content = _get_default_config_cmd() + _get_ssh_key_gen_cmd(
+        ssh_public_key)
+    init_script_content_string = f'"{init_script_content}"'
+    command = f'echo {init_script_content_string} | base64'
+    result = subprocess.run(command,
+                            shell=True,
+                            capture_output=True,
+                            text=True,
+                            check=True)
+    init_script_content_base64 = result.stdout
+    return {
+        'encodingType': 'base64',
+        'initialScriptShell': 'bash',
+        'initialScriptType': 'text',
+        'initialScriptContent': init_script_content_base64
+    }
+
+
+def _get_default_config_cmd():
+    cmd_list = ['apt-get update', 'apt-get -y install python3-pip']
+    res = ''
+    for cmd in cmd_list:
+        res += cmd + '; '
+    return res
+
+
+def _get_ssh_key_gen_cmd(ssh_public_key: str):
+    cmd_st = 'mkdir -p ~/.ssh/; touch ~/.ssh/authorized_keys;'
+    cmd_ed = 'chmod 644 ~/.ssh/authorized_keys; chmod 700 ~/.ssh/'
+    cmd = "echo '{}' &>>~/.ssh/authorized_keys;".format(ssh_public_key)  # pylint: disable=invalid-string-quote
+    return cmd_st + cmd + cmd_ed
+
+
+def _get_key_pair_id():
+    key_pairs = scp_utils.SCPClient().get_key_pairs()
+    if key_pairs['totalCount'] == 0:
+        raise RuntimeError('create key pair')
+    return key_pairs['contents'][0]['keyPairId']