Fast API, many improvements and bug fixes

[Public API] added bulk methods
[agent] fast upload/download
[jupyter] tutorials and cookbooks use new Public API

Author: max-supervisely

agent/Dockerfile

@@ -27,6 +27,8 @@ RUN pip install --no-cache-dir \
         py3exiv2==0.4.0 
 
 RUN pip install requests-toolbelt
+RUN pip install packaging
+
 
 ############### copy code ###############
 ARG MODULE_PATH

agent/VERSION

@@ -1 +1 @@
-agent:4.2.0
+agent:4.3.0

agent/src/worker/agent.py

@@ -6,7 +6,7 @@
 import threading
 from concurrent.futures import ThreadPoolExecutor, wait
 import subprocess
-
+import os
 import supervisely_lib as sly
 
 from worker import constants
@@ -36,8 +36,13 @@ def __init__(self):
         self.thread_list = []
         self.daemons_list = []
 
+        self._remove_old_agent()
+        self._validate_duplicated_agents()
+
         sly.fs.clean_dir(constants.AGENT_TMP_DIR())
-        self._stop_missed_containers()
+        self._stop_missed_containers(constants.TASKS_DOCKER_LABEL())
+        # for compatibility with old plugins
+        self._stop_missed_containers(constants.TASKS_DOCKER_LABEL_LEGACY())
 
         self.docker_api = docker.from_env(version='auto')
         self._docker_login()
@@ -48,6 +53,33 @@ def __init__(self):
         self.agent_connect_initially()
         self.logger.info('Agent connected to server.')
 
+    def _remove_old_agent(self):
+        container_id = os.getenv('REMOVE_OLD_AGENT', None)
+        if container_id is None:
+            return
+
+        dc = docker.from_env()
+        olg_agent = dc.containers.get(container_id)
+        olg_agent.remove(force=True)
+
+        agent_same_token = []
+        for cont in dc.containers.list():
+            if constants.TOKEN() in cont.name:
+                agent_same_token.append(cont)
+
+        if len(agent_same_token) > 1:
+            raise RuntimeError("Several agents with the same token are running. Please, kill them or contact support.")
+        agent_same_token[0].rename('supervisely-agent-{}'.format(constants.TOKEN()))
+
+    def _validate_duplicated_agents(self):
+        dc = docker.from_env()
+        agent_same_token = []
+        for cont in dc.containers.list():
+            if constants.TOKEN() in cont.name:
+                agent_same_token.append(cont)
+        if len(agent_same_token) > 1:
+            raise RuntimeError("Agent with the same token already exists.")
+
     def agent_connect_initially(self):
         try:
             hw_info = get_hw_info()
@@ -63,6 +95,7 @@ def agent_connect_initially(self):
         self.agent_info = {
             'hardware_info': hw_info,
             'agent_image': json.loads(docker_img_info)["Config"]["Image"],
+            'agent_version': json.loads(docker_img_info)["Config"]["Labels"]["VERSION"],
             'agent_image_digest': get_self_docker_image_digest()
         }
 
@@ -121,6 +154,7 @@ def start_task(self, task):
                 self.logger.warning('TASK_ID_ALREADY_STARTED', extra={'task_id': task['task_id']})
             else:
                 task_id = task['task_id']
+                task["agent_version"] = self.agent_info["agent_version"]
                 self.task_pool[task_id] = create_task(task, self.docker_api)
                 self.task_pool[task_id].start()
         finally:
@@ -161,9 +195,9 @@ def _remove_containers(label_filter):
             cont.remove(force=True)
         return stop_list
 
-    def _stop_missed_containers(self):
+    def _stop_missed_containers(self, ecosystem_token):
         self.logger.info('Searching for missed containers...')
-        label_filter = {'label': 'ecosystem_token={}'.format(constants.TASKS_DOCKER_LABEL())}
+        label_filter = {'label': 'ecosystem_token={}'.format(ecosystem_token)}
 
         stopped_list = Agent._remove_containers(label_filter=label_filter)
 

