From c77224af68e1d1b66adc888ba0ee5ca2d7b1b8d3 Mon Sep 17 00:00:00 2001
From: allegroai <>
Date: Mon, 3 Jul 2023 11:07:12 +0300
Subject: [PATCH] Add support for task field injection into container docker
 name

---
 .../backend_api/config/default/agent.conf     |  5 +++
 clearml_agent/commands/worker.py              | 43 ++++++++++++++++++-
 2 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/clearml_agent/backend_api/config/default/agent.conf b/clearml_agent/backend_api/config/default/agent.conf
index 3eb4123..82f5b88 100644
--- a/clearml_agent/backend_api/config/default/agent.conf
+++ b/clearml_agent/backend_api/config/default/agent.conf
@@ -269,10 +269,15 @@
 
     # Name docker containers created by the daemon using the following string format (supported from Docker 0.6.5)
     # Allowed variables are task_id, worker_id and rand_string (random lower-case letters string, up to 32 characters)
+    # Custom variables may be specified using the docker_container_name_format_fields option.
     # Note: resulting name must start with an alphanumeric character and
     #       continue with alphanumeric characters, underscores (_), dots (.) and/or dashes (-)
     # docker_container_name_format: "clearml-id-{task_id}-{rand_string:.8}"
 
+    # Specify custom variables for the docker_container_name_format option using a mapping of variable name
+    # to a (nested) task field (using "." as a task field separator, digits specify array index)
+    # docker_container_name_format_fields: { foo: "bar.moo" }
+
     # Apply top-level environment section from configuration into os.environ
     apply_environment: true
     # Top-level environment section is in the form of:
diff --git a/clearml_agent/commands/worker.py b/clearml_agent/commands/worker.py
index d17838a..973bbc6 100644
--- a/clearml_agent/commands/worker.py
+++ b/clearml_agent/commands/worker.py
@@ -319,6 +319,37 @@ def get_next_task(session, queue, get_task_info=False):
     return data
 
 
+def get_task_fields(session, task_id, fields: list, log=None) -> dict:
+    """
+    Returns dict with Task docker container setup {container: '', arguments: '', setup_shell_script: ''}
+    """
+    result = session.send_request(
+        service='tasks',
+        action='get_all',
+        json={'id': [task_id], 'only_fields': list(fields), 'search_hidden': True},
+        method=Request.def_method,
+        async_enable=False,
+    )
+    # noinspection PyBroadException
+    try:
+        results = {}
+        result = result.json()['data']['tasks'][0]
+        for field in fields:
+            cur = result
+            for part in field.split("."):
+                if part.isdigit():
+                    cur = cur[part]
+                else:
+                    cur = cur.get(part, {})
+            results[field] = cur
+        return results
+    except Exception as ex:
+        if log:
+            log.error("Failed obtaining values for task fields {}: {}", fields, ex)
+        pass
+    return {}
+
+
 def get_task_container(session, task_id):
     """
     Returns dict with Task docker container setup {container: '', arguments: '', setup_shell_script: ''}
@@ -890,11 +921,21 @@ class Worker(ServiceCommandSection):
 
             name_format = self._session.config.get('agent.docker_container_name_format', None)
             if name_format:
+                custom_fields = {}
+                name_format_fields = self._session.config.get('agent.docker_container_name_format_fields', None)
+                if name_format_fields:
+                    field_values = get_task_fields(task_session, task_id, name_format_fields.values(), log=self.log)
+                    custom_fields = {
+                        k: field_values.get(v)
+                        for k, v in name_format_fields.items()
+                    }
+
                 try:
                     name = name_format.format(
                         task_id=re.sub(r'[^a-zA-Z0-9._-]', '-', task_id),
                         worker_id=re.sub(r'[^a-zA-Z0-9._-]', '-', worker_id),
-                        rand_string="".join(sys_random.choice(string.ascii_lowercase) for _ in range(32))
+                        rand_string="".join(sys_random.choice(string.ascii_lowercase) for _ in range(32)),
+                        **custom_fields,
                     )
                 except Exception as ex:
                     print("Warning: failed generating docker container name: {}".format(ex))