From f8be7b97d43a516378df44473a2443c50298b704 Mon Sep 17 00:00:00 2001 From: Jayaram Kancherla Date: Thu, 6 Feb 2025 23:34:36 -0800 Subject: [PATCH] Add param to manifest for extra sbatch arguments (#83) --- CHANGELOG.md | 4 ++ src/cellarr/slurm/README.md | 7 +++- src/cellarr/slurm/build_cellarr_steps.py | 51 +++++++++++++++++++----- 3 files changed, 51 insertions(+), 11 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a0d6f88..d56e545 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,9 @@ # Changelog +## Version 0.5.6 + +- slrum runner's manifest now accepts an additional key: "sbatch_extra_args", to add any additional SBATCH headers to the generated scripts. + ## Version 0.5.3 - 0.5.5 - 'config' is renamed to 'context_or_config' to accept both TileDB's context or config objects. This supports in-memory TileDB contexts from [Phil](https://github.com/hanslovsky) [[#79](https://github.com/TileOme/cellarr/pull/79)]. diff --git a/src/cellarr/slurm/README.md b/src/cellarr/slurm/README.md index 8c2f264..f92d433 100644 --- a/src/cellarr/slurm/README.md +++ b/src/cellarr/slurm/README.md @@ -47,7 +47,12 @@ conda activate biocpy_miniforge ~/.conda/envs/biocpy_miniforge/bin/python --version which python python --version - """, +""", + "sbatch_extra_args": """ +#SBATCH --qos=short +#SBATCH --partition=defq +#SBATCH -t 0-02:00 +""" } ``` diff --git a/src/cellarr/slurm/build_cellarr_steps.py b/src/cellarr/slurm/build_cellarr_steps.py index 4f2d848..6beca1d 100644 --- a/src/cellarr/slurm/build_cellarr_steps.py +++ b/src/cellarr/slurm/build_cellarr_steps.py @@ -60,6 +60,7 @@ def create_slurm_script( args: Dict, dependencies: Optional[str] = None, python_env: str = "", + sbatch_extra_args: str = "", ) -> str: """Create a SLURM job submission script.""" script = f"""#!/bin/bash @@ -69,6 +70,7 @@ def create_slurm_script( #SBATCH --time={self.time_hours}:00:00 #SBATCH --mem={self.memory_gb}G #SBATCH --cpus-per-task={self.cpus_per_task} +{sbatch_extra_args} """ if dependencies: script += f"#SBATCH --dependency={dependencies}\n" @@ -91,6 +93,7 @@ def create_array_script( n_tasks: int, dependencies: Optional[str] = None, python_env: str = "", + sbatch_extra_args: str = "", ) -> str: """Create a SLURM array job submission script.""" script = f"""#!/bin/bash @@ -101,6 +104,7 @@ def create_array_script( #SBATCH --mem={self.memory_gb}G #SBATCH --cpus-per-task={self.cpus_per_task} #SBATCH --array=0-{n_tasks-1} +{sbatch_extra_args} """ if dependencies: script += f"#SBATCH --dependency={dependencies}\n" @@ -120,7 +124,9 @@ def submit_job(self, script_path: str) -> str: result = subprocess.run(["sbatch", script_path], capture_output=True, text=True) return result.stdout.strip().split()[-1] - def submit_gene_annotation_job(self, files: List[str], gene_options: Dict, python_env: str) -> str: + def submit_gene_annotation_job( + self, files: List[str], gene_options: Dict, python_env: str, sbatch_extra_args: str + ) -> str: """Submit gene annotation processing job.""" args = { "files": files, @@ -136,11 +142,12 @@ def submit_gene_annotation_job(self, files: List[str], gene_options: Dict, pytho python_script=f"{parent_dir}/process_gene_annotation.py", args=args, python_env=python_env, + sbatch_extra_args=sbatch_extra_args, ) return self.submit_job(script_path) def submit_sample_metadata_job( - self, files: List[str], sample_options: Dict, dependency: str, python_env: str + self, files: List[str], sample_options: Dict, dependency: str, python_env: str, sbatch_extra_args: str ) -> str: """Submit sample metadata processing job.""" args = { @@ -157,11 +164,14 @@ def submit_sample_metadata_job( python_script=f"{parent_dir}/process_sample_metadata.py", args=args, python_env=python_env, + sbatch_extra_args=sbatch_extra_args, # dependencies=f"afterok:{dependency}", ) return self.submit_job(script_path) - def submit_cell_metadata_job(self, files: List[str], cell_options: Dict, dependency: str, python_env: str) -> str: + def submit_cell_metadata_job( + self, files: List[str], cell_options: Dict, dependency: str, python_env: str, sbatch_extra_args: str + ) -> str: """Submit cell metadata processing job.""" args = { "files": files, @@ -177,12 +187,13 @@ def submit_cell_metadata_job(self, files: List[str], cell_options: Dict, depende python_script=f"{parent_dir}/process_cell_metadata.py", args=args, python_env=python_env, + sbatch_extra_args=sbatch_extra_args, # dependencies=f"afterok:{dependency}", ) return self.submit_job(script_path) def submit_matrix_processing( - self, files: List[str], matrix_options: Dict, dependency: str, python_env: str + self, files: List[str], matrix_options: Dict, dependency: str, python_env: str, sbatch_extra_args: str ) -> Tuple[str, str]: """Submit matrix processing as SLURM array job.""" @@ -213,6 +224,7 @@ def submit_matrix_processing( n_tasks=len(files), dependencies=f"afterok:{dependency}", python_env=python_env, + sbatch_extra_args=sbatch_extra_args, ) array_job_id = self.submit_job(array_script) @@ -230,12 +242,15 @@ def submit_matrix_processing( args=final_args, dependencies=f"afterok:{array_job_id}", python_env=python_env, + sbatch_extra_args=sbatch_extra_args, ) final_job_id = self.submit_job(final_script) return array_job_id, final_job_id - def submit_final_assembly(self, matrix_names: List[str], dependencies: List[str], python_env: str) -> str: + def submit_final_assembly( + self, matrix_names: List[str], dependencies: List[str], python_env: str, sbatch_extra_args: str + ) -> str: """Submit final assembly job.""" args = {"input_dir": str(self.output_dir), "output_dir": str(self.output_dir), "matrix_names": matrix_names} @@ -247,6 +262,7 @@ def submit_final_assembly(self, matrix_names: List[str], dependencies: List[str] args=args, dependencies=f"afterok:{','.join(dependencies)}", python_env=python_env, + sbatch_extra_args=sbatch_extra_args, ) return self.submit_job(script_path) @@ -286,15 +302,23 @@ def main(): # Submit jobs gene_job_id = builder.submit_gene_annotation_job( - manifest["files"], manifest.get("gene_options", {}), manifest["python_env"] + manifest["files"], manifest.get("gene_options", {}), manifest["python_env"], manifest["sbatch_extra_args"] ) sample_job_id = builder.submit_sample_metadata_job( - manifest["files"], manifest.get("sample_options", {}), gene_job_id, manifest["python_env"] + manifest["files"], + manifest.get("sample_options", {}), + gene_job_id, + manifest["python_env"], + manifest["sbatch_extra_args"], ) cell_job_id = builder.submit_cell_metadata_job( - manifest["files"], manifest.get("cell_options", {}), sample_job_id, manifest["python_env"] + manifest["files"], + manifest.get("cell_options", {}), + sample_job_id, + manifest["python_env"], + manifest["sbatch_extra_args"], ) # Process matrices @@ -305,13 +329,20 @@ def main(): matrix_job_ids = [] for matrix_opt in matrix_options: _, final_id = builder.submit_matrix_processing( - manifest["files"], matrix_opt, f"{cell_job_id},{gene_job_id}", manifest["python_env"] + manifest["files"], + matrix_opt, + f"{cell_job_id},{gene_job_id}", + manifest["python_env"], + manifest["sbatch_extra_args"], ) matrix_job_ids.append(final_id) # Submit final assembly builder.submit_final_assembly( - [opt["matrix_name"] for opt in matrix_options], matrix_job_ids, manifest["python_env"] + [opt["matrix_name"] for opt in matrix_options], + matrix_job_ids, + manifest["python_env"], + manifest["sbatch_extra_args"], )