Repository: codalab/codalab-worksheets

Issue:
Failed to download dependency celebA: Dependency download failed: IncompleteRead(0 bytes read)
See https://worksheets.codalab.org/bundles/0xc26e560d35d94a2aa475ac8a2acb2df9

Discussion:
adiprerepa: do we know for a fact that retries fixes the issue?
epicfaace: We're pretty sure, see https://github.com/codalab/codalab-worksheets/issues/3457#issuecomment-855001732 re. our conversation last meeting.

Patch:
diff --git a/codalab/worker/dependency_manager.py b/codalab/worker/dependency_manager.py
index d1eb5457..93a7baf8 100644
--- a/codalab/worker/dependency_manager.py
+++ b/codalab/worker/dependency_manager.py
@@ -1,18 +1,19 @@
-from contextlib import closing
-from collections import namedtuple
 import logging
 import os
 import threading
 import traceback
 import time
 import shutil
-from typing import Dict
+from contextlib import closing
+from collections import namedtuple
+from typing import Dict, Set
 
+import codalab.worker.pyjson
+from .bundle_service_client import BundleServiceClient
 from codalab.lib.formatting import size_str
 from codalab.worker.file_util import remove_path
 from codalab.worker.un_tar_directory import un_tar_directory
 from codalab.worker.fsm import BaseDependencyManager, DependencyStage, StateTransitioner
-import codalab.worker.pyjson
 from codalab.worker.worker_thread import ThreadDict
 from codalab.worker.state_committer import JsonStateCommitter
 from codalab.worker.bundle_state import DependencyKey
@@ -52,7 +53,14 @@ class DependencyManager(StateTransitioner, BaseDependencyManager):
     # the data format of how we store this)
     MAX_SERIALIZED_LEN = 60000
 
-    def __init__(self, commit_file, bundle_service, worker_dir, max_cache_size_bytes):
+    def __init__(
+        self,
+        commit_file: str,
+        bundle_service: BundleServiceClient,
+        worker_dir: str,
+        max_cache_size_bytes: int,
+        download_dependencies_max_retries: int,
+    ):
         super(DependencyManager, self).__init__()
         self.add_transition(DependencyStage.DOWNLOADING, self._transition_from_DOWNLOADING)
         self.add_terminal(DependencyStage.READY)
@@ -62,6 +70,7 @@ class DependencyManager(StateTransitioner, BaseDependencyManager):
         self._bundle_service = bundle_service
         self._max_cache_size_bytes = max_cache_size_bytes
         self.dependencies_dir = os.path.join(worker_dir, DependencyManager.DEPENDENCIES_DIR_NAME)
+        self._download_dependencies_max_retries = download_dependencies_max_retries
         if not os.path.exists(self.dependencies_dir):
             logger.info('{} doesn\'t exist, creating.'.format(self.dependencies_dir))
             os.makedirs(self.dependencies_dir, 0o770)
@@ -72,9 +81,8 @@ class DependencyManager(StateTransitioner, BaseDependencyManager):
         self._paths_lock = threading.RLock()  # Used for path name computations
 
         # File paths that are currently being used to store dependencies. Used to prevent conflicts
-        self._paths = set()
-        # DependencyKey -> DependencyState
-        self._dependencies = dict()
+        self._paths: Set[str] = set()
+        self._dependencies: Dict[DependencyKey, DependencyState] = dict()
         # DependencyKey -> WorkerThread(thread, success, failure_message)
         self._downloading = ThreadDict(fields={'success': False, 'failure_message': None})
         self._load_state()
@@ -429,44 +437,57 @@ class DependencyManager(StateTransitioner, BaseDependencyManager):
             # TODO(Ashwin): make this not fs-specific.
             dependency_path = os.path.join(self.dependencies_dir, dependency_state.path)
             logger.debug('Downloading dependency %s', dependency_state.dependency_key)
-            try:
-                # Start async download to the fileobj
-                fileobj, target_type = self._bundle_service.get_bundle_contents(
-                    dependency_state.dependency_key.parent_uuid,
-                    dependency_state.dependency_key.parent_path,
-                )
-                with closing(fileobj):
-                    # "Bug" the fileobj's read function so that we can keep
-                    # track of the number of bytes downloaded so far.
-                    old_read_method = fileobj.read
-                    bytes_downloaded = [0]
-
-                    def interruptable_read(*args, **kwargs):
-                        data = old_read_method(*args, **kwargs)
-                        bytes_downloaded[0] += len(data)
-                        update_state_and_check_killed(bytes_downloaded[0])
-                        return data
-
-                    fileobj.read = interruptable_read
-
-                    # Start copying the fileobj to filesystem dependency path
-                    self._store_dependency(dependency_path, fileobj, target_type)
-
-                logger.debug(
-                    'Finished downloading %s dependency %s to %s',
-                    target_type,
-                    dependency_state.dependency_key,
-                    dependency_path,
-                )
-                with self._dependency_locks[dependency_state.dependency_key]:
-                    self._downloading[dependency_state.dependency_key]['success'] = True
 
