Allow resubmit of failed jobs while others are running

Author: tien-tong

babs/interaction.py

@@ -1,6 +1,7 @@
 """This is the main module."""
 
 import numpy as np
+import pandas as pd
 
 from babs.base import BABS
 from babs.scheduler import (
@@ -30,20 +31,16 @@ def babs_submit(self, count=None, submit_df=None, skip_failed=False):
             default: None
         """
 
-        # Check if there are still jobs running
+        # Check if there are still jobs running/pending
         currently_running_df = self.get_currently_running_jobs_df()
-        if currently_running_df.shape[0] > 0:
-            non_cg_states = (
-                currently_running_df['state'].fillna('').ne('CG')
-                if 'state' in currently_running_df
-                else np.array([True] * currently_running_df.shape[0])
-            )
-            if non_cg_states.any():
-                raise Exception(
-                    'There are still jobs running. Please wait for them to finish or cancel them.'
-                    f' Current running jobs:\n{currently_running_df}'
-                )
-            print('All currently running jobs are in CG state; proceeding with submission.')
+        if currently_running_df.empty:
+            running_pending_df = currently_running_df
+        elif 'state' in currently_running_df:
+            running_pending_df = currently_running_df[
+                currently_running_df['state'].fillna('').isin(['PD', 'R'])
+            ]
+        else:
+            running_pending_df = currently_running_df
 
         # Find the rows that don't have results yet
         status_df = self.get_job_status_df()
@@ -54,6 +51,36 @@ def babs_submit(self, count=None, submit_df=None, skip_failed=False):
         if submit_df is not None:
             df_needs_submit = submit_df
 
+        # Remove any jobs that are still running/pending
+        if not running_pending_df.empty:
+            merge_cols = ['sub_id']
+            if 'ses_id' in df_needs_submit and 'ses_id' in running_pending_df:
+                merge_cols = ['sub_id', 'ses_id']
+            if all(col in running_pending_df for col in merge_cols):
+                skipped_df = pd.merge(
+                    df_needs_submit,
+                    running_pending_df,
+                    on=merge_cols,
+                    how='inner',
+                    suffixes=('', '_running'),
+                )
+                df_needs_submit = df_needs_submit.merge(
+                    running_pending_df[merge_cols].drop_duplicates(),
+                    on=merge_cols,
+                    how='left',
+                    indicator=True,
+                )
+                df_needs_submit = df_needs_submit[df_needs_submit['_merge'] == 'left_only']
+                df_needs_submit = df_needs_submit.drop(columns=['_merge'])
+                if not skipped_df.empty:
+                    if 'job_id_running' in skipped_df:
+                        job_ids = sorted(
+                            pd.Series(skipped_df['job_id_running']).dropna().unique().tolist()
+                        )
+                        print(f'Skipping running/pending jobs. Job IDs: {job_ids}')
+                    else:
+                        print('Skipping running/pending jobs without job IDs.')
+
         # only run `babs submit` when there are subjects/sessions not yet submitted
         if df_needs_submit.empty:
             print('No jobs to submit')

tests/test_interaction.py

@@ -1,7 +1,6 @@
 """Tests for interaction behaviors."""
 
 import pandas as pd
-import pytest
 
 from babs.interaction import BABSInteraction
 
@@ -26,7 +25,7 @@ def _minimal_status_df():
     )
 
 
-def test_babs_submit_blocks_non_cg_jobs(babs_project_subjectlevel, monkeypatch):
+def test_babs_submit_skips_running_jobs(babs_project_subjectlevel, monkeypatch, capsys):
     babs_proj = BABSInteraction(project_root=babs_project_subjectlevel)
     running_df = pd.DataFrame(
         {
@@ -39,12 +38,25 @@ def test_babs_submit_blocks_non_cg_jobs(babs_project_subjectlevel, monkeypatch):
             'cpus': [1],
             'partition': ['normal'],
             'name': ['test_array_job'],
+            'sub_id': ['sub-01'],
         }
     )
     monkeypatch.setattr(babs_proj, 'get_currently_running_jobs_df', lambda: running_df)
+    monkeypatch.setattr(babs_proj, 'get_job_status_df', _minimal_status_df)
+
+    submit_calls = []
+
+    def _mock_submit_array(analysis_path, queue, total_jobs):
+        submit_calls.append((analysis_path, queue, total_jobs))
+        return 123
+
+    monkeypatch.setattr('babs.interaction.submit_array', _mock_submit_array)
+
+    babs_proj.babs_submit(count=2)
 
-    with pytest.raises(Exception, match='There are still jobs running'):
-        babs_proj.babs_submit(count=1)
+    assert submit_calls[0][2] == 1
+    out = capsys.readouterr().out
+    assert 'Job IDs: [1]' in out
 
 
 def test_babs_submit_allows_cg_jobs(babs_project_subjectlevel, monkeypatch):