Usage Collection & Runtimes (skypilot-org#872)

* Logging CLI & YAML

* Logging stack traces & added metrics

* yapf

* added privacy policy

* formatting

* make suggested changes + added setup documentation

* more edits

* address some comments

* Fix comments

* Fix user hash

* Rename

* refactor a bit

* fix return value not returned bug

* update instance and readme

* Fix transaction and cli logging

* Remove remnant

* options to turn off logging

* add SKY_DEV=1 before pytest

* format

* fix option

* push

* Fix redaction and add overrides

* fix typo

* remove remnant

* Add new lines

* fix

* refactor metrics and usage logging

* Disable logging for testing

* reformat

* Disable metrics

* make time stamp the same

* format

* Send json instead

* Add runtime

* Add runtime

* format

* Remove metric

* address comments

* update readme

* docstr

* restore timeline

* switch back the timeline.event

* Reorganize backend functions

* fix

* Fix smoke and spot status refresh

* order of __init__

* fix

* address comments

* Use cli for status showing

* Address comments

* Add runtime logging for backend

* format

* Disable logging for github

* Revert to two functions for file mounts

* Move comment lines out

* Add schema for the logs

* Add fields

* Add readme for setting loki

* format

* Add instructions for setting up loki service

* link to IAM role creation

* Address comments

* Add region

* format

* remove IAM key

* move update functions into the class

* fix set_new_cluster

* Catch unrecognized options and command

* Address comments

* Remove note

* Add development environment

* Switch to elastic IP

* update status

* Fix time and add readme

* Add callable handle for clean_yaml

* format

* Fix status

* Fix sky_dev, user for spot and env

* Use new space for release

* Add internal field

* increase sleep for spot recovery

* Fix logs for onprem

* fix instruction

* disable logging for smoke test

* Update docs and message reset

* format

* Address comments

* Fix readme

* Fix grafana screenshot

Co-authored-by: Zhanghao Wu <[email protected]>

Author: mraheja

.github/workflows/pytest.yml

@@ -51,4 +51,4 @@ jobs:
           pip install pytest
 
       - name: Run tests with pytest
-        run: pytest ${{ matrix.test-path }}
+        run: SKY_DISABLE_USAGE_COLLECTION=1 pytest ${{ matrix.test-path }}

README.md

@@ -21,13 +21,21 @@ Use editable mode (`-e`) when installing:
 pip install -e ".[all]"
 pip install -r requirements-dev.txt
 ```
+IMPORTANT: Please `export SKY_DEV=1` before running the sky commands in the terminal, so that the developing log will not pollute the actual user logs.
+
 
 ### Submitting pull requests
 - After you commit, format your code with [`format.sh`](./format.sh).
 - In the PR description, write a `Tested:` section to describe relevant tests performed.
 - For changes that touch the core system, run the [smoke tests](#testing) and ensure they pass.
 - Follow the [Google style guide](https://google.github.io/styleguide/pyguide.html).
 
+
+### Environment Variable Options
+- `export SKY_DEV=1` to show debugging logs (logging.DEBUG) and send the logs to dev space.
+- `export SKY_DISABLE_USAGE_COLLECTION=1` to disable usage logging.
+- `export SKY_MINIMIZE_LOGGING=1` to minimize the sky outputs for demo purpose.
+
 ### Dump timeline
 
 Timeline is useful for performance analysis and debugging in Sky.

docs/source/examples/spot-jobs.rst

@@ -9,7 +9,7 @@ This feature **saves significant cost** (e.g., up to 70\% for GPU VMs) by making
 To maximize availability, Sky automatically finds available spot resources across regions and clouds.
 Here is an example of BERT training job failing over different regions across AWS and GCP.
 
-.. image:: ../imgs/spot-training.png
+.. image:: ../images/spot-training.png
   :width: 600
   :alt: BERT training on Spot V100
 

docs/source/imgs/spot-training.png renamed to docs/source/images/spot-training.png

@@ -54,6 +54,7 @@ Key features:
    reference/storage
    reference/local/index
    reference/quota
+   reference/logging
    reference/faq
 
 .. toctree::

docs/source/index.rst

@@ -0,0 +1,31 @@
+.. _logging:
+
+Usage Collection
+=================
+
+Sky collects usage stats by default. This data will only be used by the Sky team to improve its services and for research purpose.
+We will **not** sell data or buy data about you.
+
+
+What data is collected?
+-----------------------
+
+We collect non-sensitive data that helps us understand how Sky is used. We will redact your ``setup``, ``run``, and ``env`` from the collected data.
+
+.. _usage-disable:
+
+How to disable it
+-----------------
+To disable usage collection, set the ``SKY_DISABLE_USAGE_COLLECTION`` environment variable by :code:`export SKY_DISABLE_USAGE_COLLECTION=1`.
+
+
+How does it work?
+-----------------
+
+When a Sky CLI or entrypoint function is called, Sky will do the following:
+
+#. Check the environment variable ``SKY_DISABLE_USAGE_COLLECTION`` is set: 1 means disabled and 0 means enabled.
+
+#. If the environment variable is not set or set to 0, it will collect information about the cluster and task resource requirements 
+
+#. If the environment variable is set to 1, it will skip any message sending.

docs/source/reference/logging.rst

@@ -2,7 +2,9 @@
 import typing
 from typing import Dict, Optional
 
+import sky
 from sky.utils import timeline
+from sky.usage import usage_lib
 
 if typing.TYPE_CHECKING:
     from sky import resources
@@ -32,21 +34,28 @@ def check_resources_fit_cluster(self, handle: ResourceHandle,
         raise NotImplementedError
 
     @timeline.event
+    @usage_lib.messages.usage.update_runtime('provision')
     def provision(self,
                   task: 'task_lib.Task',
                   to_provision: Optional['resources.Resources'],
                   dryrun: bool,
                   stream_logs: bool,
                   cluster_name: Optional[str] = None,
                   retry_until_up: bool = False) -> ResourceHandle:
+        if cluster_name is None:
+            cluster_name = sky.backends.backend_utils.generate_cluster_name()
+        usage_lib.messages.usage.update_cluster_name(cluster_name)
+        usage_lib.messages.usage.update_actual_task(task)
         return self._provision(task, to_provision, dryrun, stream_logs,
                                cluster_name, retry_until_up)
 
     @timeline.event
+    @usage_lib.messages.usage.update_runtime('sync_workdir')
     def sync_workdir(self, handle: ResourceHandle, workdir: Path) -> None:
         return self._sync_workdir(handle, workdir)
 
     @timeline.event
+    @usage_lib.messages.usage.update_runtime('sync_file_mounts')
     def sync_file_mounts(
         self,
         handle: ResourceHandle,
@@ -56,15 +65,19 @@ def sync_file_mounts(
         return self._sync_file_mounts(handle, all_file_mounts, storage_mounts)
 
     @timeline.event
+    @usage_lib.messages.usage.update_runtime('setup')
     def setup(self, handle: ResourceHandle, task: 'task_lib.Task') -> None:
         return self._setup(handle, task)
 
     def add_storage_objects(self, task: 'task_lib.Task') -> None:
         raise NotImplementedError
 
     @timeline.event
+    @usage_lib.messages.usage.update_runtime('execute')
     def execute(self, handle: ResourceHandle, task: 'task_lib.Task',
                 detach_run: bool) -> None:
+        usage_lib.messages.usage.update_cluster_name(handle.get_cluster_name())
+        usage_lib.messages.usage.update_actual_task(task)
         return self._execute(handle, task, detach_run)
 
     @timeline.event
@@ -77,6 +90,7 @@ def teardown_ephemeral_storage(self, task: 'task_lib.Task') -> None:
         return self._teardown_ephemeral_storage(task)
 
     @timeline.event
+    @usage_lib.messages.usage.update_runtime('teardown')
     def teardown(self,
                  handle: ResourceHandle,
                  terminate: bool,
@@ -93,7 +107,7 @@ def _provision(self,
                    to_provision: Optional['resources.Resources'],
                    dryrun: bool,
                    stream_logs: bool,
-                   cluster_name: Optional[str] = None,
+                   cluster_name: str,
                    retry_until_up: bool = False) -> ResourceHandle:
         raise NotImplementedError
 

sky/backends/backend.py

@@ -5,13 +5,11 @@
 import difflib
 import enum
 import getpass
-import hashlib
 import json
 import os
 import pathlib
 import random
 import re
-import socket
 import subprocess
 import textwrap
 import threading
@@ -45,11 +43,13 @@
 from sky import spot as spot_lib
 from sky.backends import onprem_utils
 from sky.skylet import log_lib
+from sky.utils import common_utils
 from sky.utils import command_runner
 from sky.utils import subprocess_utils
 from sky.utils import timeline
 from sky.utils import ux_utils
 from sky.utils import validator
+from sky.usage import usage_lib
 
 if typing.TYPE_CHECKING:
     from sky import resources
@@ -613,7 +613,7 @@ def write_cluster_config(to_provision: 'resources.Resources',
                 # (username, last 4 chars of hash of hostname): for uniquefying
                 # users on shared-account cloud providers. Using uuid.getnode()
                 # is incorrect; observed to collide on Macs.
-                'security_group': f'sky-sg-{user_and_hostname_hash()}',
+                'security_group': f'sky-sg-{common_utils.user_and_hostname_hash()}',
                 # Azure only.
                 'azure_subscription_id': azure_subscription_id,
                 'resource_group': f'{cluster_name}-{region_name}',
@@ -640,6 +640,7 @@ def write_cluster_config(to_provision: 'resources.Resources',
     if dryrun:
         return config_dict
     _add_auth_to_cluster_config(cloud, yaml_path)
+    usage_lib.messages.usage.update_ray_yaml(yaml_path)
     # For TPU nodes. TPU VMs do not need TPU_NAME.
     if (resources_vars.get('tpu_type') is not None and
             resources_vars.get('tpu_vm') is None):
@@ -689,30 +690,7 @@ def _add_auth_to_cluster_config(cloud: clouds.Cloud, cluster_config_file: str):
         # in the local cluster config (in ~/.sky/local/...). There is no need
         # for Sky to generate authentication.
         pass
-    dump_yaml(cluster_config_file, config)
-
-
-def read_yaml(path):
-    with open(path, 'r') as f:
-        config = yaml.safe_load(f)
-    return config
-
-
-def dump_yaml(path, config):
-    # https://github.com/yaml/pyyaml/issues/127
-    class LineBreakDumper(yaml.SafeDumper):
-
-        def write_line_break(self, data=None):
-            super().write_line_break(data)
-            if len(self.indents) == 1:
-                super().write_line_break()
-
-    with open(path, 'w') as f:
-        yaml.dump(config,
-                  f,
-                  Dumper=LineBreakDumper,
-                  sort_keys=False,
-                  default_flow_style=False)
+    common_utils.dump_yaml(cluster_config_file, config)
 
 
 def get_run_timestamp() -> str:
@@ -836,7 +814,7 @@ def wait_until_ray_cluster_ready(
 
 def ssh_credential_from_yaml(cluster_yaml: str) -> Tuple[str, str, str]:
     """Returns ssh_user, ssh_private_key and ssh_control name."""
-    config = read_yaml(cluster_yaml)
+    config = common_utils.read_yaml(cluster_yaml)
     auth_section = config['auth']
     ssh_user = auth_section['ssh_user'].strip()
     ssh_private_key = auth_section.get('ssh_private_key')
@@ -934,35 +912,6 @@ def check_local_gpus() -> bool:
     return is_functional
 
 
-def user_and_hostname_hash() -> str:
-    """Returns a string containing <user>-<hostname hash last 4 chars>.
-
-    For uniquefying user clusters on shared-account cloud providers. Also used
-    for AWS security group.
-
-    Using uuid.getnode() instead of gethostname() is incorrect; observed to
-    collide on Macs.
-
-    NOTE: BACKWARD INCOMPATIBILITY NOTES
-
-    Changing this string will render AWS clusters shown in `sky status`
-    unreusable and potentially cause leakage:
-
-    - If a cluster is STOPPED, any command restarting it (`sky launch`, `sky
-      start`) will launch a NEW cluster.
-    - If a cluster is UP, a `sky launch` command reusing it will launch a NEW
-      cluster. The original cluster will be stopped and thus leaked from Sky's
-      perspective.
-    - `sky down/stop/exec` on these pre-change clusters still works, if no new
-      clusters with the same name have been launched.
-
-    The reason is AWS security group names are derived from this string, and
-    thus changing the SG name makes these clusters unrecognizable.
-    """
-    hostname_hash = hashlib.md5(socket.gethostname().encode()).hexdigest()[-4:]
-    return f'{getpass.getuser()}-{hostname_hash}'
-
-
 def generate_cluster_name():
     # TODO: change this ID formatting to something more pleasant.
     # User name is helpful in non-isolated accounts, e.g., GCP, Azure.
@@ -1309,7 +1258,7 @@ def _get_cluster_status_via_cloud_cli(
     """Returns the status of the cluster."""
     resources: sky.Resources = handle.launched_resources
     cloud = resources.cloud
-    ray_config = read_yaml(handle.cluster_yaml)
+    ray_config = common_utils.read_yaml(handle.cluster_yaml)
     return _QUERY_STATUS_FUNCS[str(cloud)](handle.cluster_name, ray_config)
 
 
@@ -1407,7 +1356,7 @@ def _update_cluster_status(
 
     The function will update the cached cluster status in the global state. For the
     design of the cluster status and transition, please refer to the
-    sky/design_docs/cluster_states.md
+    sky/design_docs/cluster_status.md
 
     Returns:
       If the cluster is terminated or does not exist, return None.