agent/src/worker/constants.py

@@ -3,6 +3,7 @@
 import os
 from urllib.parse import urlparse
 import supervisely_lib as sly
+import hashlib
 
 
 def HOST_DIR():
@@ -23,6 +24,10 @@ def TOKEN():
 
 
 def TASKS_DOCKER_LABEL():
+    return 'supervisely_{}'.format(hashlib.sha256(TOKEN().encode('utf-8')).hexdigest())
+
+
+def TASKS_DOCKER_LABEL_LEGACY():
     return 'supervisely_{}'.format(TOKEN())
 
 

agent/src/worker/data_manager.py

@@ -5,7 +5,16 @@
 from worker.agent_storage import AgentStorage
 from worker.fs_storages import EmptyStorage
 from worker import constants
-from collections import defaultdict
+
+
+def _maybe_append_image_extension(name, ext):
+    name_split = os.path.splitext(name)
+    if name_split[1] == '':
+        normalized_ext = ('.' + ext).replace('..', '.')
+        sly.image.validate_ext(normalized_ext)
+        return name + normalized_ext
+    else:
+        return name
 
 
 class DataManager(object):
@@ -34,10 +43,10 @@ def download_nn(self, name, parent_dir):
             self.logger.info('NN has been copied from local storage.')
             return
 
-        model_in_mb = int(float(model_info.size) / 1024 / 1024)
-        progress = sly.Progress('Download NN: {!r}'.format(name), model_in_mb)
+        model_in_mb = int(float(model_info.size) / 1024 / 1024 + 1)
+        progress = sly.Progress('Download NN: {!r}'.format(name), model_in_mb, self.logger)
 
-        self.public_api.model.download_to_dir(self.workspace_id, name, parent_dir, progress.iter_done_report)
+        self.public_api.model.download_to_dir(self.workspace_id, name, parent_dir, progress.iters_done_report)
         self.logger.info('NN has been downloaded from server.')
 
         if self.has_nn_storage():
@@ -87,34 +96,37 @@ def download_dataset(self, dataset, dataset_id):
                                                                          'images_to_download': len(images_to_download)})
             if len(images_to_download) + len(images_in_cache) != len(images):
                 raise RuntimeError("Error with images cache during download. Please contact support.")
-            for batch_cache in sly.batched(list(zip(images_in_cache, images_cache_paths)), constants.BATCH_SIZE_GET_IMAGES_INFO()):
-                img_cache_ids = [img_info.id for img_info, _ in batch_cache]
+
+            if len(images_in_cache) > 0:
+                img_cache_ids = [img_info.id for img_info in images_in_cache]
                 ann_info_list = self.public_api.annotation.download_batch(dataset_id, img_cache_ids, progress_anns.iters_done_report)
-                img_name_to_ann = {ann.image_name: ann.annotation for ann in ann_info_list}
-                for img_info, img_cache_path in batch_cache:
-                    dataset.add_item_file(img_info.name, img_cache_path, img_name_to_ann[img_info.name])
+                img_name_to_ann = {ann.image_id: ann.annotation for ann in ann_info_list}
+                for img_info, img_cache_path in zip(images_in_cache, images_cache_paths):
+                    item_name = _maybe_append_image_extension(img_info.name, img_info.ext)
+                    dataset.add_item_file(item_name, img_cache_path, img_name_to_ann[img_info.id])
                     progress_imgs.iter_done_report()
 
         # download images from server
-        for batch_download in sly.batched(images_to_download, constants.BATCH_SIZE_GET_IMAGES_INFO()):
+        if len(images_to_download) > 0:
             #prepare lists for api methods
             img_ids = []
             img_paths = []
-            for img_info in batch_download:
+            for img_info in images_to_download:
                 img_ids.append(img_info.id)
                 # TODO download to a temp file and use dataset api to add the image to the dataset.