-            except Exception as e:
-                with self._dependency_locks[dependency_state.dependency_key]:
-                    self._downloading[dependency_state.dependency_key]['success'] = False
-                    self._downloading[dependency_state.dependency_key][
-                        'failure_message'
-                    ] = "Dependency download failed: %s " % str(e)
+            attempt = 0
+            while attempt < self._download_dependencies_max_retries:
+                try:
+                    # Start async download to the fileobj
+                    fileobj, target_type = self._bundle_service.get_bundle_contents(
+                        dependency_state.dependency_key.parent_uuid,
+                        dependency_state.dependency_key.parent_path,
+                    )
+                    with closing(fileobj):
+                        # "Bug" the fileobj's read function so that we can keep
+                        # track of the number of bytes downloaded so far.
+                        old_read_method = fileobj.read
+                        bytes_downloaded = [0]
+
+                        def interruptable_read(*args, **kwargs):
+                            data = old_read_method(*args, **kwargs)
+                            bytes_downloaded[0] += len(data)
+                            update_state_and_check_killed(bytes_downloaded[0])
+                            return data
+
+                        fileobj.read = interruptable_read
+
+                        # Start copying the fileobj to filesystem dependency path
+                        self._store_dependency(dependency_path, fileobj, target_type)
+
+                    logger.debug(
+                        'Finished downloading %s dependency %s to %s',
+                        target_type,
+                        dependency_state.dependency_key,
+                        dependency_path,
+                    )
+                    with self._dependency_locks[dependency_state.dependency_key]:
+                        self._downloading[dependency_state.dependency_key]['success'] = True
+
+                except Exception as e:
+                    attempt += 1
+                    if attempt >= self._download_dependencies_max_retries:
+                        with self._dependency_locks[dependency_state.dependency_key]:
+                            self._downloading[dependency_state.dependency_key]['success'] = False
+                            self._downloading[dependency_state.dependency_key][
+                                'failure_message'
+                            ] = "Dependency download failed: %s " % str(e)
+                    else:
+                        logger.warning(
+                            f'Failed to download {dependency_state.dependency_key} after {attempt} attempt(s) '
+                            f'due to {str(e)}. Retrying up to {self._download_dependencies_max_retries} times...'
+                        )
+                else:
+                    # Break out of the retry loop if no exceptions were thrown
+                    break
 
         self._downloading.add_if_new(
             dependency_state.dependency_key, threading.Thread(target=download, args=[])
diff --git a/codalab/worker/main.py b/codalab/worker/main.py
index 1f3f0272..b2421be3 100644
--- a/codalab/worker/main.py
+++ b/codalab/worker/main.py
@@ -167,6 +167,12 @@ def parse_args():
         action='store_true',
         help="Exit the worker if it encounters an exception (rather than sleeping).",
     )
+    parser.add_argument(
+        '--download-dependencies-max-retries',
+        type=int,
+        default=3,
+        help='The number of times to retry downloading dependencies after a failure (defaults to 3).',
+    )
     return parser.parse_args()
 
 
@@ -244,6 +250,7 @@ def main():
             bundle_service,
             args.work_dir,
             args.max_work_dir_size,
+            args.download_dependencies_max_retries,
         )
     # Set up local directories
     if not os.path.exists(args.work_dir):
diff --git a/codalab/worker_manager/main.py b/codalab/worker_manager/main.py
index 1430b653..4c6d87a8 100644
--- a/codalab/worker_manager/main.py
+++ b/codalab/worker_manager/main.py
@@ -70,6 +70,12 @@ def main():
         default=10 * 60,
         type=int,
     )
