Support bundle checkpoint / preemptible workers

codalab/bundles/run_bundle.py

@@ -69,6 +69,10 @@ class RunBundle(DerivedBundle):
     METADATA_SPECS.append(MetadataSpec('exitcode', int, 'Exitcode of the process.', generated=True))
     METADATA_SPECS.append(MetadataSpec('job_handle', str, 'Identifies the job handle (internal).', generated=True, hide_when_anonymous=True))
     METADATA_SPECS.append(MetadataSpec('remote', str, 'Where this job is/was run (internal).', generated=True, hide_when_anonymous=True))
+    METADATA_SPECS.append(MetadataSpec('remote_history', list, 'All workers where this job has been run (internal); multiple values indicate'
+                                       'that the bundle was preempted and moved to a different worker.',
+                                       generated=True, hide_when_anonymous=True))
+    METADATA_SPECS.append(MetadataSpec('preemptible', bool, 'Whether the bundle is currently running / finished on a preemptible worker.', generated=True, hide_when_anonymous=True, default=False))
     # fmt: on
 
     @classmethod

codalab/model/bundle_model.py

@@ -915,7 +915,15 @@ def transition_bundle_preparing(self, bundle, user_id, worker_id, start_time, re
 
             bundle_update = {
                 'state': State.PREPARING,
-                'metadata': {'started': start_time, 'last_updated': start_time, 'remote': remote},
+                'metadata': {
+                    'started': start_time,
+                    'last_updated': start_time,
+                    'remote': remote,
+                    'remote_history': getattr(bundle.metadata, "remote_history", [])
+                    + [
+                        remote
+                    ],  # Store the history of which workers ran this bundle before in the bundle metadata.
+                },
             }
             self.update_bundle(bundle, bundle_update, connection)
 
@@ -987,6 +995,8 @@ def transition_bundle_worker_offline(self, bundle):
         Transitions bundle to WORKER_OFFLINE state:
             Updates the last_updated metadata.
             Removes the corresponding row from worker_run if it exists.
+
+        If the bundle is preemptible, move the bundle to the STAGED state instead.
         """
         with self.engine.begin() as connection:
             # Check that it still exists and is running
@@ -1000,15 +1010,20 @@ def transition_bundle_worker_offline(self, bundle):
                 # The user deleted the bundle or the bundle finished
                 return False
 
+            if getattr(bundle.metadata, "preemptible", False):
+                bundle_update = {
+                    'state': State.STAGED,
+                    'metadata': {'last_updated': int(time.time())},
+                }
+            else:
+                bundle_update = {
+                    'state': State.WORKER_OFFLINE,
+                    'metadata': {'last_updated': int(time.time())},
+                }
             # Delete row in worker_run
             connection.execute(
                 cl_worker_run.delete().where(cl_worker_run.c.run_uuid == bundle.uuid)
             )
-
-            bundle_update = {
-                'state': State.WORKER_OFFLINE,
-                'metadata': {'last_updated': int(time.time())},
-            }
             self.update_bundle(bundle, bundle_update, connection)
         return True
 

codalab/model/tables.py

@@ -526,6 +526,7 @@
         'exit_after_num_runs', Integer, nullable=False
     ),  # Number of jobs allowed to run on worker.
     Column('is_terminating', Boolean, nullable=False),
+    Column('preemptible', Boolean, nullable=False),
     mysql_charset=TABLE_DEFAULT_CHARSET,
 )
 

codalab/model/worker_model.py

@@ -52,6 +52,7 @@ def worker_checkin(
         tag_exclusive,
         exit_after_num_runs,
         is_terminating,
+        preemptible,
     ):
         """
         Adds the worker to the database, if not yet there. Returns the socket ID
@@ -69,6 +70,7 @@ def worker_checkin(
                 'tag_exclusive': tag_exclusive,
                 'exit_after_num_runs': exit_after_num_runs,
                 'is_terminating': is_terminating,
+                'preemptible': preemptible,
             }
 
             # Populate the group for this worker, if group_name is valid
@@ -206,6 +208,7 @@ def get_workers(self):
                 'tag_exclusive': row.tag_exclusive,
                 'exit_after_num_runs': row.exit_after_num_runs,
                 'is_terminating': row.is_terminating,
+                'preemptible': row.preemptible,
             }
             for row in worker_rows
         }

codalab/worker/main.py

@@ -190,6 +190,9 @@ def parse_args():
         default=1,
         help='The shared memory size of the run container in GB (defaults to 1).',
     )
+    parser.add_argument(
+        '--preemptible', action='store_true', help='Whether the worker is preemptible.',
+    )
     return parser.parse_args()
 
 
@@ -325,6 +328,7 @@ def main():
         delete_work_dir_on_exit=args.delete_work_dir_on_exit,
         exit_on_exception=args.exit_on_exception,
         shared_memory_size_gb=args.shared_memory_size_gb,
+        preemptible=args.preemptible,
     )
 
     # Register a signal handler to ensure safe shutdown.

codalab/worker/worker.py

@@ -79,6 +79,7 @@ def __init__(
         # A flag indicating if the worker will exit if it encounters an exception
         exit_on_exception=False,  # type: bool
         shared_memory_size_gb=1,  # type: int
+        preemptible=False,  # type: bool
     ):
         self.image_manager = image_manager
         self.dependency_manager = dependency_manager
