Merge 4941546 into 5250c6a

ark/utils/deepcell_service_utils.py

@@ -1,20 +1,23 @@
-from ark.utils import io_utils
-import requests
+from concurrent.futures import ThreadPoolExecutor
+from json import JSONDecodeError
+import os
 from pathlib import Path
+import requests
+from requests.adapters import HTTPAdapter
+from requests.exceptions import RetryError
+from requests.packages.urllib3.util.retry import Retry
 import time
 from tqdm.notebook import tqdm
 from urllib.parse import unquote_plus
-import os
-from zipfile import ZipFile, ZIP_DEFLATED
 import warnings
-from concurrent.futures import ThreadPoolExecutor
+from zipfile import ZipFile, ZIP_DEFLATED
 
-from ark.utils import misc_utils
+from ark.utils import io_utils, misc_utils
 
 
 def create_deepcell_output(deepcell_input_dir, deepcell_output_dir, fovs=None,
                            suffix='_feature_0', host='https://deepcell.org', job_type='mesmer',
-                           scale=1.0, timeout=3600, zip_size=10, parallel=False):
+                           scale=1.0, timeout=3600, zip_size=5, parallel=False):
     """Handles all of the necessary data manipulation for running deepcell tasks.
     Creates .zip files (to be used as input for DeepCell),
     calls run_deepcell_task method,
@@ -113,10 +116,15 @@ def zip_write(zip_path):
 
         # pass the zip file to deepcell.org
         print('Uploading files to DeepCell server.')
-        run_deepcell_direct(
+        status = run_deepcell_direct(
             zip_path, deepcell_output_dir, host, job_type, scale, timeout
         )
 
+        # ensure execution is halted if run_deepcell_direct returned non-zero exit code
+        if status != 0:
+            print("The following FOVs could not be processed: %s" % ','.join(fov_group))
+            return
+
         # extract the .tif output
         print("Extracting tif files from DeepCell response.")
         zip_names = io_utils.list_files(deepcell_output_dir, substrs=[".zip"])
@@ -144,7 +152,7 @@ def zip_write(zip_path):
 
 
 def run_deepcell_direct(input_dir, output_dir, host='https://deepcell.org',
-                        job_type='mesmer', scale=1.0, timeout=3600):
+                        job_type='mesmer', scale=1.0, timeout=3600, num_retries=5):
     """Uses direct calls to DeepCell API and saves output to output_dir.
 
     Args:
@@ -163,6 +171,8 @@ def run_deepcell_direct(input_dir, output_dir, host='https://deepcell.org',
         timeout (int):
             Approximate seconds until timeout.
             Default: 1 hour (3600)
+        num_retries (int):
+            The maximum number of times to call the Deepcell API in case of failure
     """
 
     # upload zip file
@@ -175,11 +185,47 @@ def run_deepcell_direct(input_dir, output_dir, host='https://deepcell.org',
         }
         f.seek(0)
 
-    upload_response = requests.post(
-        upload_url,
-        timeout=timeout,
-        files=upload_fields
-    ).json()
+    # define and mount a retry instance to call the Deepcell API again if needed
+    retry_strategy = Retry(
+        total=num_retries,
+        status_forcelist=[404, 500, 502, 503, 504],
+        method_whitelist=['HEAD', 'GET', 'POST', 'PUT', 'DELETE', 'OPTIONS', 'TRACE']
+    )
+    adapter = HTTPAdapter(max_retries=retry_strategy)
+
+    http = requests.Session()
+    http.mount('https://', adapter)
+    http.mount('http://', adapter)
+
+    total_retries = 0
+    while total_retries < num_retries:
+        # handles the case if the main endpoint can't be reached
+        try:
+            upload_response = http.post(
+                upload_url,
+                timeout=timeout,
+                files=upload_fields
+            )
+        except RetryError as re:
+            print(re)
+            return 1
+
+        # handles the case if the endpoint returns an invalid JSON
+        # indicating an internal API error
+        try:
+            upload_response = upload_response.json()
+        except JSONDecodeError as jde:
+            total_retries += 1
+            continue
+
+        # if we reach the end no errors were encountered on this attempt
+        break
+
+    # if the JSON could not be decoded num_retries number of times
+    if total_retries == num_retries:
+        print("The JSON response from DeepCell could not be decoded after %d attempts" %
+              num_retries)
+        return 1
 
     # call prediction
     predict_url = host + '/api/predict'
@@ -252,4 +298,4 @@ def run_deepcell_direct(input_dir, output_dir, host='https://deepcell.org',
         }
     )
 
-    return
+    return 0

ark/utils/deepcell_service_utils_test.py

@@ -26,6 +26,8 @@ def mocked_run_deepcell(in_zip_path, output_dir, host, job_type, scale, timeout)
             zipObj.write(filename, os.path.basename(filename))
             os.remove(filename)
 
+    return 0
+
 
 def test_create_deepcell_output(mocker: MockerFixture):
     with tempfile.TemporaryDirectory() as temp_dir: