added upd skip if no need

supervisely · Nov 20, 2023 · ce477f1 · ce477f1
1 parent 2bbd3c6
commit ce477f1
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 9 deletions.
diff --git a/agent/worker/agent.py b/agent/worker/agent.py
@@ -18,7 +18,7 @@
 import supervisely_lib as sly
 
 from worker.agent_utils import TaskDirCleaner, AppDirCleaner
-from worker.task_update import check_and_pull_sly_net
+from worker.task_update import check_and_pull_sly_net_if_needed
 
 warnings.filterwarnings(action="ignore", category=UserWarning)
 
@@ -107,6 +107,10 @@ def _remove_old_agent(self):
         agent_same_token[0].rename(agent_name_start)
 
     def _update_net_client(self, dc: docker.DockerClient):
+        need_update_env = os.getenv("UPDATE_SLY_NET_AFTER_RESTART", None)
+        if need_update_env == "0":
+            return
+
         net_container_name = "supervisely-net-client-{}".format(constants.TOKEN())
         sly_net_hub_name = "supervisely/sly-net-client:latest"
         sly_net_container = None
@@ -128,10 +132,14 @@ def _update_net_client(self, dc: docker.DockerClient):
             )
             return
         else:
-            # pull if update to old agent
-            agnet_checked_sly_net = os.getenv("AGNET_CHECKED_SLY_NET", None)
-            if agnet_checked_sly_net is None:
-                check_and_pull_sly_net(dc, sly_net_container, self.logger, sly_net_hub_name)
+            # pull if update too old agent
+            if need_update_env is None:
+                need_update = check_and_pull_sly_net_if_needed(
+                    dc, sly_net_container, self.logger, sly_net_hub_name
+                )
+
+        if need_update is False:
+            return
 
         network = "supervisely-net-{}".format(constants.TOKEN())
         command = sly_net_container.attrs.get("Args")

diff --git a/agent/worker/task_update.py b/agent/worker/task_update.py
@@ -62,7 +62,6 @@ def task_main_func(self):
                 new_envs.append(val)
         cur_envs = new_envs
         cur_envs.append("REMOVE_OLD_AGENT={}".format(cur_container_id))
-        cur_envs.append("AGNET_CHECKED_SLY_NET=1")
 
         # Pull net-client if needed
         net_container_name = "supervisely-net-client-{}".format(constants.TOKEN())
@@ -79,9 +78,13 @@ def task_main_func(self):
                 "Something goes wrong: can't find sly-net-client attached to this agent"
             )
         else:
-            check_and_pull_sly_net(
+            need_update = check_and_pull_sly_net_if_needed(
                 self._docker_api, sly_net_container, self.logger, sly_net_hub_name
             )
+            if need_update is True:
+                cur_envs.append("UPDATE_SLY_NET_AFTER_RESTART=1")
+            else:
+                cur_envs.append("UPDATE_SLY_NET_AFTER_RESTART=0")
 
         # start new agent
         container = self._docker_api.containers.run(
@@ -118,18 +121,20 @@ def run_shell_command(cmd, print_output=False):
     return res
 
 
-def check_and_pull_sly_net(
+def check_and_pull_sly_net_if_needed(
     dc: docker.DockerClient,
     cur_container: Container,
     logger: Logger,
     sly_net_hub_name: str = "supervisely/sly-net-client:latest",
-):
+) -> bool:
     ic = ImageCollection(dc)
     docker_hub_image_info = ic.get_registry_data(sly_net_hub_name)
     name_with_digest: str = cur_container.image.attrs.get("RepoDigests", [""])[0]
 
     if name_with_digest.endswith(docker_hub_image_info.id):
         logger.info("sly-net-client is already updated")
+        return False
     else:
         logger.info("Found new version of sly-net-client. Pulling...")
         sly.docker_utils._docker_pull_progress(dc, sly_net_hub_name, logger)
+        return True