@@ -114,6 +115,7 @@ def __init__(
         self.terminate_and_restage = False
         self.pass_down_termination = pass_down_termination
         self.exit_on_exception = exit_on_exception
+        self.preemptible = preemptible
 
         self.checkin_frequency_seconds = checkin_frequency_seconds
         self.last_checkin_successful = False
@@ -422,6 +424,7 @@ def checkin(self):
             'tag_exclusive': self.tag_exclusive,
             'exit_after_num_runs': self.exit_after_num_runs - self.num_runs,
             'is_terminating': self.terminate or self.terminate_and_restage,
+            'preemptible': self.preemptible,
         }
         try:
             response = self.bundle_service.checkin(self.id, request)

codalab/worker_manager/main.py

@@ -114,6 +114,11 @@ def main():
         type=int,
         help="The shared memory size in GB of the run container started by the CodaLab Workers.",
     )
+    parser.add_argument(
+        '--worker-preemptible',
+        action='store_true',
+        help='Whether the CodaLab workers are preemptible.',
+    )
     subparsers = parser.add_subparsers(
         title='Worker Manager to run',
         description='Which worker manager to run (AWS Batch etc.)',

codalab/worker_manager/worker_manager.py

@@ -151,6 +151,8 @@ def build_command(self, worker_id: str, work_dir: str) -> List[str]:
             )
         if self.args.worker_shared_memory_size_gb:
             command.extend(['--shared-memory-size-gb', str(self.args.worker_shared_memory_size_gb)])
+        if self.args.worker_preemptible:
+            command.extend(['--worker-preemptible'])
 
         return command
 

docs/Checkpoints.md

@@ -0,0 +1,14 @@
+# Checkpoints
+
+!!! warning "Unreleased feature"
+        This feature is still in development. It does not fully work yet and will only work once [#3710](https://github.com/codalab/codalab-worksheets/issues/3710) is resolved.
+
+CodaLab supports the concept of bundle checkpoints, which allow bundles to be stopped and then resumed on different workers. This is especially useful if you want to use preemptible compute nodes, such as [EC2 Spot](https://aws.amazon.com/ec2/spot/) instances. When a worker is preempted, the bundle should be resumed on another worker.
+
+## Workflow steps
+
+This section describes a sample workflow for running bundles on a pool of preemptible workers. The workflow uses the tag `preemptiblepool1` to organize compute, but you can change this tag name as needed.
+
+1. Start a worker manager that runs workers on preemptible compute notes, using the following flags: `--worker-preemptible --worker-tag-exclusive --worker-tag="preemptiblepool1"`. Ensure that all the workers started by the worker manager are configured to share the same network disk for their bundle run working directory / dependency cache.
+1. User can run a bundle by running `cl run --tag="preemptiblepool1"`.
+1. The bundle manager will assign the bundle to one of the preemptible workers with tag `preemptiblepool1`. If the worker gets preempted, the bundle transitions back to the `STAGED` state and will then be assigned to another preemptible worker with tag `preemptiblepool1`. The history of workers where a bundle has run is stored in the `remote_history` bundle metadata field.

mkdocs.yml

@@ -37,6 +37,7 @@ nav:
     - Execution: Execution.md
     - Competitions: Competitions.md
     - Storage Externalization: Externalization.md
+    - Checkpoints: Checkpoints.md
   - Reference:
     - CLI Reference: CLI-Reference.md
     - REST API Reference: REST-API-Reference.md

tests/unit/rest/workers_test.py

@@ -17,6 +17,7 @@ def test_checkin(self):
             'tag_exclusive': False,
             'exit_after_num_runs': 999999999,
             'is_terminating': False,
+            'preemptible': True,
         }
         response = self.app.post_json('/rest/workers/test_worker/checkin', body)
         self.assertEqual(response.status_int, 200)

tests/unit/server/bundle_manager/__init__.py

@@ -244,6 +244,7 @@ def mock_worker_checkin(
             tag_exclusive=False,
             exit_after_num_runs=999999999,
             is_terminating=False,
+            preemptible=False,
         )
         # Mock a reply from the worker
         self.bundle_manager._worker_model.send_json_message = Mock(return_value=True)

tests/unit/server/bundle_manager/schedule_run_bundles_test.py

@@ -71,6 +71,18 @@ def test_finalizing_bundle_goes_offline_if_no_worker_claims(self):
         bundle = self.bundle_manager._model.get_bundle(bundle.uuid)
         self.assertEqual(bundle.state, State.WORKER_OFFLINE)
 
+    def test_reassign_stuck_running_preemptible_bundles(self):
+        """If no workers exist to claim a bundle, and the bundle is running on a preemptible worker, it should go to the STAGED state in preparation for being reassigned to another worker."""
+        bundle = self.create_run_bundle(
+            State.RUNNING, {"preemptible": True, "remote_history": ["remote1"], "remote": "remote1"}
+        )
+        self.save_bundle(bundle)
+        self.bundle_manager._schedule_run_bundles()
+
+        bundle = self.bundle_manager._model.get_bundle(bundle.uuid)
+        self.assertEqual(bundle.state, State.STAGED)
+        self.assertEqual(bundle.metadata.remote_history, ["remote1"])
+
     def test_finalizing_bundle_gets_finished(self):
         """If a worker checks in with a "finalizing" message, the bundle should transition
         to the FINALIZING and then FINISHED state."""