Fix "no" gpu flag (#49)

supervisely · Nov 15, 2023 · 97f7fd4 · 97f7fd4
1 parent 243b238
commit 97f7fd4
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 20 deletions.
diff --git a/agent/worker/task_app.py b/agent/worker/task_app.py
@@ -73,12 +73,12 @@ def from_config(cls, config: dict) -> GPUFlag:
 
     @classmethod
     def from_str(cls, config_val: Optional[str]) -> GPUFlag:
-        if config_val == "preferred":
+        if config_val is None or config_val.lower() == "no":
+            return GPUFlag.skipped
+        elif config_val.lower() == "preferred":
             return GPUFlag.preferred
-        elif config_val == "required":
+        elif config_val.lower() == "required":
             return GPUFlag.required
-        elif config_val is None:
-            return GPUFlag.skipped
         raise ValueError(f"Unknown gpu flag found in config: {config_val}")
 
 

diff --git a/agent/worker/task_update.py b/agent/worker/task_update.py
@@ -19,31 +19,38 @@ def docker_api(self, val):
         self._docker_api = val
 
     def task_main_func(self):
-        if constants.TOKEN() != self.info['config']['access_token']:
-            raise RuntimeError('Current token != new token')
+        if constants.TOKEN() != self.info["config"]["access_token"]:
+            raise RuntimeError("Current token != new token")
 
         docker_inspect_cmd = "curl -s --unix-socket /var/run/docker.sock http://localhost/containers/$(hostname)/json"
-        docker_img_info = subprocess.Popen([docker_inspect_cmd],
-                                           shell=True, executable="/bin/bash",
-                                           stdout=subprocess.PIPE).communicate()[0]
+        docker_img_info = subprocess.Popen(
+            [docker_inspect_cmd], shell=True, executable="/bin/bash", stdout=subprocess.PIPE
+        ).communicate()[0]
         docker_img_info = json.loads(docker_img_info)
 
         cur_container_id = docker_img_info["Config"]["Hostname"]
         cur_volumes = docker_img_info["HostConfig"]["Binds"]
         cur_envs = docker_img_info["Config"]["Env"]
 
-        if docker_img_info["Config"]["Labels"].get("com.docker.compose.project", None) == "supervisely":
-            raise RuntimeError('Docker container was started from docker-compose. Please, use docker-compose to upgrade.')
+        if (
+            docker_img_info["Config"]["Labels"].get("com.docker.compose.project", None)
+            == "supervisely"
+        ):
+            raise RuntimeError(
+                "Docker container was started from docker-compose. Please, use docker-compose to upgrade."
+            )
             return
 
-        sly.docker_utils._docker_pull_progress(self._docker_api, self.info['docker_image'], self.logger)
+        sly.docker_utils._docker_pull_progress(
+            self._docker_api, self.info["docker_image"], self.logger
+        )
 
         new_volumes = {}
         for vol in cur_volumes:
             parts = vol.split(":")
             src = parts[0]
             dst = parts[1]
-            new_volumes[src] = {'bind': dst, 'mode': 'rw'}
+            new_volumes[src] = {"bind": dst, "mode": "rw"}
 
         new_envs = []
         for val in cur_envs:
@@ -53,20 +60,23 @@ def task_main_func(self):
         cur_envs.append("REMOVE_OLD_AGENT={}".format(cur_container_id))
 
         container = self._docker_api.containers.run(
-            self.info['docker_image'],
-            runtime=self.info['config']['docker_runtime'],
+            self.info["docker_image"],
+            runtime=self.info["config"]["docker_runtime"],
             detach=True,
-            name='supervisely-agent-{}-{}'.format(constants.TOKEN(), sly.rand_str(5)),
+            name="supervisely-agent-{}-{}".format(constants.TOKEN(), sly.rand_str(5)),
             remove=False,
             restart_policy={"Name": "unless-stopped"},
             volumes=new_volumes,
             environment=cur_envs,
             stdin_open=False,
-            tty=False
+            tty=False,
         )
         container.reload()
-        self.logger.debug('After spawning. Container status: {}'.format(str(container.status)))
-        self.logger.info('Docker container is spawned', extra={'container_id': container.id, 'container_name': container.name})
+        self.logger.debug("After spawning. Container status: {}".format(str(container.status)))
+        self.logger.info(
+            "Docker container is spawned",
+            extra={"container_id": container.id, "container_name": container.name},
+        )
 
 
 def run_shell_command(cmd, print_output=False):
@@ -80,4 +90,4 @@ def run_shell_command(cmd, print_output=False):
         res.append(clean_line)
         if print_output:
             print(clean_line)
-    return res
+    return res