Streaming model download

Author: Kelvin-M

backend/substrapp/tasks/tasks.py

@@ -21,11 +21,11 @@
 from celery.task import Task
 
 from backend.celery import app
-from substrapp.utils import get_hash, get_owner, create_directory, uncompress_content, compute_hash
+from substrapp.utils import get_hash, get_owner, create_directory, uncompress_content
 from substrapp.ledger_utils import (log_start_tuple, log_success_tuple, log_fail_tuple,
                                     query_tuples, LedgerError, LedgerStatusError, get_object_from_ledger)
-from substrapp.tasks.utils import (ResourcesManager, compute_docker, get_asset_content, list_files,
-                                   get_k8s_client, timeit)
+from substrapp.tasks.utils import (ResourcesManager, compute_docker, get_asset_content, get_and_put_asset_content,
+                                   list_files, get_k8s_client, timeit)
 from substrapp.tasks.exception_handler import compute_error_code
 
 logger = logging.getLogger(__name__)
@@ -136,52 +136,61 @@ def find_training_step_tuple_from_key(tuple_key):
         f'Key {tuple_key}: no tuple found for training step: model: {metadata}')
 
 
-def get_model_content(tuple_type, tuple_key, tuple_, out_model):
+def get_and_put_model_content(tuple_type, tuple_key, tuple_, out_model, model_dst_path):
     """Get out model content."""
     owner = tuple_get_owner(tuple_type, tuple_)
-    return get_asset_content(
+    return get_and_put_asset_content(
         out_model['storageAddress'],
         owner,
         out_model['hash'],
         salt=tuple_key,
+        content_dst_path=model_dst_path
     )
 
 
-def get_local_model_content(tuple_key, out_model):
+def get_and_put_local_model_content(tuple_key, out_model, model_dst_path):
     """Get local model content."""
     from substrapp.models import Model
 
     model = Model.objects.get(pk=out_model['hash'])
-    model_content = model.file.read()
-    computed_hash = compute_hash(model_content, key=tuple_key)
 
-    if computed_hash != out_model['hash']:
-        raise Exception(f"Local model fetch error: hash doesn't match {out_model['hash']} vs {computed_hash}")
+    # verify that local db model file is not corrupted
+    if get_hash(model.file.path, tuple_key) != out_model['hash']:
+        raise Exception('Local Model Hash in Subtuple is not the same as in local db')
 
-    return model_content
+    if not os.path.exists(model_dst_path):
+        os.link(model.file.path, model_dst_path)
+    else:
+        # verify that local subtuple model file is not corrupted
+        if get_hash(model_dst_path, tuple_key) != out_model['hash']:
+            raise Exception('Local Model Hash in Subtuple is not the same as in local medias')
 
 
 @timeit
-def get_and_put_model_content(parent_tuple_type, authorized_types, input_model, directory):
+def fetch_model(parent_tuple_type, authorized_types, input_model, directory):
 
     tuple_type, metadata = find_training_step_tuple_from_key(input_model['traintupleKey'])
 
     if tuple_type not in authorized_types:
         raise TasksError(f'{parent_tuple_type.capitalize()}: invalid input model: type={tuple_type}')
 
+    model_dst_path = path.join(directory, f'model/{input_model["traintupleKey"]}')
+
     if tuple_type == TRAINTUPLE_TYPE:
-        model = get_model_content(tuple_type, input_model['traintupleKey'], metadata, metadata['outModel'])
+        get_and_put_model_content(
+            tuple_type, input_model['traintupleKey'], metadata, metadata['outModel'], model_dst_path
+        )
     elif tuple_type == AGGREGATETUPLE_TYPE:
-        model = get_model_content(tuple_type, input_model['traintupleKey'], metadata, metadata['outModel'])
+        get_and_put_model_content(
+            tuple_type, input_model['traintupleKey'], metadata, metadata['outModel'], model_dst_path
+        )
     elif tuple_type == COMPOSITE_TRAINTUPLE_TYPE:
-        model = get_model_content(
-            tuple_type, input_model['traintupleKey'], metadata, metadata['outTrunkModel']['outModel']
+        get_and_put_model_content(
+            tuple_type, input_model['traintupleKey'], metadata, metadata['outTrunkModel']['outModel'], model_dst_path
         )
     else:
         raise TasksError(f'Traintuple: invalid input model: type={tuple_type}')
 
-    _put_model(directory, model, input_model['hash'], input_model['traintupleKey'])
-
 
 def prepare_traintuple_input_models(directory, tuple_):
     """Get traintuple input models content."""
@@ -193,7 +202,7 @@ def prepare_traintuple_input_models(directory, tuple_):
 
     models = []
     for input_model in input_models:
-        proc = Thread(target=get_and_put_model_content,
+        proc = Thread(target=fetch_model,
                       args=(TRAINTUPLE_TYPE, authorized_types, input_model, directory))
         models.append(proc)
         proc.start()
@@ -212,7 +221,7 @@ def prepare_aggregatetuple_input_models(directory, tuple_):
     models = []
 
     for input_model in input_models:
-        proc = Thread(target=get_and_put_model_content,
+        proc = Thread(target=fetch_model,
                       args=(AGGREGATETUPLE_TYPE, authorized_types, input_model, directory))
         models.append(proc)
         proc.start()
@@ -235,33 +244,27 @@ def prepare_composite_traintuple_input_models(directory, tuple_):
     if tuple_type != COMPOSITE_TRAINTUPLE_TYPE:
         raise TasksError(f'CompositeTraintuple: invalid head input model: type={tuple_type}')
     # get the output head model
-    head_model_content = get_local_model_content(head_model_key, metadata['outHeadModel']['outModel'])
+    head_model_dst_path = path.join(directory, f'model/{PREFIX_HEAD_FILENAME}{head_model_key}')
+    get_and_put_local_model_content(
+        head_model_key, metadata['outHeadModel']['outModel'], head_model_dst_path
+    )
 
     # get trunk model
     trunk_model_key = trunk_model['traintupleKey']
     tuple_type, metadata = find_training_step_tuple_from_key(trunk_model_key)
+    trunk_model_dst_path = path.join(directory, f'model/{PREFIX_TRUNK_FILENAME}{trunk_model_key}')
     # trunk model must refer to a composite traintuple or an aggregatetuple
     if tuple_type == COMPOSITE_TRAINTUPLE_TYPE:  # get output trunk model
-        trunk_model_content = get_model_content(
-            tuple_type, trunk_model_key, metadata, metadata['outTrunkModel']['outModel'],
+        get_and_put_model_content(
+            tuple_type, trunk_model_key, metadata, metadata['outTrunkModel']['outModel'], trunk_model_dst_path
         )
     elif tuple_type == AGGREGATETUPLE_TYPE:
-        trunk_model_content = get_model_content(
-            tuple_type, trunk_model_key, metadata, metadata['outModel'],
+        get_and_put_model_content(
+            tuple_type, trunk_model_key, metadata, metadata['outModel'], trunk_model_dst_path
         )
     else:
         raise TasksError(f'CompositeTraintuple: invalid trunk input model: type={tuple_type}')
 
-    # put head and trunk models
-    _put_model(directory, head_model_content,
-               tuple_['inHeadModel']['hash'],
-               tuple_['inHeadModel']['traintupleKey'],
-               filename_prefix=PREFIX_HEAD_FILENAME)
-    _put_model(directory, trunk_model_content,
-               tuple_['inTrunkModel']['hash'],
-               tuple_['inTrunkModel']['traintupleKey'],
-               filename_prefix=PREFIX_TRUNK_FILENAME)
-
 
 def prepare_testtuple_input_models(directory, tuple_):
     """Get testtuple input models content."""
@@ -272,52 +275,26 @@ def prepare_testtuple_input_models(directory, tuple_):
 
     if traintuple_type == TRAINTUPLE_TYPE:
         metadata = get_object_from_ledger(traintuple_key, 'queryTraintuple')
-        model = get_model_content(traintuple_type, traintuple_key, metadata, metadata['outModel'])
-        model_hash = metadata['outModel']['hash']
-        _put_model(directory, model, model_hash, traintuple_key)
+        model_dst_path = path.join(directory, f'model/{traintuple_key}')
+        get_and_put_model_content(
+            traintuple_type, traintuple_key, metadata, metadata['outModel'], model_dst_path
+        )
 
     elif traintuple_type == COMPOSITE_TRAINTUPLE_TYPE:
         metadata = get_object_from_ledger(traintuple_key, 'queryCompositeTraintuple')
-        head_content = get_local_model_content(traintuple_key, metadata['outHeadModel']['outModel'])
-        trunk_content = get_model_content(
-            traintuple_type, traintuple_key, metadata, metadata['outTrunkModel']['outModel'],
+        head_model_dst_path = path.join(directory, f'model/{PREFIX_HEAD_FILENAME}{traintuple_key}')
+        get_and_put_local_model_content(traintuple_key, metadata['outHeadModel']['outModel'],
+                                        head_model_dst_path)
+
+        model_dst_path = path.join(directory, f'model/{PREFIX_TRUNK_FILENAME}{traintuple_key}')
+        get_and_put_model_content(
+            traintuple_type, traintuple_key, metadata, metadata['outTrunkModel']['outModel'], model_dst_path
         )
-        _put_model(directory, head_content, metadata['outHeadModel']['outModel']['hash'],
-                   traintuple_key, filename_prefix=PREFIX_HEAD_FILENAME)
-        _put_model(directory, trunk_content, metadata['outTrunkModel']['outModel']['hash'],
-                   traintuple_key, filename_prefix=PREFIX_TRUNK_FILENAME)
 
     else:
         raise TasksError(f"Testtuple from type '{traintuple_type}' not supported")
 
 
-def _put_model(subtuple_directory, model_content, model_hash, traintuple_hash, filename_prefix=''):
-    if not model_content:
-        raise Exception('Model content should not be empty')
-
-    from substrapp.models import Model
-
-    # store a model in local subtuple directory from input model content
-    model_dst_path = path.join(subtuple_directory, f'model/{filename_prefix}{traintuple_hash}')
-    model = None
-    try:
-        model = Model.objects.get(pk=model_hash)
-    except ObjectDoesNotExist:  # write it to local disk
-        with open(model_dst_path, 'wb') as f:
-            f.write(model_content)
-    else:
-        # verify that local db model file is not corrupted
-        if get_hash(model.file.path, traintuple_hash) != model_hash:
-            raise Exception('Model Hash in Subtuple is not the same as in local db')
-
-        if not os.path.exists(model_dst_path):
-            os.link(model.file.path, model_dst_path)
-        else:
-            # verify that local subtuple model file is not corrupted
-            if get_hash(model_dst_path, traintuple_hash) != model_hash:
-                raise Exception('Model Hash in Subtuple is not the same as in local medias')
-
-
 @timeit
 def prepare_models(directory, tuple_type, tuple_):
     """Prepare models for tuple execution.

backend/substrapp/tasks/utils.py

@@ -7,7 +7,7 @@
 from subprocess import check_output
 from django.conf import settings
 from requests.auth import HTTPBasicAuth
-from substrapp.utils import get_owner, get_remote_file_content, NodeError
+from substrapp.utils import get_owner, get_remote_file_content, get_and_put_remote_file_content, NodeError
 
 from kubernetes import client, config
 
@@ -47,6 +47,11 @@ def get_asset_content(url, node_id, content_hash, salt=None):
     return get_remote_file_content(url, authenticate_worker(node_id), content_hash, salt=salt)
 
 
+def get_and_put_asset_content(url, node_id, content_hash, salt=None, content_dst_path=None):
+    return get_and_put_remote_file_content(url, authenticate_worker(node_id), content_hash, salt=salt,
+                                           content_dst_path=content_dst_path)
+
+
 @timeit
 def get_cpu_count(client):
     # Get CPU count from docker container through the API

backend/substrapp/utils.py

@@ -158,7 +158,7 @@ class NodeError(Exception):
     pass
 
 
-def get_remote_file(url, auth, **kwargs):
+def get_remote_file(url, auth, content_dst_path=None, **kwargs):
     kwargs.update({
         'headers': {'Accept': 'application/json;version=0.0'},
         'auth': auth
@@ -168,14 +168,21 @@ def get_remote_file(url, auth, **kwargs):
         kwargs['verify'] = False
 
     try:
-        response = requests.get(url, **kwargs)
+        if kwargs.get('stream', False) and content_dst_path is not None:
+            chunk_size = 1024 * 1024
+            with open(content_dst_path, 'wb') as fp:
+                response = requests.get(url, **kwargs)
+                fp.writelines(response.iter_content(chunk_size))
+        else:
+            response = requests.get(url, **kwargs)
     except (requests.exceptions.ConnectionError, requests.exceptions.Timeout) as e:
         raise NodeError(f'Failed to fetch {url}') from e
 
     return response
 
 
 def get_remote_file_content(url, auth, content_hash, salt=None):
+
     response = get_remote_file(url, auth)
 
     if response.status_code != status.HTTP_200_OK:
@@ -186,3 +193,16 @@ def get_remote_file_content(url, auth, content_hash, salt=None):
     if computed_hash != content_hash:
         raise NodeError(f"url {url}: hash doesn't match {content_hash} vs {computed_hash}")
     return response.content
+
+
+def get_and_put_remote_file_content(url, auth, content_hash, salt=None, content_dst_path=None):
+
+    response = get_remote_file(url, auth, content_dst_path, stream=True)
+
+    if response.status_code != status.HTTP_200_OK:
+        logger.error(response.text)
+        raise NodeError(f'Url: {url} returned status code: {response.status_code}')
+
+    computed_hash = get_hash(content_dst_path, key=salt)
+    if computed_hash != content_hash:
+        raise NodeError(f"url {url}: hash doesn't match {content_hash} vs {computed_hash}")