Don't wait till training data processing gets finished (#7)

* Release lock as soon as data processing is registered

* Fix typo

* Remove timeout test case for training data upload

Author: tidylobster

aws/traffic_shadowing/src/errors.py

@@ -22,7 +22,3 @@ class ModelRegistrationFailed(Exception):
 
 class DataUploadFailed(Exception):
     pass
-
-
-class TimeOut(Exception):
-    pass

aws/traffic_shadowing/src/model_pool.py

@@ -111,14 +111,13 @@ def create_model(
             metadata: dict = None
     ) -> Model:
         """
-        Register an external model in the Hydrosphere platform
-        and uploads training data.
+        Register an external model and send training data.
         """
-        registration_response = self._register_model(name, schema, metadata)
+        response = self._register_model(name, schema, metadata)
         model = Model(
-            name=registration_response["model"]["name"],
-            version=registration_response["modelVersion"],
-            model_version_id=registration_response["id"]
+            name=response["model"]["name"],
+            version=response["modelVersion"],
+            model_version_id=response["id"]
         )
         self._upload_training_data(model.model_version_id, training_file)
         self._wait_for_data_processing(model.model_version_id)
@@ -127,39 +126,41 @@ def create_model(
     def _wait_for_data_processing(
             self,
             model_version_id,
-            timeout: int = 120,
             retry: int = 3
     ) -> requests.Response:
-        """Wait till the data gets processed."""
+        """Wait till the data gets registered or finishes the processing."""
+
+        def tick(sleep: int = 5) -> bool:
+            nonlocal retry
+            if retry == 0: 
+                return False
+            retry -= 1
+            time.sleep(sleep)
+            return True
+
         url = urllib.parse.urljoin(
             self.endpoint, f"/monitoring/profiles/batch/{model_version_id}/status")
-
         result = None
         while True:
             result = requests.get(url)
             if result.status_code != 200:
-                if retry > 0:
-                    retry -= 1
-                    time.sleep(5)
+                if tick(): 
                     continue
                 else:
                     raise errors.ApiNotAvailable(
                         "Could not fetch the status of the data processing task.")
 
             body = result.json()
             status = DataProfileStatus[body["kind"]]
-            if status == DataProfileStatus.Processing:
-                if timeout > 0:
-                    seconds = min(10, timeout)
-                    timeout -= seconds
-                    time.sleep(seconds)
+
+            if status in (DataProfileStatus.Success, DataProfileStatus.Processing): 
+                # Don't wait until all data gets processed, release immediately
+                break 
+            elif status == DataProfileStatus.NotRegistered:
+                # If profile hasn't been registered yet, wait for a little more
+                if tick(): 
                     continue
-                else:
-                    raise errors.TimeOut("Data processing timed out.")
-            elif status == DataProfileStatus.Success:
-                break
-            else:
-                raise errors.DataUploadFailed(f"Failed to upload the data: {status}")
+            raise errors.DataUploadFailed(f"Failed to upload the data: {status}")
         return result
 
     def _upload_training_data(self, model_version_id: int, training_file: str) -> requests.Response:
@@ -189,8 +190,7 @@ def _register_model(self, name: str, schema: SchemaDescription, metadata: dict)
             response = result.json()
             self.logger.info("Registered a new model: %s", response["model"]["name"])
         else:
-            raise errors.ModelRegistrationFailed(
-                f"Could not register a model: {result.content}")
+            raise errors.ModelRegistrationFailed(f"Could not register a model: {result.content}")
         return response
 
     def _create_feature(self, column: ColumnDescription) -> dict:

aws/traffic_shadowing/tests/test_model_pool.py

@@ -3,7 +3,7 @@
 import requests_mock
 from src.model_pool import ModelPool
 from src.errors import (
-    ModelNotFound, DataUploadFailed, TimeOut, ApiNotAvailable
+    ModelNotFound, DataUploadFailed, ApiNotAvailable
 )
 from tests.stubs.http.hydrosphere import (
     ListModelsStub, ListModelVersionsStub, RegisterExternalModelStub,
@@ -106,17 +106,6 @@ def test_wait_training_data_processed():
         assert result.status_code == 200
 
 
-def test_wait_training_data_processed_timeout():
-    with requests_mock.mock(real_http=False) as mock:
-        mock.get(**WaitTrainingDataProcessingStub(
-            kind="Processing",
-            model_version_id=MODEL_VERSION_ID,
-        ).generate_response())
-        with pytest.raises(TimeOut):
-            pool = ModelPool(HYDROSPHERE_ENDPOINT)
-            pool._wait_for_data_processing(MODEL_VERSION_ID, timeout=1, retry=0)
-
-
 def test_wait_training_data_processed_fail():
     with requests_mock.mock(real_http=False) as mock:
         mock.get(**WaitTrainingDataProcessingStub(
@@ -125,7 +114,7 @@ def test_wait_training_data_processed_fail():
         ).generate_response())
         with pytest.raises(DataUploadFailed):
             pool = ModelPool(HYDROSPHERE_ENDPOINT)
-            pool._wait_for_data_processing(MODEL_VERSION_ID, timeout=1, retry=0)
+            pool._wait_for_data_processing(MODEL_VERSION_ID, retry=0)
 
 
 def test_wait_training_data_processed_not_registered():
@@ -136,7 +125,7 @@ def test_wait_training_data_processed_not_registered():
         ).generate_response())
         with pytest.raises(DataUploadFailed):
             pool = ModelPool(HYDROSPHERE_ENDPOINT)
-            pool._wait_for_data_processing(MODEL_VERSION_ID, timeout=1, retry=0)
+            pool._wait_for_data_processing(MODEL_VERSION_ID, retry=0)
 
 
 def test_wait_training_data_processed_unavailable():
@@ -147,7 +136,7 @@ def test_wait_training_data_processed_unavailable():
         ).generate_client_error())
         with pytest.raises(ApiNotAvailable):
             pool = ModelPool(HYDROSPHERE_ENDPOINT)
-            pool._wait_for_data_processing(MODEL_VERSION_ID, timeout=1, retry=0)
+            pool._wait_for_data_processing(MODEL_VERSION_ID, retry=0)
 
 
 def test_create_model():