From 980b7ac706efa89ba619f42b0b77f57bf90b96c6 Mon Sep 17 00:00:00 2001
From: Andrew Aikawa <asai@berkeley.edu>
Date: Tue, 17 Sep 2024 16:20:03 -0700
Subject: [PATCH] [usage] log entrypoints to posthog (#4)

* add posthog events"

* lint
---
 sky/setup_files/setup.py | 1 +
 sky/usage/constants.py   | 7 +++++++
 sky/usage/usage_lib.py   | 8 ++++++++
 3 files changed, 16 insertions(+)

diff --git a/sky/setup_files/setup.py b/sky/setup_files/setup.py
index 604060c68ae..c6634bc2599 100644
--- a/sky/setup_files/setup.py
+++ b/sky/setup_files/setup.py
@@ -156,6 +156,7 @@ def parse_readme(readme: str) -> str:
     "typing_extensions",
     'filelock >= 3.6.0',
     'packaging',
+    'posthog',
     'psutil',
     'pulp',
     # Cython 3.0 release breaks PyYAML 5.4.* (https://github.com/yaml/pyyaml/issues/601)
diff --git a/sky/usage/constants.py b/sky/usage/constants.py
index c2f8c6d067b..1656f627ae4 100644
--- a/sky/usage/constants.py
+++ b/sky/usage/constants.py
@@ -1,11 +1,18 @@
 """Constants for usage collection."""
 
+import os
+
 LOG_URL = 'http://usage.skypilot.co:9090/loki/api/v1/push'  # pylint: disable=line-too-long
 
 USAGE_MESSAGE_SCHEMA_VERSION = 1
 
 PRIVACY_POLICY_PATH = '~/.sky/privacy_policy'
 
+POSTHOG_API_KEY = os.environ.get(
+    'POSTHOG_API_KEY', 'phc_4UgX80BfVNmYRZ2o3dJLyRMGkv1CxBozPAcPnD29uP4')
+
+POSTHOG_HOST = os.environ.get('POSTHOG_HOST', 'https://us.i.posthog.com')
+
 USAGE_POLICY_MESSAGE = (
     'SkyPilot collects usage data to improve its services. '
     '`setup` and `run` commands are not collected to '
diff --git a/sky/usage/usage_lib.py b/sky/usage/usage_lib.py
index a6c10da5c7a..a30fb02d0b9 100644
--- a/sky/usage/usage_lib.py
+++ b/sky/usage/usage_lib.py
@@ -12,6 +12,7 @@
 from typing import Any, Callable, Dict, List, Optional, Union
 
 import click
+import posthog
 import requests
 
 import sky
@@ -27,6 +28,8 @@
     from sky import task as task_lib
 
 logger = sky_logging.init_logger(__name__)
+posthog_client = posthog.Posthog(constants.POSTHOG_API_KEY,
+                                 constants.POSTHOG_HOST)
 
 
 def _get_current_timestamp_ns() -> int:
@@ -320,6 +323,11 @@ def _send_to_loki(message_type: MessageType):
                              data=payload,
                              headers=headers,
                              timeout=0.5)
+    posthog_client.capture(common_utils.get_user_hash(), 'trainy-sky-event', {
+        **prom_labels,
+        **json.loads(str(message)), 'timestamp': str(log_timestamp)
+    })
+
     if response.status_code != 204:
         logger.debug(
             f'Grafana Loki failed with response: {response.text}\n{payload}')