-                img_paths.append(dataset.deprecated_make_img_path(img_info.name, img_info.ext))
+                img_paths.append(
+                    os.path.join(dataset.img_dir, _maybe_append_image_extension(img_info.name, img_info.ext)))
 
             # download annotations
             ann_info_list = self.public_api.annotation.download_batch(dataset_id, img_ids, progress_anns.iters_done_report)
-            img_name_to_ann = {ann.image_name: ann.annotation for ann in ann_info_list}
-            self.public_api.image.download_batch(dataset_id, img_ids, img_paths, progress_imgs.iters_done_report)
-            for img_info, img_path in zip(batch_download, img_paths):
-                dataset.add_item_file(img_info.name, img_path, img_name_to_ann[img_info.name])
+            img_name_to_ann = {ann.image_id: ann.annotation for ann in ann_info_list}
+            self.public_api.image.download_paths(dataset_id, img_ids, img_paths, progress_imgs.iters_done_report)
+            for img_info, img_path in zip(images_to_download, img_paths):
+                dataset.add_item_file(img_info.name, img_path, img_name_to_ann[img_info.id])
 
             if self.has_images_storage():
                 progress_cache = sly.Progress('Dataset {!r}: cache images'.format(dataset.name), len(img_paths), self.logger)
-                img_hashes = [img_info.hash for img_info in batch_download]
+                img_hashes = [img_info.hash for img_info in images_to_download]
                 self.storage.images.write_objects(img_paths, img_hashes, progress_cache.iter_done_report)
 
     # @TODO: remove legacy stuff
