From 14ac584577c1d7ca6181489adbdb61071f7d9678 Mon Sep 17 00:00:00 2001
From: allegroai <>
Date: Wed, 31 Mar 2021 23:53:58 +0300
Subject: [PATCH] Support k8s glue container env vars merging

---
 clearml_agent/glue/k8s.py     | 18 ++++++++++++++++--
 clearml_agent/helper/dicts.py | 12 +++++++++---
 2 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/clearml_agent/glue/k8s.py b/clearml_agent/glue/k8s.py
index e00834a..2e9fc33 100644
--- a/clearml_agent/glue/k8s.py
+++ b/clearml_agent/glue/k8s.py
@@ -433,7 +433,7 @@ class K8sIntegration(Worker):
                     script_encoded.encode('ascii')
                 ).decode('ascii'))
 
-        container = merge_dicts(
+        container = self._merge_containers(
             container,
             dict(name=name, image=docker_image,
                  command=['/bin/bash'],
@@ -441,7 +441,7 @@ class K8sIntegration(Worker):
         )
 
         if template['spec']['containers']:
-            template['spec']['containers'][0] = merge_dicts(template['spec']['containers'][0], container)
+            template['spec']['containers'][0] = self._merge_containers(template['spec']['containers'][0], container)
         else:
             template['spec']['containers'].append(container)
 
@@ -591,3 +591,17 @@ class K8sIntegration(Worker):
     @classmethod
     def get_ssh_server_bash(cls, ssh_port_number):
         return ' ; '.join(line.format(port=ssh_port_number) for line in cls.BASH_INSTALL_SSH_CMD)
+
+    @staticmethod
+    def _merge_containers(c1, c2):
+        def merge_env(k, d1, d2, not_set):
+            if k != "env":
+                return not_set
+            # Merge environment lists, second list overrides first
+            return list({
+                item['name']: item for envs in (d1, d2) for item in envs
+            }.values())
+
+        return merge_dicts(
+            c1, c2, custom_merge_func=merge_env
+        )
diff --git a/clearml_agent/helper/dicts.py b/clearml_agent/helper/dicts.py
index 649fb9d..4417b6e 100644
--- a/clearml_agent/helper/dicts.py
+++ b/clearml_agent/helper/dicts.py
@@ -1,17 +1,23 @@
-from typing import Callable, Dict, Any
+from typing import Callable, Dict, Any, Optional
+
+_not_set = object()
 
 
 def filter_keys(filter_, dct):  # type: (Callable[[Any], bool], Dict) -> Dict
     return {key: value for key, value in dct.items() if filter_(key)}
 
 
-def merge_dicts(dict1, dict2):
+def merge_dicts(dict1, dict2, custom_merge_func=None):
+    # type: (Any, Any, Optional[Callable[[str, Any, Any, Any], Any]]) -> Any
     """ Recursively merges dict2 into dict1 """
     if not isinstance(dict1, dict) or not isinstance(dict2, dict):
         return dict2
     for k in dict2:
         if k in dict1:
-            dict1[k] = merge_dicts(dict1[k], dict2[k])
+            res = None
+            if custom_merge_func:
+                res = custom_merge_func(k, dict1[k], dict2[k], _not_set)
+            dict1[k] = merge_dicts(dict1[k], dict2[k], custom_merge_func) if res is None else res
         else:
             dict1[k] = dict2[k]
     return dict1