Tigru/fix/experiments killing each other (#159)

* reclaiming rescheduled experiments is now preferred; rescheduled experiments may not be claimed by other jobs anymore

* added filter on rescheduled exps in cancel_empty_pending_jobs

* added comments and added RESCHEDULED to detect_killed

Author: Grutschus

src/seml/commands/manage.py

@@ -87,12 +87,27 @@ def cancel_empty_pending_jobs(db_collection_name: str, *sacred_ids: int):
         # There are still pending experiments, we don't want to cancel the jobs.
         return
     pending_exps = list(collection.find({'_id': {'$in': sacred_ids}}, {'slurm'}))
+    rescheduled_exps = list(
+        collection.find(
+            {'_id': {'$in': sacred_ids}, 'status': {'$in': States.RESCHEDULED}},
+            {'execution'},
+        )
+    )
+
+    # We exclude all SLURM jobs that are needed for continuing a RESCHEDULED experiment
+    rescheduled_ids = []
+    for exp in rescheduled_exps:
+        execution = exp['execution']
+        if 'array_id' in execution:
+            rescheduled_ids.append(execution['array_id'])
+
     array_ids = {
         conf['array_id']
         for exp in pending_exps
         for conf in exp['slurm']
         if 'array_id' in conf
     }
+    array_ids = array_ids.difference(rescheduled_ids)
     # Only cancel the pending jobs
     cancel_slurm_jobs(*array_ids, state=SETTINGS.SLURM_STATES.PENDING[0])
 
@@ -659,7 +674,7 @@ def detect_killed(db_collection_name: str, print_detected: bool = True):
     cluster = get_cluster_name()
     exps = collection.find(
         {
-            'status': {'$in': [*States.PENDING, *States.RUNNING]},
+            'status': {'$in': [*States.PENDING, *States.RUNNING, *States.RESCHEDULED]},
             'execution.cluster': cluster,  # only check experiments that are running on the current cluster
             # Previously we only checked for started experiments by including the following line:
             # 'host': {'$exists': True},  # only check experiments that have been started

src/seml/commands/start.py

@@ -1118,6 +1118,7 @@ def claim_experiment(db_collection_name: str, exp_ids: Sequence[int]):
     """
     collection = get_collection(db_collection_name)
     array_id, task_id = get_current_slurm_array_id()
+    exp = None
     if array_id is not None and task_id is not None:
         # We are running in slurm
         array_id, task_id = int(array_id), int(task_id)
@@ -1127,14 +1128,30 @@ def claim_experiment(db_collection_name: str, exp_ids: Sequence[int]):
             'execution.task_id': task_id,
             'execution.cluster': cluster_name,
         }
+        # First, we check whether this SLURM job is responsible for a RESCHEDULED experiment
+        # If so, we claim this first
         exp = collection.find_one_and_update(
             {
                 '_id': {'$in': list(exp_ids)},
-                'status': {'$in': States.PENDING + States.RESCHEDULED},
+                'status': {'$in': States.RESCHEDULED},
+                'execution.array_id': array_id,
+                'execution.task_id': task_id,
+                'execution.cluster': cluster_name,
             },
             {'$set': {'status': States.RUNNING[0], **update}},
             {'_id': 1, 'slurm': 1},
         )
+        # Only after we have checked that this job is not responsible for a RESCHEDULED experiment
+        # can we pick up a pending one.
+        if exp is None:
+            exp = collection.find_one_and_update(
+                {
+                    '_id': {'$in': list(exp_ids)},
+                    'status': {'$in': States.PENDING},
+                },
+                {'$set': {'status': States.RUNNING[0], **update}},
+                {'_id': 1, 'slurm': 1},
+            )
         if exp is None:
             exit(3)
         # Set slurm output file