@@ -146,44 +158,31 @@ def upload_project(self, parent_dir, project_name, new_title, legacy=False, add_
         self.logger.info('PROJECT_CREATED',extra={'event_type': sly.EventType.PROJECT_CREATED, 'project_id': project_id})
 
     def upload_dataset(self, dataset, dataset_id):
-        progress = None
+        progress_cache = None
         items_count = len(dataset)
-        hash_to_img_paths = defaultdict(list)
-        hash_to_ann_paths = defaultdict(list)
-        hash_to_item_names = defaultdict(list)
+
+        item_names = []
+        img_paths = []
+        ann_paths = []
         for item_name in dataset:
+            item_names.append(item_name)
             item_paths = dataset.get_item_paths(item_name)
-            img_hash = sly.fs.get_file_hash(item_paths.img_path)
-            hash_to_img_paths[img_hash].append(item_paths.img_path)
-            hash_to_ann_paths[img_hash].append(item_paths.ann_path)
-            hash_to_item_names[img_hash].append(item_name)
+            img_paths.append(item_paths.img_path)
+            ann_paths.append(item_paths.ann_path)
+
             if self.has_images_storage():
-                if progress is None:
-                    progress = sly.Progress('Dataset {!r}: cache images'.format(dataset.name), items_count, self.logger)
+                if progress_cache is None:
+                    progress_cache = sly.Progress('Dataset {!r}: cache images'.format(dataset.name), items_count, self.logger)
+
+                img_hash = sly.fs.get_file_hash(item_paths.img_path)
                 self.storage.images.write_object(item_paths.img_path, img_hash)
-                progress.iter_done_report()
+                progress_cache.iter_done_report()
+
+        progress = sly.Progress('Dataset {!r}: upload images'.format(dataset.name), items_count, self.logger)
+        image_infos = self.public_api.image.upload_paths(dataset_id, item_names, img_paths, progress.iters_done_report)
 
-        progress_img = sly.Progress('Dataset {!r}: upload images'.format(dataset.name), items_count, self.logger)
-        progress_ann = sly.Progress('Dataset {!r}: upload annotations'.format(dataset.name), items_count, self.logger)
-
-        def add_images_annotations(hashes, pb_img_cb, pb_ann_cb):
-            names = [name for hash in hashes for name in hash_to_item_names[hash]]
-            unrolled_hashes = [hash for hash in hashes for _ in range(len(hash_to_item_names[hash]))]
-            ann_paths = [path for hash in hashes for path in hash_to_ann_paths[hash]]
-            remote_infos = self.public_api.image.add_batch(dataset_id, names, unrolled_hashes, pb_img_cb)
-            self.public_api.annotation.upload_batch_paths(dataset_id, [info.id for info in remote_infos], ann_paths, pb_ann_cb)
-
-        # add already uploaded images + attach annotations
-        remote_hashes = self.public_api.image.check_existing_hashes(list(hash_to_img_paths.keys()))
-        if len(remote_hashes) > 0:
-            add_images_annotations(remote_hashes, progress_img.iters_done_report, progress_ann.iters_done_report)
-
-        # upload new images + add annotations
-        new_hashes = list(set(hash_to_img_paths.keys()) - set(remote_hashes))
-        img_paths = [path for hash in new_hashes for path in hash_to_img_paths[hash]]
-        self.public_api.image.upload_batch_paths(img_paths, progress_img.iters_done_report)
-        if len(new_hashes) > 0:
-            add_images_annotations(new_hashes, None, progress_ann.iters_done_report)
+        progress = sly.Progress('Dataset {!r}: upload annotations'.format(dataset.name), items_count, self.logger)
+        self.public_api.annotation.upload_paths([info.id for info in image_infos], ann_paths, progress.iters_done_report)
 
     def upload_archive(self, task_id, dir_to_archive, archive_name):
         self.logger.info("PACK_TO_ARCHIVE ...")

agent/src/worker/task_custom.py

@@ -29,16 +29,14 @@ def init_additional(self):
 
     def download_step(self):
         for model_info in self.info['models']:
-            self.data_mgr.download_nn(model_info['title'], model_info['id'], model_info['hash'], self.dir_model)
+            self.data_mgr.download_nn(model_info['title'], self.dir_model)
 
         self.logger.info("DOWNLOAD_DATA")
         json.dump(self.info['config'], open(self.config_path1, 'w'))  # Deprecated 'task_settings.json'
         json.dump(self.info['config'], open(self.config_path2, 'w'))  # New style task_config.json
 
         for pr_info in self.info['projects']:
-            project = sly.api_proto.Project(id=pr_info['id'], title=pr_info['title'])
-            datasets = [sly.api_proto.Dataset(id=ds['id'], title=ds['title']) for ds in pr_info['datasets']]
-            self.data_mgr.download_project(self.dir_data, project, datasets)
+            self.data_mgr.download_project(self.dir_data, pr_info['title'])
 
         self.report_step_done(TaskStep.DOWNLOAD)
 

agent/src/worker/task_dockerized.py

@@ -4,6 +4,7 @@
 from threading import Lock
 import json
 from docker.errors import DockerException, ImageNotFound as DockerImageNotFound
+from packaging import version
 
 import supervisely_lib as sly
 
@@ -41,8 +42,8 @@ def __init__(self, *args, **kwargs):
 
         self._container = None
         self._container_lock = Lock()  # to drop container from different threads
-        self.docker_image_name = self.info['docker_image']
-        if ':' not in self.docker_image_name:
+        self.docker_image_name = self.info.get('docker_image', None)
+        if self.docker_image_name is not None and ':' not in self.docker_image_name:
             self.docker_image_name += ':latest'
         self.docker_pulled = False  # in task
 
@@ -118,6 +119,34 @@ def _docker_pull(self):
             raise DockerException('Unable to pull image: not enough free disk space or something wrong with DockerHub.'
                                   ' Please, run the task again or email support.')
         self.logger.info('Docker image has been pulled', extra={'pulled': {'tags': pulled_img.tags, 'id': pulled_img.id}})
+        self._validate_version(self.info["agent_version"], pulled_img.labels.get("VERSION", None))
+
+    def _validate_version(self, agent_image, plugin_image):
+        self.logger.info('Check if agent and plugin versions are compatible')
+
+        def get_version(docker_image):
+            if docker_image is None:
+                return None
+            image_parts = docker_image.split(":")
+            if len(image_parts) != 2:
+                return None
+            return image_parts[1]
+
+        agent_version = get_version(agent_image.strip())
+        plugin_version = get_version(plugin_image.strip())
+
+        if agent_version is None or plugin_version is None:
+            self.logger.info('Unknown version')
+
+        av = version.parse(agent_version)
+        pv = version.parse(plugin_version)
+
+        if type(av) is version.LegacyVersion or type(pv) is version.LegacyVersion:
+            self.logger.info('Invalid semantic version, can not compare')
+            return
+
+        if av.release[0] < pv.release[0]:
+            self.logger.critical('Agent version is lower than plugin version. Please, update agent.')
 
     def _docker_image_exists(self):
         try:
@@ -138,17 +167,16 @@ def spawn_container(self, add_envs=None):
             add_envs = {}
         self._container_lock.acquire()
         try:
-            #@TODO: DEBUG_COPY_IMAGES only for compatibility with old plugins
             self._container = self._docker_api.containers.run(
                 self.docker_image_name,
                 runtime=self.docker_runtime,
                 entrypoint=["sh", "-c", "python -u {}".format(self.entrypoint)],
                 detach=True,
-                name='sly_task_{}_{}'.format(constants.TOKEN(), self.info['task_id']),
+                name='sly_task_{}_{}'.format(self.info['task_id'], constants.TASKS_DOCKER_LABEL()),
                 remove=False,
                 volumes={self.dir_task_host: {'bind': '/sly_task_data',
                                               'mode': 'rw'}},
-                environment={'LOG_LEVEL': 'DEBUG', 'LANG': 'C.UTF-8', 'DEBUG_COPY_IMAGES': 1, **add_envs},
+                environment={'LOG_LEVEL': 'DEBUG', 'LANG': 'C.UTF-8', **add_envs},
                 labels={'ecosystem': 'supervisely',
                         'ecosystem_token': constants.TASKS_DOCKER_LABEL(),
                         'task_id': str(self.info['task_id'])},
@@ -158,7 +186,7 @@ def spawn_container(self, add_envs=None):
             )
             self._container.reload()
             self.logger.debug('After spawning. Container status: {}'.format(str(self._container.status)))
-            self.logger.info('Docker container spawned',extra={'container_id': self._container.id, 'container_name': self._container.name})
+            self.logger.info('Docker container is spawned',extra={'container_id': self._container.id, 'container_name': self._container.name})
         finally:
             self._container_lock.release()
 

agent/src/worker/task_factory.py

@@ -13,11 +13,9 @@
 from worker.task_upload_images import TaskUploadImages
 from worker.task_import_local import TaskImportLocal
 from worker.task_custom import TaskCustom
-#from worker.task_pipeline import TaskPipeline
+from worker.task_update import TaskUpdate
 
 
-import importlib
-
 _task_class_mapping = {
     'export':               TaskDTL,
     'import':               TaskImport,
@@ -30,7 +28,7 @@
     'upload_images':        TaskUploadImages,
     'import_agent':         TaskImportLocal,
     'custom': 				TaskCustom,
-    #'pipeline':             TaskPipeline
+    'update_agent':         TaskUpdate
 }
 
 
@@ -40,7 +38,6 @@ def create_task(task_msg, docker_api):
         sly.logger.critical('unknown task type', extra={'task_msg': task_msg})
         raise RuntimeError('unknown task type')
     task_obj = task_cls(task_msg)
-    #@TODO: check condition for pipelines and custom tasks
-    if issubclass(task_cls, TaskDockerized) or (task_msg['task_type'] == 'pipeline'):
+    if issubclass(task_cls, TaskDockerized) or (task_msg['task_type'] == 'update_agent'):
         task_obj.docker_api = docker_api
     return task_obj