From 6470b16b70165cdab0159ce5613521917e861d79 Mon Sep 17 00:00:00 2001
From: allegroai <>
Date: Thu, 15 Jul 2021 10:26:09 +0300
Subject: [PATCH] Add k8s set task container if using default image/arguments

---
 clearml_agent/glue/k8s.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/clearml_agent/glue/k8s.py b/clearml_agent/glue/k8s.py
index 24d65d9..4abccce 100644
--- a/clearml_agent/glue/k8s.py
+++ b/clearml_agent/glue/k8s.py
@@ -17,7 +17,7 @@ from typing import Text, List, Callable, Any, Collection, Optional, Union
 import yaml
 
 from clearml_agent.commands.events import Events
-from clearml_agent.commands.worker import Worker, get_task_container
+from clearml_agent.commands.worker import Worker, get_task_container, set_task_container
 from clearml_agent.definitions import ENV_DOCKER_IMAGE
 from clearml_agent.errors import APIError
 from clearml_agent.helper.base import safe_remove_file
@@ -281,6 +281,9 @@ class K8sIntegration(Worker):
                 ENV_DOCKER_IMAGE.get() or self._session.config.get("agent.default_docker.image", "nvidia/cuda")
             )
             container['arguments'] = self._session.config.get("agent.default_docker.arguments", None)
+            set_task_container(
+                self._session, task_id, docker_image=container['image'], docker_arguments=container['arguments']
+            )
 
         # get the clearml.conf encoded file
         # noinspection PyProtectedMember