allow resubmit failed, non-running jobs

tien-tong · tien-tong · commit b5ae1b1a6228 · 2026-02-08T02:24:09.000-05:00
diff --git a/babs/base.py b/babs/base.py
@@ -22,6 +22,7 @@
     read_yaml,
     results_branch_dataframe,
     results_status_columns,
+    scheduler_status_columns,
     status_dtypes,
     update_job_batch_status,
     update_results_status,
@@ -444,15 +445,51 @@ def get_currently_running_jobs_df(self):
         Index: []
 
         """
+
+        def _empty_running():
+            cols = scheduler_status_columns + ['sub_id']
+            if self.processing_level == 'session':
+                cols = cols + ['ses_id']
+            return pd.DataFrame(columns=cols)
+
+        job_status_df = self.get_job_status_df()
         last_submitted_jobs_df = self.get_latest_submitted_jobs_df()
-        if last_submitted_jobs_df.empty:
-            return EMPTY_JOB_SUBMIT_DF
-        job_ids = last_submitted_jobs_df['job_id'].unique()
-        if not len(job_ids) == 1:
-            raise Exception(f'Expected 1 job id, got {len(job_ids)}')
-        job_id = job_ids[0]
-        currently_running_df = request_all_job_status(self.queue, job_id)
-        return identify_running_jobs(last_submitted_jobs_df, currently_running_df)
+
+        # Rows that are submitted but don't have results yet (candidates for "running")
+        if not job_status_df.empty:
+            sub = job_status_df['submitted'].fillna(False)
+            no_res = ~job_status_df['has_results'].fillna(False)
+            job_status_df = job_status_df.loc[sub & no_res].copy()
+
+        # Use status rows (submitted, no results) or last submit file for job_id -> sub/ses
+        mapping_df = job_status_df if not job_status_df.empty else last_submitted_jobs_df.copy()
+        if mapping_df.empty:
+            return _empty_running()
+
+        # Keep only columns needed to join scheduler output with subject/session
+        mapping_cols = ['job_id', 'task_id', 'sub_id']
+        if 'ses_id' in mapping_df:
+            mapping_cols.append('ses_id')
+        mapping_df = mapping_df[mapping_cols].copy()
+        # Drop rows with missing or invalid job/task ids so we only query real jobs
+        mapping_df = mapping_df[
+            mapping_df['job_id'].notna()
+            & mapping_df['task_id'].notna()
+            & (mapping_df['job_id'] > 0)
+            & (mapping_df['task_id'] > 0)
+        ]
+        if mapping_df.empty:
+            return _empty_running()
+
+        # Ask scheduler for each distinct job_id, keep only non-empty responses
+        job_ids = sorted({int(j) for j in mapping_df['job_id'].unique()})
+        running_dfs = [request_all_job_status(self.queue, j) for j in job_ids]
+        running_dfs = [d for d in running_dfs if not d.empty]
+        if not running_dfs:
+            return _empty_running()
+
+        # Attach sub_id (and ses_id) to scheduler rows and return
+        return identify_running_jobs(mapping_df, pd.concat(running_dfs, ignore_index=True))
 
     def get_job_status_df(self):
         """
diff --git a/babs/cli.py b/babs/cli.py
@@ -351,6 +351,14 @@ def _parse_submit():
         ' If this flag is specified, it will override the `--select` flag.',
         type=PathExists,
     )
+    parser.add_argument(
+        '--skip-running-jobs',
+        action='store_true',
+        help=(
+            'Allow submission when there are running/pending jobs by skipping '
+            'those jobs instead of raising errrors.'
+        ),
+    )
 
     return parser
 
@@ -376,6 +384,7 @@ def babs_submit_main(
     count: int | None,
     select: list | None,
     inclusion_file: Path | None,
+    skip_running_jobs: bool = False,
 ):
     """This is the core function of ``babs submit``.
 
@@ -389,6 +398,8 @@ def babs_submit_main(
         list of subject IDs and session IDs to be submitted.
     inclusion_file: Path
         path to a CSV file that lists the subjects (and sessions) to analyze.
+    skip_running_jobs: bool
+        whether to allow submission when there are running/pending jobs
     """
     import pandas as pd
 
@@ -406,7 +417,11 @@ def babs_submit_main(
     else:
         df_job_specified = None
 
-    babs_proj.babs_submit(count=count, submit_df=df_job_specified)
+    babs_proj.babs_submit(
+        count=count,
+        submit_df=df_job_specified,
+        skip_running_jobs=skip_running_jobs,
+    )
 
 
 def _parse_status():
diff --git a/babs/interaction.py b/babs/interaction.py
@@ -15,7 +15,7 @@
 class BABSInteraction(BABS):
     """Implement interactions with a BABS project - submitting jobs and checking status."""
 
-    def babs_submit(self, count=None, submit_df=None, skip_failed=False):
+    def babs_submit(self, count=None, submit_df=None, skip_failed=False, skip_running_jobs=False):
         """
         This function submits jobs that don't have results yet and prints out job status.
 
@@ -28,22 +28,34 @@ def babs_submit(self, count=None, submit_df=None, skip_failed=False):
         submit_df: pd.DataFrame
             dataframe of jobs to be submitted
             default: None
+        skip_running_jobs: bool
+            whether to allow submission when there are running/pending jobs
         """
 
         # Check if there are still jobs running
         currently_running_df = self.get_currently_running_jobs_df()
+        running_pending_df = currently_running_df.copy()
         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 not skip_running_jobs:
+                    raise Exception(
+                        'There are still jobs running. '
+                        'Please wait for them to finish or cancel them. '
+                        'Current running jobs:\n'
+                        f'{currently_running_df}'
+                    )
+                if 'state' in currently_running_df:
+                    running_pending_df = currently_running_df[
+                        currently_running_df['state'].isin(['PD', 'R'])
+                    ]
+            else:
+                running_pending_df = currently_running_df.iloc[0:0]
+                print('All currently running jobs are in CG state; proceeding with submission.')
 
         # Find the rows that don't have results yet
         status_df = self.get_job_status_df()
@@ -54,6 +66,38 @@ def babs_submit(self, count=None, submit_df=None, skip_failed=False):
         if submit_df is not None:
             df_needs_submit = submit_df
 
+        if skip_running_jobs and not running_pending_df.empty:
+            # Build (sub_id,) or (sub_id, ses_id) keys for set lookup
+            if self.processing_level == 'session':
+                running_keys = set(
+                    zip(
+                        running_pending_df['sub_id'],
+                        running_pending_df['ses_id'],
+                        strict=False,
+                    )
+                )
+                submit_keys = list(
+                    zip(df_needs_submit['sub_id'], df_needs_submit['ses_id'], strict=False)
+                )
+            else:
+                running_keys = set(running_pending_df['sub_id'].tolist())
+                submit_keys = df_needs_submit['sub_id'].tolist()
+
+            # Mark which of the to-submit rows are still running/pending
+            if running_keys:
+                skip_mask = [key in running_keys for key in submit_keys]
+            else:
+                skip_mask = [False] * len(submit_keys)
+
+            if any(skip_mask):
+                # Report skipped job IDs and filter them out of the submission list
+                skip_job_ids = sorted(running_pending_df['job_id'].dropna().unique().tolist())
+                print(
+                    'Skipping running/pending jobs from job IDs: '
+                    + ', '.join(str(job_id) for job_id in skip_job_ids)
+                )
+                df_needs_submit = df_needs_submit.loc[~np.array(skip_mask)].reset_index(drop=True)
+
         # only run `babs submit` when there are subjects/sessions not yet submitted
         if df_needs_submit.empty:
             print('No jobs to submit')
diff --git a/tests/test_babs_workflow.py b/tests/test_babs_workflow.py
@@ -99,7 +99,11 @@ def test_babs_init_raw_bids(
 
     # babs submit:
     babs_submit_opts = argparse.Namespace(
-        project_root=project_root, select=None, inclusion_file=None, count=1
+        project_root=project_root,
+        select=None,
+        inclusion_file=None,
+        count=1,
+        skip_running_jobs=False,
     )
     with mock.patch.object(argparse.ArgumentParser, 'parse_args', return_value=babs_submit_opts):
         _enter_submit()
@@ -124,7 +128,11 @@ def test_babs_init_raw_bids(
 
     # Submit the last job:
     babs_submit_opts = argparse.Namespace(
-        project_root=project_root, select=None, inclusion_file=None, count=None
+        project_root=project_root,
+        select=None,
+        inclusion_file=None,
+        count=None,
+        skip_running_jobs=False,
     )
     with mock.patch.object(argparse.ArgumentParser, 'parse_args', return_value=babs_submit_opts):
         _enter_submit()
diff --git a/tests/test_interaction.py b/tests/test_interaction.py
@@ -4,6 +4,7 @@
 import pytest
 
 from babs.interaction import BABSInteraction
+from babs.utils import scheduler_status_columns
 
 
 def _minimal_status_df():
@@ -26,6 +27,26 @@ def _minimal_status_df():
     )
 
 
+def _status_df_for_submit():
+    return pd.DataFrame(
+        {
+            'sub_id': ['sub-01', 'sub-02', 'sub-03'],
+            'submitted': [True, True, False],
+            'has_results': [False, False, False],
+            'is_failed': [False, True, False],
+            'job_id': [10, 11, -1],
+            'task_id': [1, 1, -1],
+            'state': ['R', '', ''],
+            'time_used': ['0:01', '', ''],
+            'time_limit': ['5-00:00:00', '', ''],
+            'nodes': [1, 0, 0],
+            'cpus': [1, 0, 0],
+            'partition': ['normal', '', ''],
+            'name': ['test_array_job', '', ''],
+        }
+    )
+
+
 def test_babs_submit_blocks_non_cg_jobs(babs_project_subjectlevel, monkeypatch):
     babs_proj = BABSInteraction(project_root=babs_project_subjectlevel)
     running_df = pd.DataFrame(
@@ -76,3 +97,81 @@ def _mock_submit_array(analysis_path, queue, total_jobs):
     babs_proj.babs_submit(count=1)
 
     assert submit_calls
+
+
+def test_babs_submit_allows_running_skips_jobs(babs_project_subjectlevel, monkeypatch, capsys):
+    babs_proj = BABSInteraction(project_root=babs_project_subjectlevel)
+    running_df = pd.DataFrame(
+        {
+            'job_id': [10],
+            'task_id': [1],
+            'state': ['R'],
+            'time_used': ['0:01'],
+            'time_limit': ['5-00:00:00'],
+            'nodes': [1],
+            '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', _status_df_for_submit)
+
+    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(skip_running_jobs=True)
+
+    captured = capsys.readouterr()
+    assert submit_calls
+    assert submit_calls[0][2] == 2
+    assert 'Skipping running/pending jobs from job IDs' in captured.out
+    assert '10' in captured.out
+
+
+def test_get_currently_running_jobs_df_multiple_job_ids(babs_project_subjectlevel, monkeypatch):
+    babs_proj = BABSInteraction(project_root=babs_project_subjectlevel)
+    status_df = pd.DataFrame(
+        {
+            'sub_id': ['sub-01', 'sub-02'],
+            'submitted': [True, True],
+            'has_results': [False, False],
+            'is_failed': [False, False],
+            'job_id': [10, 20],
+            'task_id': [1, 2],
+        }
+    )
+    monkeypatch.setattr(babs_proj, 'get_job_status_df', lambda: status_df)
+    monkeypatch.setattr(babs_proj, 'get_latest_submitted_jobs_df', pd.DataFrame)
+
+    calls = []
+
+    def _mock_request_all_job_status(queue, job_id):
+        calls.append(job_id)
+        task_id = 1 if job_id == 10 else 2
+        return pd.DataFrame(
+            {
+                'job_id': [job_id],
+                'task_id': [task_id],
+                'state': ['R'],
+                'time_used': ['0:01'],
+                'time_limit': ['5-00:00:00'],
+                'nodes': [1],
+                'cpus': [1],
+                'partition': ['normal'],
+                'name': ['test_array_job'],
+            }
+        )[scheduler_status_columns]
+
+    monkeypatch.setattr('babs.base.request_all_job_status', _mock_request_all_job_status)
+
+    running_df = babs_proj.get_currently_running_jobs_df()
+
+    assert set(calls) == {10, 20}
+    assert set(running_df['sub_id']) == {'sub-01', 'sub-02'}
