Repository: codalab/codalab-worksheets

Issue:
Make worker check-in frequency and worker timeout configurable
**Describe the bug**
- Add a delay for worker check-in, which can be specified when creating a CodaLab worker
- Make `WORKER_TIMEOUT` in `bundle_manager.py` configurable 


Discussion:

Patch:
diff --git a/codalab/bin/bundle_manager.py b/codalab/bin/bundle_manager.py
index 59cfc4d1..108ac56f 100644
--- a/codalab/bin/bundle_manager.py
+++ b/codalab/bin/bundle_manager.py
@@ -13,9 +13,15 @@ def main():
         type=int,
         default=0.5,
     )
+    parser.add_argument(
+        '--worker-timeout-seconds',
+        help='Number of seconds to wait after a worker check-in before determining a worker is offline',
+        type=int,
+        default=60,
+    )
     args = parser.parse_args()
 
-    manager = BundleManager(CodaLabManager())
+    manager = BundleManager(CodaLabManager(), args.worker_timeout_seconds)
     # Register a signal handler to ensure safe shutdown.
     for sig in [signal.SIGTERM, signal.SIGINT, signal.SIGHUP]:
         signal.signal(sig, lambda signup, frame: manager.signal())
diff --git a/codalab/server/bundle_manager.py b/codalab/server/bundle_manager.py
index 1860aa48..e67d0543 100644
--- a/codalab/server/bundle_manager.py
+++ b/codalab/server/bundle_manager.py
@@ -24,7 +24,6 @@ from codalab.worker.bundle_state import State, RunResources
 
 logger = logging.getLogger(__name__)
 
-WORKER_TIMEOUT_SECONDS = 60
 SECONDS_PER_DAY = 60 * 60 * 24
 # Fail unresponsive bundles in uploading, staged and running state after this many days.
 BUNDLE_TIMEOUT_DAYS = 60
@@ -43,7 +42,7 @@ class BundleManager(object):
     Assigns run bundles to workers and makes make bundles.
     """
 
-    def __init__(self, codalab_manager):
+    def __init__(self, codalab_manager, worker_timeout_seconds=60):
         config = codalab_manager.config.get('workers')
         if not config:
             print('config.json file missing a workers section.', file=sys.stderr)
@@ -65,6 +64,7 @@ class BundleManager(object):
         def parse(to_value, field):
             return to_value(config[field]) if field in config else None
 
+        self._worker_timeout_seconds = worker_timeout_seconds
         self._max_request_time = parse(formatting.parse_duration, 'max_request_time') or 0
         self._max_request_memory = parse(formatting.parse_size, 'max_request_memory') or 0
         self._min_request_memory = (
@@ -263,7 +263,7 @@ class BundleManager(object):
         """
         for worker in workers.workers():
             if datetime.datetime.utcnow() - worker['checkin_time'] > datetime.timedelta(
-                seconds=WORKER_TIMEOUT_SECONDS
+                seconds=self._worker_timeout_seconds
             ):
                 logger.info(
                     'Cleaning up dead worker (%s, %s)', worker['user_id'], worker['worker_id']
@@ -322,7 +322,7 @@ class BundleManager(object):
             failure_message = None
             if not workers.is_running(bundle.uuid):
                 failure_message = 'No worker claims bundle'
-            if now - bundle.metadata.last_updated > WORKER_TIMEOUT_SECONDS:
+            if now - bundle.metadata.last_updated > self._worker_timeout_seconds:
                 failure_message = 'Worker offline'
             if failure_message is not None:
                 logger.info('Bringing bundle offline %s: %s', bundle.uuid, failure_message)
@@ -773,7 +773,9 @@ class BundleManager(object):
         READY / FAILED, no worker_run DB entry:
             Finished.
         """
-        workers = WorkerInfoAccessor(self._model, self._worker_model, WORKER_TIMEOUT_SECONDS - 5)
+        workers = WorkerInfoAccessor(
+            self._model, self._worker_model, self._worker_timeout_seconds - 5
+        )
 
         # Handle some exceptional cases.
         self._cleanup_dead_workers(workers)
diff --git a/codalab/worker/main.py b/codalab/worker/main.py
index acce7940..3ba18444 100644
--- a/codalab/worker/main.py
+++ b/codalab/worker/main.py
@@ -120,6 +120,12 @@ def parse_args():
         type=int,
         default=0,
     )
+    parser.add_argument(
+        '--checkin-frequency-seconds',
+        help='Number of seconds to wait between worker check-ins',
+        type=int,
+        default=5,
+    )
     parser.add_argument(
         '--id',
         default='%s(%d)' % (socket.gethostname(), os.getpid()),
@@ -259,6 +265,7 @@ def main():
         args.exit_when_idle,
         args.exit_after_num_runs,
         args.idle_seconds,
+        args.checkin_frequency_seconds,
         bundle_service,
         args.shared_file_system,
         args.tag_exclusive,
diff --git a/codalab/worker/worker.py b/codalab/worker/worker.py
index d3224c2f..6757d539 100644
--- a/codalab/worker/worker.py
+++ b/codalab/worker/worker.py
@@ -65,6 +65,7 @@ class Worker:
         exit_when_idle,  # type: str
         exit_after_num_runs,  # type: int
         idle_seconds,  # type: int
+        checkin_frequency_seconds,  # type: int
         bundle_service,  # type: BundleServiceClient
         shared_file_system,  # type: bool
         tag_exclusive,  # type: bool
@@ -113,6 +114,7 @@ class Worker:
         self.pass_down_termination = pass_down_termination
         self.exit_on_exception = exit_on_exception
 
+        self.checkin_frequency_seconds = checkin_frequency_seconds
         self.last_checkin_successful = False
         self.last_time_ran = None  # type: Optional[bool]
 
@@ -234,8 +236,18 @@ class Worker:
         now = time.time()
         if len(self.runs) > 0 or self.last_time_ran is None:
             self.last_time_ran = now
-        is_idle = now - self.last_time_ran > self.idle_seconds
-        return self.exit_when_idle and is_idle and self.last_checkin_successful
+
+        idle_duration_seconds = now - self.last_time_ran
+        if (
+            self.exit_when_idle
+            and idle_duration_seconds > self.idle_seconds
+            and self.last_checkin_successful
+        ):
+            logger.warning(
+                "Worker was idle for {} seconds. Exiting...".format(idle_duration_seconds)
+            )
+            return True
+        return False
 
     def check_num_runs_stop(self):
         """
@@ -261,6 +273,8 @@ class Worker:
                 self.save_state()
                 if self.check_idle_stop() or self.check_num_runs_stop():
                     self.terminate = True
+                else:
+                    time.sleep(self.checkin_frequency_seconds)
             except Exception:
                 self.last_checkin_successful = False
                 if using_sentry():
diff --git a/codalab/worker_manager/main.py b/codalab/worker_manager/main.py
index 77ccbc35..780ffd3a 100644
--- a/codalab/worker_manager/main.py
+++ b/codalab/worker_manager/main.py
@@ -59,6 +59,11 @@ def main():
         default=1 * 60,
         type=int,
     )
