Merge pull request #290 from runpod/user-agent-tracking

User agent tracking

Author: justinmerrell

CHANGELOG.md

@@ -1,5 +1,11 @@
 # Change Log
 
+## Release 1.6.1 (TBD)
+
+### Added
+
+- User-Agent for better analytics tracking.
+
 ## Release 1.6.0 (1/29/24)
 
 ### Fixed

runpod/api/graphql.py

@@ -8,18 +8,23 @@
 import requests
 
 from runpod import error
+from runpod.user_agent import USER_AGENT
 
 HTTP_STATUS_UNAUTHORIZED = 401
 
+
 def run_graphql_query(query: str) -> Dict[str, Any]:
     '''
     Run a GraphQL query
     '''
     from runpod import api_key  # pylint: disable=import-outside-toplevel, cyclic-import
     url = f"https://api.runpod.io/graphql?api_key={api_key}"
+
     headers = {
         "Content-Type": "application/json",
+        "User-Agent": USER_AGENT,
     }
+
     data = json.dumps({"query": query})
     response = requests.post(url, headers=headers, data=data, timeout=30)
 

runpod/serverless/modules/rp_ping.py

@@ -9,12 +9,13 @@
 import requests
 from urllib3.util.retry import Retry
 
+from runpod.version import __version__ as runpod_version
 from runpod.serverless.modules.rp_logger import RunPodLogger
-from .worker_state import Jobs, WORKER_ID
-from ...version import __version__ as runpod_version
+from runpod.serverless.modules.worker_state import Jobs, WORKER_ID
+
 
 log = RunPodLogger()
-jobs = Jobs() # Contains the list of jobs that are currently running.
+jobs = Jobs()  # Contains the list of jobs that are currently running.
 
 
 class Heartbeat:

runpod/serverless/worker.py

@@ -8,6 +8,7 @@
 
 import aiohttp
 
+from runpod.user_agent import USER_AGENT
 from runpod.serverless.modules import (
     rp_logger, rp_local, rp_handler, rp_ping,
     rp_scale
@@ -24,7 +25,10 @@
 
 def _get_auth_header() -> Dict[str, str]:
     """ Returns the authorization header with the API key. """
-    return {"Authorization": f"{os.environ.get('RUNPOD_AI_API_KEY')}"}
+    return {
+        "Authorization": f"{os.environ.get('RUNPOD_AI_API_KEY')}",
+        "User-Agent": USER_AGENT
+    }
 
 
 def _is_local(config) -> bool:

runpod/user_agent.py

@@ -0,0 +1,32 @@
+""" User-Agent for RunPod-Python-SDK """
+
+import os
+import platform
+
+from runpod.version import __version__ as runpod_version
+
+
+def construct_user_agent():
+    """ Constructs the User-Agent string for the RunPod-Python-SDK
+
+    Example:
+        RunPod-Python-SDK/0.1.0 (Linux 5.4.0-54-generic; x86_64) Language/Python 3.8.5
+    """
+    os_info = f"{platform.system()} {platform.release()}; {platform.machine()}"
+    python_version = platform.python_version()
+    integration_method = os.getenv('RUNPOD_UA_INTEGRATION')
+
+    ua_components = [
+        f"RunPod-Python-SDK/{runpod_version}",
+        f"({os_info})",
+        f"Language/Python {python_version}"
+    ]
+
+    if integration_method:
+        ua_components.append(f"Integration/{integration_method}")
+
+    user_agent = " ".join(ua_components)
+    return user_agent
+
+
+USER_AGENT = construct_user_agent()

tests/test_serverless/test_worker.py

@@ -2,13 +2,15 @@
 # pylint: disable=protected-access
 
 import os
+import platform
 import argparse
 from unittest.mock import patch, mock_open, Mock, MagicMock
 
 from unittest import IsolatedAsyncioTestCase
 import nest_asyncio
 
 import runpod
+from runpod._version import __version__ as runpod_version
 from runpod.serverless.modules.rp_logger import RunPodLogger
 from runpod.serverless import _signal_handler
 
@@ -34,10 +36,14 @@ def test_get_auth_header(self):
         '''
         Test _get_auth_header
         '''
+        os_info = f"{platform.system()} {platform.release()}; {platform.machine()}"
         with patch("runpod.serverless.worker.os") as mock_os:
             mock_os.environ.get.return_value = "test"
             assert runpod.serverless.worker._get_auth_header(
-            ) == {'Authorization': 'test'}
+            ) == {
+                'Authorization': 'test',
+                'User-Agent': f'RunPod-Python-SDK/{runpod_version} ({os_info}) Language/Python {platform.python_version()}'  # pylint: disable=line-too-long
+            }
 
     def test_is_local(self):
         '''

tests/test_user_agent.py

@@ -0,0 +1,50 @@
+""" Tests for the user_agent module. """
+
+import unittest
+from unittest.mock import patch
+import os
+
+from runpod import __version__ as runpod_version
+from runpod.user_agent import construct_user_agent
+
+
+class TestConstructUserAgent(unittest.TestCase):
+    """Test the construct_user_agent function."""
+
+    @patch('runpod.user_agent.platform.system', return_value='Windows')
+    @patch('runpod.user_agent.platform.release', return_value='10')
+    @patch('runpod.user_agent.platform.machine', return_value='AMD64')
+    @patch('runpod.user_agent.platform.python_version', return_value='3.8.10')
+    def test_user_agent_without_integration(
+            self, mock_python_version, mock_machine, mock_release, mock_system):
+        """Test the User-Agent string without specifying an integration method."""
+        if 'RUNPOD_UA_INTEGRATION' in os.environ:
+            del os.environ['RUNPOD_UA_INTEGRATION']
+
+        expected_ua = f"RunPod-Python-SDK/{runpod_version} (Windows 10; AMD64) Language/Python 3.8.10"  # pylint: disable=line-too-long
+        self.assertEqual(construct_user_agent(), expected_ua)
+
+        assert mock_python_version.called
+        assert mock_machine.called
+        assert mock_release.called
+        assert mock_system.called
+
+    @patch('runpod.user_agent.platform.system', return_value='Linux')
+    @patch('runpod.user_agent.platform.release', return_value='5.4')
+    @patch('runpod.user_agent.platform.machine', return_value='x86_64')
+    @patch('runpod.user_agent.platform.python_version', return_value='3.9.5')
+    @patch.dict(os.environ, {'RUNPOD_UA_INTEGRATION': 'SkyPilot'})
+    def test_user_agent_with_integration(
+            self, mock_python_version, mock_machine, mock_release, mock_system):
+        """Test the User-Agent string with an integration method specified."""
+        expected_ua = f"RunPod-Python-SDK/{runpod_version} (Linux 5.4; x86_64) Language/Python 3.9.5 Integration/SkyPilot"  # pylint: disable=line-too-long
+        self.assertEqual(construct_user_agent(), expected_ua)
+
+        assert mock_python_version.called
+        assert mock_machine.called
+        assert mock_release.called
+        assert mock_system.called
+
+
+if __name__ == '__main__':
+    unittest.main()