Skip to content

Commit a3251a8

Browse files
agrouazeagrouaze
committed
full linted
1 parent ec73008 commit a3251a8

3 files changed

Lines changed: 167 additions & 94 deletions

File tree

‎src/turboblast/blaster.py‎

Lines changed: 135 additions & 76 deletions
Original file line numberDiff line numberDiff line change
@@ -1,56 +1,98 @@
1+
#!/usr/bin/python
12
import argparse
23
import datetime
34
import functools
5+
import logging
46
import os
7+
import shlex
58
import subprocess
69
import sys
10+
from pathlib import Path
711

8-
import submitit
12+
import submitit # type: ignore[import-not-found]
913

14+
from turboblast.logo import logo
1015

11-
def process_line(slurmexe, options_one_line):
12-
"""
13-
Function executed by each task in the job array.
16+
# Configure the logger globally
17+
# This format matches standard logging practices: [Date Time] [LEVEL] Message
18+
logging.basicConfig(
19+
level=logging.INFO,
20+
format="%(asctime)s [%(levelname)s] %(message)s",
21+
datefmt="%Y-%m-%d %H:%M:%S",
22+
handlers=[logging.StreamHandler(sys.stdout)],
23+
)
24+
25+
logger = logging.getLogger(__name__)
26+
27+
28+
def process_line(slurmexe: str, options_one_line: str) -> None:
29+
"""Executes a single task in the Slurm job array by calling a bash script.
30+
31+
This function is serialized and executed on the Slurm compute node by submitit.
32+
It constructs the command, sets up the environment, and streams both standard
33+
output and standard error directly to the submitit log files.
34+
35+
Args:
36+
slurmexe (str): The path to the bash executable/script to run.
37+
options_one_line (str): A string containing the command-line arguments
38+
to pass to the bash script (e.g., "--input file.txt --output dir/").
39+
40+
Raises:
41+
subprocess.CalledProcessError: If the bash command exits with a non-zero
42+
status code, ensuring submitit and Slurm mark the task as FAILED.
1443
"""
15-
print(f"[Task options: {options_one_line}] Starting computation", flush=True)
44+
logger.info("Starting computation for options: %s", options_one_line)
1645

1746
# Construct the command
18-
# We use shlex.split if options are complex, but for simple strings list structure is better
19-
# However, since shell=False is safer, we construct the command list
20-
cmd = ["bash", slurmexe] + options_one_line.split()
47+
# shlex.split is safer than string.split() if your options contain quoted strings
48+
# RUF005: Use unpacking instead of list concatenation
49+
cmd = ["bash", slurmexe, *shlex.split(options_one_line)]
2150

22-
print(f"Executing: {' '.join(cmd)}", flush=True)
51+
logger.info("Executing command: %s", " ".join(cmd))
2352

24-
# We use subprocess.call or run WITHOUT capture_output.
25-
# This allows logs to stream directly to the submitit log files in real-time.
2653
# 1. Prepare the environment correctly
27-
# Copy the current environment (including PATH, LD_LIBRARY_PATH, etc.)
2854
full_env = os.environ.copy()
29-
# Add your specific variable
3055
full_env["PYTHONUNBUFFERED"] = "1"
56+
3157
try:
32-
# check=True will raise a CalledProcessError if return code != 0
33-
# subprocess.run(cmd, shell=False, check=True, env=env)
3458
subprocess.run(
3559
cmd,
3660
shell=False,
3761
check=True,
3862
stdout=sys.stdout,
39-
stderr=sys.stderr,
63+
# Redirect stderr to stdout to merge logs chronologically.
64+
# This ensures INFO and WARNING/ERROR messages don't get out of sync.
65+
stderr=subprocess.STDOUT,
4066
env=full_env,
4167
)
42-
sys.stdout.flush() # Force a flush after the process finishes
43-
print(
44-
f"Task with options {options_one_line} completed successfully.", flush=True
45-
)
68+
# Flush to ensure everything is written to the submitit .out file immediately
69+
sys.stdout.flush()
70+
logger.info("Task completed successfully for options: %s", options_one_line)
71+
4672
except subprocess.CalledProcessError as e:
47-
print(f"Error in task with options: {options_one_line}", file=sys.stderr)
48-
print(f"Command returned exit status {e.returncode}", file=sys.stderr)
49-
raise e # Re-raise exception so submitit marks the job as FAILED
73+
# TRY400: Use logging.exception instead of logging.error inside except block
74+
logger.exception(
75+
"Task failed for options: %s (exit status %s)",
76+
options_one_line,
77+
e.returncode,
78+
)
79+
# TRY201: Use bare raise
80+
raise
81+
82+
83+
def parser_args() -> argparse.Namespace:
84+
"""Parses command-line arguments for the submitit client.
5085
86+
Returns:
87+
argparse.Namespace: An object containing all the parsed command-line
88+
arguments and their values.
89+
"""
5190

52-
def parser_args():
53-
parser = argparse.ArgumentParser(description="Submitit array job example")
91+
# formatter_class=argparse.RawDescriptionHelpFormatter is REQUIRED
92+
# to preserve the newlines and spaces of the ASCII art logo!
93+
parser = argparse.ArgumentParser(
94+
description=logo, formatter_class=argparse.RawDescriptionHelpFormatter
95+
)
5496
parser.add_argument(
5597
"--num-tasks",
5698
type=int,
@@ -81,7 +123,7 @@ def parser_args():
81123
)
82124
parser.add_argument(
83125
"--output-dir",
84-
help="Directory to store submitit logs (suffixed with out/YYYYMMDDTHHMMSS/)",
126+
help="Directory to store submitit logs",
85127
default="submitit_logs_array",
86128
)
87129
parser.add_argument(
@@ -90,91 +132,108 @@ def parser_args():
90132
return parser.parse_args()
91133

92134

93-
def main(args):
135+
def main(args: argparse.Namespace) -> None:
136+
"""Configures and submits the job arrays to the Slurm cluster.
137+
138+
This function reads the input file, configures the submitit AutoExecutor
139+
with the required Slurm parameters (memory, partition, time, etc.), and
140+
submits the jobs in chunks to respect Slurm array limits.
141+
142+
Args:
143+
args (argparse.Namespace): The parsed command-line arguments containing
144+
paths and Slurm configuration variables.
145+
"""
94146
# Create log directory
95-
output_dir_with_date = os.path.join(
96-
args.output_dir, datetime.datetime.today().strftime("%Y%m%dT%H%M%S")
97-
)
98-
os.makedirs(output_dir_with_date, exist_ok=True)
147+
# DTZ002: Use datetime.now with a timezone instead of today()
148+
output_dir_with_date = Path(args.output_dir) / datetime.datetime.now(
149+
datetime.timezone.utc
150+
).strftime("%Y%m%dT%H%M%S")
151+
152+
# Use instance method for directory creation
153+
output_dir_with_date.mkdir(parents=True, exist_ok=True)
154+
155+
logger.info("Submitit logs will be stored in: %s", output_dir_with_date)
156+
logger.info("Bash script to execute: %s", args.bash_slurm_exec)
99157

100158
executor = submitit.AutoExecutor(folder=output_dir_with_date, cluster="slurm")
101-
print("args.bash_slurm_exec", args.bash_slurm_exec)
159+
102160
# Slurm Configuration
103161
executor.update_parameters(
104162
timeout_min=args.timeout_min,
105163
mem_gb=args.mem_gb,
106164
cpus_per_task=args.cpus_per_task,
107165
slurm_partition=args.slurm_partition,
108166
slurm_array_parallelism=args.slurm_array_parallelism,
109-
# Naming the job makes it easier to find in squeue
110-
slurm_job_name=os.path.basename(args.bash_slurm_exec).replace(".sh", ""),
167+
slurm_job_name=Path(args.bash_slurm_exec).name.replace(".sh", ""),
111168
slurm_additional_parameters={"export": "ALL,PYTHONUNBUFFERED=1"},
112169
)
113170

114171
# Read inputs
115-
with open(args.listing_input) as f:
116-
# Filter out empty lines just in case
172+
# Use Path object method to correctly open the file (fixes mypy overload error)
173+
with Path(args.listing_input).open("r") as f:
117174
array_inputs = [line.strip() for line in f if line.strip()]
118175

119176
if not array_inputs:
120-
print("Error: Input listing file is empty.")
177+
logger.error("Input listing file is empty. Aborting submission.")
121178
return
122179

123-
print(f"Submitting {len(array_inputs)} tasks...")
124-
125-
# if an image must be pulled from registry, honestly it is much easier with a .sif on fs.
126-
# reason why the block below is commented out.
127-
# Pull the image on the submission node first
128-
# img_url = "oras://registry.hpc.ifremer.fr/lops-siam-sentinel1-workbench/unifiedwvalticolocs:2026.2.17"
129-
# img_sif = "unifiedwvalticolocs_2026.2.17.sif"
130-
131-
# if not os.path.exists(img_sif):
132-
# print(f"Pulling image {img_url}...")
133-
# subprocess.run(["apptainer", "pull", img_sif, img_url], check=True)
180+
# Define the chunk size (e.g., 1000)
181+
chunk_size = 1000
182+
all_jobs = []
134183

135-
# KEY CHANGE: Use partial to bind the bash script path to the function
136-
# The map_array will only vary the second argument (options_one_line)
184+
logger.info("Total tasks to submit: %d", len(array_inputs))
137185

138-
# Define the chunk size (set this slightly below your MaxArraySize, e.g., 500 or 1000)
139-
chunk_size = 1000 # Adjust based on your cluster's MaxArraySize and the total number of tasks, here I am not sure this is the limit but it failed with 1379 lines
140-
all_jobs = []
186+
# Calculate total number of chunks for logging
187+
total_chunks = (len(array_inputs) + chunk_size - 1) // chunk_size
188+
logger.info(
189+
"Submitting tasks in chunks of %d (Total chunks: %d)...",
190+
chunk_size,
191+
total_chunks,
192+
)
141193

142-
print(f"Total tasks to submit: {len(array_inputs)}")
194+
process_func = functools.partial(process_line, args.bash_slurm_exec)
143195

144196
# Loop through the inputs in chunks
145-
chuncked_inputs = range(0, len(array_inputs), chunk_size)
146-
print(f"Submitting tasks in chunks of {chunk_size}...")
147-
print(f"Total chunks to submit: {len(list(chuncked_inputs))}")
148-
for i in chuncked_inputs:
197+
for chunk_idx, i in enumerate(range(0, len(array_inputs), chunk_size), start=1):
149198
chunk = array_inputs[i : i + chunk_size]
150-
print(f"Submitting chunk starting at index {i} (size: {len(chunk)})...")
199+
logger.info(
200+
"Submitting chunk %d/%d (size: %d, starting at index %d)...",
201+
chunk_idx,
202+
total_chunks,
203+
len(chunk),
204+
i,
205+
)
151206

152-
# This will create a NEW Slurm Job ID for every 1000 tasks
153-
process_func = functools.partial(process_line, args.bash_slurm_exec)
207+
# This will create a NEW Slurm Job Array for every 1000 tasks
154208
jobs = executor.map_array(process_func, chunk)
155209

156-
print(f"Submitted. Job ID for this chunk: {jobs[0].job_id}")
210+
logger.info(
211+
"Chunk %d submitted successfully. Job Array ID: %s",
212+
chunk_idx,
213+
jobs[0].job_id,
214+
)
157215
all_jobs.extend(jobs)
158216

159-
print(
160-
f"Successfully submitted all {len(all_jobs)} tasks across multiple Job Arrays."
217+
logger.info(
218+
"Successfully submitted all %d tasks across %d Job Arrays.",
219+
len(all_jobs),
220+
total_chunks,
161221
)
162222

163-
# process_func = functools.partial(process_line, args.bash_slurm_exec)
164-
165-
# # Submit the array
166-
# jobs = executor.map_array(process_func, array_inputs)
223+
if all_jobs:
224+
logger.info("First Job Array Main Job ID: %s", all_jobs[0].job_id)
225+
logger.info(
226+
"To monitor specific task logs, use: tail -f %s/<JOB_ID>_<TASK_ID>_0.out",
227+
output_dir_with_date,
228+
)
167229

168-
# Get the ID of the array job
169-
# jobs[0] gives access to the whole array group usually, or iterate to see IDs
170-
print(f"Job Array submitted. Main Job ID: {jobs[0].job_id}")
171-
print(f"Logs are being written to: {output_dir_with_date}")
172-
print(
173-
f"You can monitor specific logs using: tail -f {output_dir_with_date}/<JOB_ID>_<TASK_ID>_0.out"
174-
)
175230

231+
def entrypoint() -> None:
232+
"""Script entrypoint.
176233
177-
def entrypoint():
234+
Calls the argument parser and passes the resulting arguments to the
235+
main execution function.
236+
"""
178237
args = parser_args()
179238
main(args)
180239

‎src/turboblast/logo.py‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
logo = r"""
2+
_____ _ _____________ ___________ _ ___ _____ _____
3+
|_ _| | | | ___ \ ___ \| _ | ___ \ | / _ \ / ___|_ _|
4+
| | | | | | |_/ / |_/ /| | | | |_/ / | / /_\ \\ `--. | |
5+
| | | | | | /| ___ \| | | | ___ \ | | _ | `--. \ | |
6+
| | | |_| | |\ \| |_/ /\ \_/ / |_/ / |____| | | |/\__/ / | |
7+
\_/ \___/\_| \_\____/ \___/\____/\_____/\_| |_/\____/ \_/
8+
9+
Submitit Client for High-Throughput Slurm Job Arrays
10+
======================================================================
11+
"""

‎src/turboblast/submitit_logs_analysis.sh‎

Lines changed: 21 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
#!/bin/bash
22

3-
LOG_DIR=$1
3+
LOG_DIR="$1"
44

55
if [ -z "$LOG_DIR" ]; then
66
echo "Usage: ./submitit_analyze.sh <path_to_submitit_logs>"
@@ -14,12 +14,13 @@ fi
1414

1515
# Function to calculate percentage
1616
get_perc() {
17-
local count=$1
18-
local total=$2
17+
local count="$1"
18+
local total="$2"
1919
if [ "$total" -eq 0 ]; then
2020
echo "0.0"
2121
else
22-
echo | awk "{printf \"%.1f\", ($count / $total) * 100}"
22+
# SC2027/SC2086 fix: Using -v to pass bash variables safely into awk
23+
awk -v c="$count" -v t="$total" 'BEGIN { printf "%.1f", (c / t) * 100 }'
2324
fi
2425
}
2526

@@ -28,40 +29,42 @@ echo " ANALYZING SUBMITIT LOGS IN: $LOG_DIR"
2829
echo "----------------------------------------------------------------"
2930

3031
# Count total unique tasks (based on .out files)
31-
TOTAL_LOGS=$(ls -1 $LOG_DIR/*.out 2>/dev/null | wc -l)
32+
# SC2012 fix: Use 'find' instead of 'ls' to safely count files
33+
TOTAL_LOGS=$(find "$LOG_DIR" -maxdepth 1 -name "*.out" 2>/dev/null | wc -l)
3234

3335
if [ "$TOTAL_LOGS" -eq 0 ]; then
3436
echo "No log files found in $LOG_DIR"
3537
exit 0
3638
fi
3739

3840
# 1. Success Count
39-
SUCCESS=$(grep -l -iE "completed successfully|success" $LOG_DIR/*.out 2>/dev/null | wc -l)
40-
SUCCESS_P=$(get_perc $SUCCESS $TOTAL_LOGS)
41+
SUCCESS=$(grep -l -iE "completed successfully|success" "$LOG_DIR"/*.out 2>/dev/null | wc -l)
42+
# SC2086 fixes below: Double quote all variables
43+
SUCCESS_P=$(get_perc "$SUCCESS" "$TOTAL_LOGS")
4144

4245
# 2. Permission Issues
43-
PERM_ERR=$(grep -l -iE "Permission denied|AccessDenied|EACCES" $LOG_DIR/*.{out,err} 2>/dev/null | sort -u | wc -l)
44-
PERM_P=$(get_perc $PERM_ERR $TOTAL_LOGS)
46+
PERM_ERR=$(grep -l -iE "Permission denied|AccessDenied|EACCES" "$LOG_DIR"/*.{out,err} 2>/dev/null | sort -u | wc -l)
47+
PERM_P=$(get_perc "$PERM_ERR" "$TOTAL_LOGS")
4548

4649
# 3. Memory / Killed Issues
47-
KILLED_ERR=$(grep -l -iE "Out of memory|Killed|OOM killer|slurm_step_terminate" $LOG_DIR/*.{out,err} 2>/dev/null | sort -u | wc -l)
48-
KILLED_P=$(get_perc $KILLED_ERR $TOTAL_LOGS)
50+
KILLED_ERR=$(grep -l -iE "Out of memory|Killed|OOM killer|slurm_step_terminate" "$LOG_DIR"/*.{out,err} 2>/dev/null | sort -u | wc -l)
51+
KILLED_P=$(get_perc "$KILLED_ERR" "$TOTAL_LOGS")
4952

5053
# 4. General Python Tracebacks
51-
PYTHON_ERR=$(grep -l "Traceback (most recent call last):" $LOG_DIR/*.err 2>/dev/null | sort -u | wc -l)
52-
PYTHON_P=$(get_perc $PYTHON_ERR $TOTAL_LOGS)
54+
PYTHON_ERR=$(grep -l "Traceback (most recent call last):" "$LOG_DIR"/*.err 2>/dev/null | sort -u | wc -l)
55+
PYTHON_P=$(get_perc "$PYTHON_ERR" "$TOTAL_LOGS")
5356

5457
# 5. Socket / Slurm Communication Errors
55-
SOCKET_ERR=$(grep -l -iE "socket|missing socket|confirm allocation" $LOG_DIR/*.{out,err} 2>/dev/null | sort -u | wc -l)
56-
SOCKET_P=$(get_perc $SOCKET_ERR $TOTAL_LOGS)
58+
SOCKET_ERR=$(grep -l -iE "socket|missing socket|confirm allocation" "$LOG_DIR"/*.{out,err} 2>/dev/null | sort -u | wc -l)
59+
SOCKET_P=$(get_perc "$SOCKET_ERR" "$TOTAL_LOGS")
5760

5861
# 6. Apptainer specific errors
59-
CONTAINER_ERR=$(grep -l -iE "apptainer: error|FATAL:|image not found" $LOG_DIR/*.{out,err} 2>/dev/null | sort -u | wc -l)
60-
CONTAINER_P=$(get_perc $CONTAINER_ERR $TOTAL_LOGS)
62+
CONTAINER_ERR=$(grep -l -iE "apptainer: error|FATAL:|image not found" "$LOG_DIR"/*.{out,err} 2>/dev/null | sort -u | wc -l)
63+
CONTAINER_P=$(get_perc "$CONTAINER_ERR" "$TOTAL_LOGS")
6164

6265
# Calculate General Failures
6366
FAIL_TOTAL=$((TOTAL_LOGS - SUCCESS))
64-
FAIL_P=$(get_perc $FAIL_TOTAL $TOTAL_LOGS)
67+
FAIL_P=$(get_perc "$FAIL_TOTAL" "$TOTAL_LOGS")
6568

6669
echo "Total Tasks Scanned: $TOTAL_LOGS"
6770
echo "----------------------------------------------------------------"

0 commit comments

Comments
 (0)