+    parser.add_argument(
+        '--worker-checkin-frequency-seconds',
+        type=int,
+        help='If specified, the CodaLab worker will wait this many seconds between check-ins',
+    )
     parser.add_argument(
         '--worker-exit-after-num-runs',
         type=int,
diff --git a/codalab/worker_manager/worker_manager.py b/codalab/worker_manager/worker_manager.py
index 8a98b66b..8bf4be49 100644
--- a/codalab/worker_manager/worker_manager.py
+++ b/codalab/worker_manager/worker_manager.py
@@ -139,6 +139,10 @@ class WorkerManager(object):
             command.extend(['--tag-exclusive'])
         if self.args.worker_pass_down_termination:
             command.extend(['--pass-down-termination'])
+        if self.args.worker_checkin_frequency_seconds:
+            command.extend(
+                ['--checkin-frequency-seconds', str(self.args.worker_checkin_frequency_seconds)]
+            )
 
         return command
 
diff --git a/codalab_service.py b/codalab_service.py
index 4fc6e83e..c4fa2485 100755
--- a/codalab_service.py
+++ b/codalab_service.py
@@ -259,6 +259,13 @@ CODALAB_ARGUMENTS = [
             'Ingest URL for logging exceptions with Sentry. If not provided, Sentry is not used.'
         ),
     ),
+    # Bundle Manager
+    CodalabArg(
+        name='bundle_manager_worker_timeout_seconds',
+        help='Number of seconds to wait after a worker check-in before determining a worker is offline',
+        type=int,
+        default=60,
+    ),
     # Worker manager
     CodalabArg(
         name='worker_manager_type',
diff --git a/docker_config/compose_files/docker-compose.yml b/docker_config/compose_files/docker-compose.yml
index 16ae9ff7..a32cab18 100644
--- a/docker_config/compose_files/docker-compose.yml
+++ b/docker_config/compose_files/docker-compose.yml
@@ -30,6 +30,7 @@ x-codalab-env: &codalab-env
   - CODALAB_SERVER=${CODALAB_SERVER}
   - CODALAB_SHARED_FILE_SYSTEM=${CODALAB_SHARED_FILE_SYSTEM}
   - 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_IDLE_SECONDS=${CODALAB_WORKER_MANAGER_IDLE_SECONDS}
   - CODALAB_WORKER_MANAGER_SECONDS_BETWEEN_WORKERS=${CODALAB_WORKER_MANAGER_SECONDS_BETWEEN_WORKERS}
@@ -99,7 +100,9 @@ services:
 
   bundle-manager:
     image: codalab/server:${CODALAB_VERSION}
-    command: cl-bundle-manager
+    command: |
+      cl-bundle-manager
+      --worker-timeout-seconds ${CODALAB_BUNDLE_MANAGER_WORKER_TIMEOUT_SECONDS}
     <<: *codalab-base
     <<: *codalab-server
     depends_on:


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 ebc2f655..38e3d6ea 100644
--- a/tests/unit/worker_manager/slurm_batch_worker_manager_test.py
+++ b/tests/unit/worker_manager/slurm_batch_worker_manager_test.py
@@ -23,6 +23,7 @@ class SlurmBatchWorkerManagerTest(unittest.TestCase):
             worker_exit_on_exception=False,
             worker_tag_exclusive=False,
             worker_pass_down_termination=False,
+            worker_checkin_frequency_seconds=30,
             password_file=None,
             slurm_work_dir=None,
             exit_after_num_failed=None,
@@ -34,11 +35,12 @@ class SlurmBatchWorkerManagerTest(unittest.TestCase):
         # --pass-down-termination should always be set for Slurm worker managers
         self.assertTrue('--pass-down-termination' in command)
 
-        expected_command_str: str = (
-            '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 --pass-down-termination'
+        expected_command_str = (
+            "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"
         )
         self.assertEqual(' '.join(command), expected_command_str)