Cvs 62354/adopt download engine for onezoo (#3087)

* add sha256

* more flexible file source processing

* fix linter

* fix linter

* remove whitespaces

Co-authored-by: Julia Kamelina <[email protected]>

* remove whitespaces

Co-authored-by: Julia Kamelina <[email protected]>

* remove useless dosc

Co-authored-by: Julia Kamelina <[email protected]>

Author: leramorozova

tools/model_tools/src/openvino/model_zoo/_configuration.py

@@ -51,7 +51,7 @@ class Model:
     def __init__(
         self, name, subdirectory, files, postprocessing, mo_args, framework,
         description, license_url, precisions, quantization_output_precisions,
-        task_type, conversion_to_onnx_args, composite_model_name
+        task_type, conversion_to_onnx_args, converter_to_onnx, composite_model_name
     ):
         self.name = name
         self.subdirectory = subdirectory
@@ -65,12 +65,17 @@ def __init__(
         self.quantization_output_precisions = quantization_output_precisions
         self.task_type = task_type
         self.conversion_to_onnx_args = conversion_to_onnx_args
-        self.converter_to_onnx = _common.KNOWN_FRAMEWORKS[framework]
+        self.converter_to_onnx = converter_to_onnx
         self.composite_model_name = composite_model_name
         self.model_stages = {}
 
     @classmethod
-    def deserialize(cls, model, name, subdirectory, composite_model_name):
+    def deserialize(cls, model, name, subdirectory, composite_model_name, known_frameworks=None, known_task_types=None):
+        if known_frameworks is None:
+            known_frameworks = _common.KNOWN_FRAMEWORKS
+        if known_task_types is None:
+            known_task_types = _common.KNOWN_TASK_TYPES
+
         with validation.deserialization_context('In model "{}"'.format(name)):
             if not RE_MODEL_NAME.fullmatch(name):
                 raise validation.DeserializationError('Invalid name, must consist only of letters, digits or ._-')
@@ -93,10 +98,10 @@ def deserialize(cls, model, name, subdirectory, composite_model_name):
                     postprocessings.append(postprocessing.Postproc.deserialize(postproc))
 
             framework = validation.validate_string_enum('"framework"', model['framework'],
-                _common.KNOWN_FRAMEWORKS.keys())
+                known_frameworks.keys())
 
             conversion_to_onnx_args = model.get('conversion_to_onnx_args', None)
-            if _common.KNOWN_FRAMEWORKS[framework]:
+            if known_frameworks[framework]:
                 if not conversion_to_onnx_args:
                     raise validation.DeserializationError('"conversion_to_onnx_args" is absent. '
                                                           'Framework "{}" is supported only by conversion to ONNX.'
@@ -155,11 +160,11 @@ def deserialize(cls, model, name, subdirectory, composite_model_name):
             license_url = validation.validate_string('"license"', model['license'])
 
             task_type = validation.validate_string_enum('"task_type"', model['task_type'],
-                _common.KNOWN_TASK_TYPES)
+                known_task_types)
 
             return cls(name, subdirectory, files, postprocessings, mo_args, framework,
                 description, license_url, precisions, quantization_output_precisions,
-                task_type, conversion_to_onnx_args, composite_model_name)
+                task_type, conversion_to_onnx_args, known_frameworks[framework], composite_model_name)
 
 class CompositeModel:
     def __init__(self, name, subdirectory, task_type, model_stages, description, framework,
@@ -177,23 +182,29 @@ def __init__(self, name, subdirectory, task_type, model_stages, description, fra
         self.composite_model_name = composite_model_name
 
     @classmethod
-    def deserialize(cls, model, name, subdirectory, stages):
+    def deserialize(cls, model, name, subdirectory, stages, known_frameworks=None, known_task_types=None):
+        if known_frameworks is None:
+            known_frameworks = _common.KNOWN_FRAMEWORKS
+        if known_task_types is None:
+            known_task_types = _common.KNOWN_TASK_TYPES
+
         with validation.deserialization_context('In model "{}"'.format(name)):
             if not RE_MODEL_NAME.fullmatch(name):
                 raise validation.DeserializationError('Invalid name, must consist only of letters, digits or ._-')
 
-            task_type = validation.validate_string_enum('"task_type"', model['task_type'], _common.KNOWN_TASK_TYPES)
+            task_type = validation.validate_string_enum('"task_type"', model['task_type'], known_task_types)
 
             description = validation.validate_string('"description"', model['description'])
 
             license_url = validation.validate_string('"license"', model['license'])
 
             framework = validation.validate_string_enum('"framework"', model['framework'],
-                _common.KNOWN_FRAMEWORKS.keys())
+                                                        known_frameworks)
 
             model_stages = []
             for model_subdirectory, model_part in stages.items():
-                model_stages.append(Model.deserialize(model_part, model_subdirectory.name, model_subdirectory, name))
+                model_stages.append(Model.deserialize(model_part, model_subdirectory.name, model_subdirectory, name,
+                                                      known_frameworks=known_frameworks, known_task_types=known_task_types))
 
             quantization_output_precisions = model_stages[0].quantization_output_precisions
             precisions = model_stages[0].precisions

tools/model_tools/src/openvino/model_zoo/download_engine/downloader.py

@@ -12,22 +12,68 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import contextlib
 import functools
 import requests
 import ssl
+import sys
+import threading
 import time
 import types
 
+from pathlib import Path
+from typing import Set
+
+from openvino.model_zoo import _common, _concurrency, _reporting
 from openvino.model_zoo.download_engine import cache
 
 DOWNLOAD_TIMEOUT = 5 * 60
 
+
+# There is no evidence that the requests.Session class is thread-safe,
+# so for safety, we use one Session per thread. This class ensures that
+# each thread gets its own Session.
+class ThreadSessionFactory:
+    def __init__(self, exit_stack):
+        self._lock = threading.Lock()
+        self._thread_local = threading.local()
+        self._exit_stack = exit_stack
+
+    def __call__(self):
+        try:
+            session = self._thread_local.session
+        except AttributeError:
+            with self._lock: # ExitStack might not be thread-safe either
+                session = self._exit_stack.enter_context(requests.Session())
+            self._thread_local.session = session
+        return session
+
+
 class Downloader:
-    def __init__(self, output_dir=None, cache_dir=None, num_attempts=1, timeout=DOWNLOAD_TIMEOUT):
+    def __init__(self, requested_precisions: str = None, output_dir: Path = None,
+                 cache_dir: Path = None, num_attempts: int = 1, timeout: int = DOWNLOAD_TIMEOUT):
         self.output_dir = output_dir
         self.cache = cache.NullCache() if cache_dir is None else cache.DirCache(cache_dir)
         self.num_attempts = num_attempts
         self.timeout = timeout
+        self.requested_precisions = requested_precisions
+
+    @property
+    def requested_precisions(self) -> Set[str]:
+        return self._requested_precisions
+
+    @requested_precisions.setter
+    def requested_precisions(self, value: str = None):
+        if value is None:
+            _requested_precisions = _common.KNOWN_PRECISIONS
+        else:
+            _requested_precisions = set(value.split(','))
+
+        unknown_precisions = _requested_precisions - _common.KNOWN_PRECISIONS
+        if unknown_precisions:
+            sys.exit('Unknown precisions specified: {}.'.format(', '.join(sorted(unknown_precisions))))
+
+        self._requested_precisions = _requested_precisions
 
     def _process_download(self, reporter, chunk_iterable, size, progress, file):
         start_time = time.monotonic()
@@ -72,12 +118,12 @@ def _try_download(self, reporter, file, start_download, size, hasher):
 
             try:
                 reporter.job_context.check_interrupted()
-                chunk_iterable, continue_offset = start_download(offset=progress.size, timeout=self.timeout)
+                chunk_iterable, continue_offset = start_download(offset=progress.size, timeout=self.timeout, size=size, checksum=hasher)
 
                 if continue_offset not in {0, progress.size}:
                     # Somehow we neither restarted nor continued from where we left off.
                     # Try to restart.
-                    chunk_iterable, continue_offset = start_download(offset=0, timeout=self.timeout)
+                    chunk_iterable, continue_offset = start_download(offset=0, timeout=self.timeout, size=size, checksum=hasher)
                     if continue_offset != 0:
                         reporter.log_error("Remote server refuses to send whole file, aborting")
                         return None
@@ -86,7 +132,7 @@ def _try_download(self, reporter, file, start_download, size, hasher):
                     file.seek(0)
                     file.truncate()
                     progress.size = 0
-                    progress.hasher = hasher()
+                    progress.hasher = hasher.type()
 
                 self._process_download(reporter, chunk_iterable, size, progress, file)
 
@@ -131,6 +177,14 @@ def _try_update_cache(reporter, cache, hash, source):
         except Exception:
             reporter.log_warning('Failed to update the cache', exc_info=True)
 
+    @staticmethod
+    def make_reporter(progress_format: str, context=None):
+        if context is None:
+            context = _reporting.DirectOutputContext()
+        return _reporting.Reporter(context,
+                                   enable_human_output=progress_format == 'text',
+                                   enable_json_output=progress_format == 'json')
+
     def _try_retrieve(self, reporter, destination, model_file, start_download):
         destination.parent.mkdir(parents=True, exist_ok=True)
 
@@ -142,7 +196,7 @@ def _try_retrieve(self, reporter, destination, model_file, start_download):
         success = False
 
         with destination.open('w+b') as f:
-            actual_hash = self._try_download(reporter, f, start_download, model_file.size, model_file.checksum.type)
+            actual_hash = self._try_download(reporter, f, start_download, model_file.size, model_file.checksum)
 
         if actual_hash and cache.verify_hash(reporter, actual_hash, model_file.checksum.value, destination):
             self._try_update_cache(reporter, self.cache, model_file.checksum.value, destination)
@@ -151,7 +205,10 @@ def _try_retrieve(self, reporter, destination, model_file, start_download):
         reporter.print()
         return success
 
-    def download_model(self, reporter, session_factory, requested_precisions, model, known_precisions):
+    def _download_model(self, reporter, session_factory, model, known_precisions: set = None):
+        if known_precisions is None:
+            known_precisions = _common.KNOWN_PRECISIONS
+
         session = session_factory()
 
         reporter.print_group_heading('Downloading {}', model.name)
@@ -164,7 +221,7 @@ def download_model(self, reporter, session_factory, requested_precisions, model,
         for model_file in model.files:
             if len(model_file.name.parts) == 2:
                 p = model_file.name.parts[0]
-                if p in known_precisions and p not in requested_precisions:
+                if p in known_precisions and p not in self.requested_precisions:
                     continue
 
             model_file_reporter = reporter.with_event_context(model=model.name, model_file=model_file.name.as_posix())
@@ -173,7 +230,8 @@ def download_model(self, reporter, session_factory, requested_precisions, model,
             destination = output / model_file.name
 
             if not self._try_retrieve(model_file_reporter, destination, model_file,
-                    functools.partial(model_file.source.start_download, session, cache.CHUNK_SIZE)):
+                    functools.partial(model_file.source.start_download, session, cache.CHUNK_SIZE,
+                                      size=model_file.size, checksum=model_file.checksum)):
                 try:
                     destination.unlink()
                 except FileNotFoundError:
@@ -198,3 +256,25 @@ def download_model(self, reporter, session_factory, requested_precisions, model,
             reporter.print()
 
         return True
+
+    def download_model(self, model, reporter, session):
+        if model.model_stages:
+            results = []
+            for model_stage in model.model_stages:
+                results.append(self._download_model(reporter, session, model_stage))
+            return sum(results) == len(model.model_stages)
+        else:
+            return self._download_model(reporter, session, model)
+
+    def bulk_download_model(self, models, reporter, jobs: int, progress_format: str) -> Set[str]:
+        with contextlib.ExitStack() as exit_stack:
+            session_factory = ThreadSessionFactory(exit_stack)
+            if jobs == 1:
+                results = [self.download_model(model, reporter, session_factory) for model in models]
+            else:
+                results = _concurrency.run_in_parallel(jobs,
+                    lambda context, model: self.download_model(model, self.make_reporter(progress_format, context),
+                                                               session_factory),
+                    models)
+
+        return {model.name for model, successful in zip(models, results) if not successful}

tools/model_tools/src/openvino/model_zoo/download_engine/file_source.py

@@ -67,7 +67,7 @@ def __init__(self, url):
     def deserialize(cls, source):
         return cls(validation.validate_string('"url"', source['url']))
 
-    def start_download(self, session, chunk_size, offset, timeout):
+    def start_download(self, session, chunk_size, offset, timeout, **kwargs):
         response = session.get(self.url, stream=True, timeout=timeout,
             headers=self.http_range_headers(offset))
         response.raise_for_status()
@@ -84,7 +84,7 @@ def __init__(self, id):
     def deserialize(cls, source):
         return cls(validation.validate_string('"id"', source['id']))
 
-    def start_download(self, session, chunk_size, offset, timeout):
+    def start_download(self, session, chunk_size, offset, timeout, **kwargs):
         range_headers = self.http_range_headers(offset)
         URL = 'https://docs.google.com/uc?export=download'
         response = session.get(URL, params={'id': self.id}, headers=range_headers,