diff --git a/utilix/batchq.py b/utilix/batchq.py index a2766ea..95b894c 100644 --- a/utilix/batchq.py +++ b/utilix/batchq.py @@ -5,7 +5,7 @@ import subprocess import re import tempfile -from typing import Any, Dict, List, Literal, Optional +from typing import Any, Dict, List, Union, Literal, Optional from pydantic import BaseModel, Field, validator from simple_slurm import Slurm # type: ignore from . import logger @@ -429,7 +429,7 @@ def _get_lc_nodes(self) -> List[str]: print(f"An error occurred while executing nodestatus: {e}") return [] - def submit(self) -> None: + def submit(self) -> Union[int, None]: """Submit the job to the SLURM queue.""" os.makedirs(TMPDIR[self.partition], exist_ok=True) # Initialize a dictionary with mandatory parameters @@ -475,7 +475,7 @@ def submit(self) -> None: print(f"Generated slurm script:\n{slurm.script()}") if self.dry_run: - return + return None # Submit the job try: @@ -488,6 +488,8 @@ def submit(self) -> None: except Exception as e: print(f"An error occurred while submitting the job: {str(e)}") + return job_id + def submit_job( jobstring: str, @@ -511,7 +513,7 @@ def submit_job( dependency: Optional[str] = None, verbose: bool = False, bypass_validation: Optional[List[str]] = [], -) -> None: +) -> Union[int, None]: """Submit a job to the SLURM queue. Args: @@ -563,7 +565,8 @@ def submit_job( verbose=verbose, bypass_validation=bypass_validation, ) - job.submit() + job_id = job.submit() + return job_id def count_jobs(string: str = "") -> int: