Support bundle checkpoint / preemptible workers

codalab/bundles/run_bundle.py

@@ -69,7 +69,9 @@ 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('remotes', list, 'Previous workers where this job was 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('remotes', list, 'Previous workers where this job was 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,
+                    'remotes': (bundle.metadata.remotes or [])
+                    + [
+                        remote
+                    ],  # Store the history of which workers ran this bundle before in the bundle metadata.
+                },
             }
             self.update_bundle(bundle, bundle_update, connection)
 
@@ -987,7 +995,7 @@ 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 worker is preemptible, move the bundle to the STAGED state instead.
         """
         with self.engine.begin() as connection:
@@ -1002,14 +1010,11 @@ def transition_bundle_worker_offline(self, bundle):
                 # The user deleted the bundle or the bundle finished
                 return False
 
-            worker = self.get_bundle_worker(bundle.uuid)
-            if worker['preemptible']:
+            if bundle.metadata.preemptible:
                 bundle_update = {
                     'state': State.STAGED,
                     'metadata': {'last_updated': int(time.time())},
-                    'remotes': worker.id
                 }
-                # TODO: we should store the history of which workers ran this bundle before in the bundle metadata.
             else:
                 bundle_update = {
                     'state': State.WORKER_OFFLINE,