+    parser.add_argument(
+        '--worker-download-dependencies-max-retries',
+        type=int,
+        default=3,
+        help='The number of times a CodaLab worker will retry downloading dependencies after a failure (defaults to 3).',
+    )
     parser.add_argument(
         '--worker-checkin-frequency-seconds',
         type=int,
diff --git a/codalab/worker_manager/worker_manager.py b/codalab/worker_manager/worker_manager.py
index a828a5f9..0f093e71 100644
--- a/codalab/worker_manager/worker_manager.py
+++ b/codalab/worker_manager/worker_manager.py
@@ -128,6 +128,13 @@ class WorkerManager(object):
             command.extend(['--group', self.args.worker_group])
         if self.args.worker_exit_after_num_runs and self.args.worker_exit_after_num_runs > 0:
             command.extend(['--exit-after-num-runs', str(self.args.worker_exit_after_num_runs)])
+        if self.args.worker_download_dependencies_max_retries:
+            command.extend(
+                [
+                    '--download-dependencies-max-retries',
+                    str(self.args.worker_download_dependencies_max_retries),
+                ]
+            )
         if self.args.worker_max_work_dir_size:
             command.extend(['--max-work-dir-size', self.args.worker_max_work_dir_size])
         if self.args.worker_delete_work_dir_on_exit:
diff --git a/codalab_service.py b/codalab_service.py
index 229a65c6..ebba1a34 100755
--- a/codalab_service.py
+++ b/codalab_service.py
@@ -279,9 +279,15 @@ CODALAB_ARGUMENTS = [
     # Worker manager
     CodalabArg(
         name='worker_manager_type',
-        help='Type of worker manager (azure-batch or aws-batch)',
+        help='Type of worker manager (azure-batch, aws-batch or kubernetes-batch)',
         default='azure-batch',
     ),
+    CodalabArg(
+        name='worker_manager_worker_download_dependencies_max_retries',
+        help='The number of times a CodaLab worker stated by the worker manager '
+        'will retry downloading dependencies after a failure (defaults to 3).',
+        default=3,
+    ),
     CodalabArg(
         name='worker_manager_worker_work_dir_prefix',
         help='Prefix to use for each worker\'s working directory of the worker manager',
diff --git a/docker_config/compose_files/docker-compose.yml b/docker_config/compose_files/docker-compose.yml
index a8538641..f54b7ba3 100644
--- a/docker_config/compose_files/docker-compose.yml
+++ b/docker_config/compose_files/docker-compose.yml
@@ -32,6 +32,7 @@ x-codalab-env: &codalab-env
   - CODALAB_LINK_MOUNTS=${CODALAB_LINK_MOUNTS}
   - CODALAB_BUNDLE_MANAGER_WORKER_TIMEOUT_SECONDS=${CODALAB_BUNDLE_MANAGER_WORKER_TIMEOUT_SECONDS}
   - CODALAB_WORKER_MANAGER_TYPE=${CODALAB_WORKER_MANAGER_TYPE}
+  - CODALAB_WORKER_MANAGER_WORKER_DOWNLOAD_DEPENDENCIES_MAX_RETRIES=${CODALAB_WORKER_MANAGER_WORKER_DOWNLOAD_DEPENDENCIES_MAX_RETRIES}
   - CODALAB_WORKER_MANAGER_WORKER_WORK_DIR_PREFIX=${CODALAB_WORKER_MANAGER_WORKER_WORK_DIR_PREFIX}
   - CODALAB_WORKER_MANAGER_WORKER_MAX_WORK_DIR_SIZE=${CODALAB_WORKER_MANAGER_WORKER_MAX_WORK_DIR_SIZE}
   - CODALAB_WORKER_MANAGER_IDLE_SECONDS=${CODALAB_WORKER_MANAGER_IDLE_SECONDS}
@@ -142,6 +143,7 @@ services:
       --server ${CODALAB_SERVER}
       --min-workers ${CODALAB_WORKER_MANAGER_MIN_CPU_WORKERS}
       --max-workers ${CODALAB_WORKER_MANAGER_MAX_CPU_WORKERS}
+      --worker-download-dependencies-max-retries ${CODALAB_WORKER_MANAGER_WORKER_DOWNLOAD_DEPENDENCIES_MAX_RETRIES}
       --worker-work-dir-prefix ${CODALAB_WORKER_MANAGER_WORKER_WORK_DIR_PREFIX}
       --worker-max-work-dir-size ${CODALAB_WORKER_MANAGER_WORKER_MAX_WORK_DIR_SIZE}
       --search request_gpus=0
@@ -167,6 +169,7 @@ services:
       --server ${CODALAB_SERVER}
       --min-workers ${CODALAB_WORKER_MANAGER_MIN_GPU_WORKERS}
       --max-workers ${CODALAB_WORKER_MANAGER_MAX_GPU_WORKERS}
+      --worker-download-dependencies-max-retries ${CODALAB_WORKER_MANAGER_WORKER_DOWNLOAD_DEPENDENCIES_MAX_RETRIES}
       --worker-work-dir-prefix ${CODALAB_WORKER_MANAGER_WORKER_WORK_DIR_PREFIX}
       --worker-max-work-dir-size ${CODALAB_WORKER_MANAGER_WORKER_MAX_WORK_DIR_SIZE}
       --search request_gpus=1
@@ -193,6 +196,7 @@ services:
       --server ${CODALAB_SERVER}
       --min-workers ${CODALAB_WORKER_MANAGER_MIN_CPU_WORKERS}
       --max-workers ${CODALAB_WORKER_MANAGER_MAX_CPU_WORKERS}
+      --worker-download-dependencies-max-retries ${CODALAB_WORKER_MANAGER_WORKER_DOWNLOAD_DEPENDENCIES_MAX_RETRIES}
       --worker-work-dir-prefix ${CODALAB_WORKER_MANAGER_WORKER_WORK_DIR_PREFIX}
       --worker-max-work-dir-size ${CODALAB_WORKER_MANAGER_WORKER_MAX_WORK_DIR_SIZE}
       --search request_gpus=0
@@ -220,6 +224,7 @@ services:
       --server ${CODALAB_SERVER}
       --min-workers ${CODALAB_WORKER_MANAGER_MIN_GPU_WORKERS}
       --max-workers ${CODALAB_WORKER_MANAGER_MAX_GPU_WORKERS}
+      --worker-download-dependencies-max-retries ${CODALAB_WORKER_MANAGER_WORKER_DOWNLOAD_DEPENDENCIES_MAX_RETRIES}
       --worker-work-dir-prefix ${CODALAB_WORKER_MANAGER_WORKER_WORK_DIR_PREFIX}
       --worker-max-work-dir-size ${CODALAB_WORKER_MANAGER_WORKER_MAX_WORK_DIR_SIZE}
       --search request_gpus=1
@@ -248,6 +253,7 @@ services:
       --server ${CODALAB_SERVER}
       --min-workers ${CODALAB_WORKER_MANAGER_MIN_CPU_WORKERS}
       --max-workers ${CODALAB_WORKER_MANAGER_MAX_CPU_WORKERS}
+      --worker-download-dependencies-max-retries ${CODALAB_WORKER_MANAGER_WORKER_DOWNLOAD_DEPENDENCIES_MAX_RETRIES}
       --worker-work-dir-prefix ${CODALAB_WORKER_MANAGER_WORKER_WORK_DIR_PREFIX}
       --worker-max-work-dir-size ${CODALAB_WORKER_MANAGER_WORKER_MAX_WORK_DIR_SIZE}
       --search request_gpus=0
@@ -277,6 +283,7 @@ services:
       --server ${CODALAB_SERVER}
       --min-workers ${CODALAB_WORKER_MANAGER_MIN_GPU_WORKERS}
       --max-workers ${CODALAB_WORKER_MANAGER_MAX_GPU_WORKERS}
+      --worker-download-dependencies-max-retries ${CODALAB_WORKER_MANAGER_WORKER_DOWNLOAD_DEPENDENCIES_MAX_RETRIES}
       --worker-work-dir-prefix ${CODALAB_WORKER_MANAGER_WORKER_WORK_DIR_PREFIX}
       --worker-max-work-dir-size ${CODALAB_WORKER_MANAGER_WORKER_MAX_WORK_DIR_SIZE}
       --search request_gpus=1


Test patch:
diff --git a/tests/unit/worker_manager/slurm_batch_worker_manager_test.py b/tests/unit/worker_manager/slurm_batch_worker_manager_test.py
index 644f69b5..191c4623 100644
--- a/tests/unit/worker_manager/slurm_batch_worker_manager_test.py
+++ b/tests/unit/worker_manager/slurm_batch_worker_manager_test.py
@@ -15,6 +15,7 @@ class SlurmBatchWorkerManagerTest(unittest.TestCase):
             partition='some_partition',
             worker_executable='cl-worker',
             worker_idle_seconds='888',
+            worker_download_dependencies_max_retries=5,
             worker_tag='some_tag',
             worker_group='some_group',
             worker_exit_after_num_runs=8,
@@ -40,8 +41,8 @@ class SlurmBatchWorkerManagerTest(unittest.TestCase):
             "cl-worker --server some_server --verbose --exit-when-idle --idle-seconds 888 "
             "--work-dir /some/path/some_user-codalab-SlurmBatchWorkerManager-scratch/some_worker_id "
             "--id $(hostname -s)-some_worker_id --network-prefix cl_worker_some_worker_id_network --tag some_tag "
-            "--group some_group --exit-after-num-runs 8 --max-work-dir-size 88g --checkin-frequency-seconds 30 "
-            "--pass-down-termination"
+            "--group some_group --exit-after-num-runs 8 --download-dependencies-max-retries 5"
+            " --max-work-dir-size 88g --checkin-frequency-seconds 30 --pass-down-termination"
         )
         self.assertEqual(' '.join(command), expected_command_str)
 
@@ -53,6 +54,7 @@ class SlurmBatchWorkerManagerTest(unittest.TestCase):
             partition='some_partition',
             worker_executable='cl-worker',
             worker_idle_seconds='888',
+            worker_download_dependencies_max_retries=5,
             worker_tag='some_tag',
             worker_group='some_group',
             worker_exit_after_num_runs=8,