From a69766bd8bcbcb527513cbd857847001a0631059 Mon Sep 17 00:00:00 2001
From: allegroai <>
Date: Thu, 15 Sep 2022 20:16:01 +0300
Subject: [PATCH] Add CLEARML_AGENT_CHILD_AGENTS_COUNT_CMD to allow overriding
 child agent count command in k8s

---
 clearml_agent/commands/worker.py | 17 ++++++++---------
 clearml_agent/definitions.py     |  1 +
 2 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/clearml_agent/commands/worker.py b/clearml_agent/commands/worker.py
index d79a531..aa1c800 100644
--- a/clearml_agent/commands/worker.py
+++ b/clearml_agent/commands/worker.py
@@ -71,6 +71,7 @@ from clearml_agent.definitions import (
     ENV_AGENT_SKIP_PYTHON_ENV_INSTALL,
     WORKING_STANDALONE_DIR,
     ENV_DEBUG_INFO,
+    ENV_CHILD_AGENTS_COUNT_CMD,
 )
 from clearml_agent.definitions import WORKING_REPOSITORY_DIR, PIP_EXTRA_INDICES
 from clearml_agent.errors import (
@@ -3589,15 +3590,13 @@ class Worker(ServiceCommandSection):
     def _get_child_agents_count_for_worker(self):
         """Get the amount of running child agents. In case of any error return 0"""
         parent_worker_label = self._parent_worker_label.format(self.worker_id)
-        cmd = [
-            'docker',
-            'ps',
-            '--filter',
-            'label={}'.format(parent_worker_label),
-            '--format',
-            # get some fields for debugging
-            '{"ID":"{{ .ID }}", "Image": "{{ .Image }}", "Names":"{{ .Names }}", "Labels":"{{ .Labels }}"}'
-        ]
+
+        default_cmd = 'docker ps --filter label={parent_worker_label} --format ' \
+                      '{{"ID":"{{{{ .ID }}}}", "Image": "{{{{ .Image }}}}", ' \
+                      '"Names":"{{{{ .Names }}}}", "Labels":"{{{{ .Labels }}}}"}}'
+        child_agents_cmd = ENV_CHILD_AGENTS_COUNT_CMD.get() or default_cmd
+
+        cmd = shlex.split(child_agents_cmd.format(parent_worker_label=parent_worker_label))
         try:
             output = Argv(*cmd).get_output(
                 stderr=subprocess.STDOUT
diff --git a/clearml_agent/definitions.py b/clearml_agent/definitions.py
index e8dceee..d23b53a 100644
--- a/clearml_agent/definitions.py
+++ b/clearml_agent/definitions.py
@@ -149,6 +149,7 @@ ENV_DOCKER_HOST_MOUNT = EnvironmentConfig('CLEARML_AGENT_K8S_HOST_MOUNT', 'CLEAR
 ENV_VENV_CACHE_PATH = EnvironmentConfig('CLEARML_AGENT_VENV_CACHE_PATH')
 ENV_EXTRA_DOCKER_ARGS = EnvironmentConfig('CLEARML_AGENT_EXTRA_DOCKER_ARGS', type=list)
 ENV_DEBUG_INFO = EnvironmentConfig('CLEARML_AGENT_DEBUG_INFO')
+ENV_CHILD_AGENTS_COUNT_CMD = EnvironmentConfig('CLEARML_AGENT_CHILD_AGENTS_COUNT_CMD')
 
 ENV_CUSTOM_BUILD_SCRIPT = EnvironmentConfig('CLEARML_AGENT_CUSTOM_BUILD_SCRIPT')
 """