add ua for k8s client

BalaBalaYi · Jun 7, 2024 · c90eb05 · c90eb05
1 parent aff7c11
commit c90eb05
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/dlrover/python/scheduler/kubernetes.py b/dlrover/python/scheduler/kubernetes.py
@@ -40,6 +40,7 @@
 }
 
 JOB_SUFFIX = "-edljob-"
+USER_AGENT = "easyDL/29.0.0/python"
 
 
 def convert_memory_to_mb(memory: str):
@@ -148,6 +149,7 @@ def __init__(self, namespace):
  self.client = client.CoreV1Api()
  self.api_instance = client.CustomObjectsApi()
  self.api_client = client.ApiClient()
+ self.api_client.user_agent = USER_AGENT
  self._namespace = namespace
 
  @retry_k8s_request

diff --git a/dlrover/python/tests/test_k8s_utils.py b/dlrover/python/tests/test_k8s_utils.py
@@ -17,6 +17,7 @@
 
 from dlrover.python.common.node import NodeResource
 from dlrover.python.scheduler.kubernetes import (
+ USER_AGENT,
  get_main_container,
  k8sClient,
  k8sServiceFactory,
@@ -90,3 +91,4 @@ def test_client(self):
  client = k8sClient.singleton_instance("default")
  succeed = client.cordon_node("test-node-0")
  self.assertFalse(succeed)
+ self.assertEqual(client.api_client.user_agent, USER_AGENT)