Files
2025-09-10 10:37:55 +02:00

282 lines
10 KiB
Python

import argparse
import os
import re
import subprocess
import time
def get_job_data():
# get job ids, partitions, runtime and total time
result = subprocess.run(
["squeue", "--user=rr41qemu", "--format=%i,%P,%j,%T,%M,%l", "--noheader"],
capture_output=True,
text=True,
# shell=True,
)
if result.returncode != 0:
print("Error getting job ids.")
return []
job_data = dict()
for line in result.stdout.splitlines():
if line.strip():
job_id, partition, job_name, state, runtime, total_time = line.split(",")
job_data[job_id] = {
"partition": partition,
"job_name": job_name,
"state": state,
"runtime": runtime,
"total_time": total_time,
}
return job_data
def tail_progress(log_file, n=1) -> str:
# get last modified time
output = ""
c_time = os.path.getmtime(log_file)
try:
with open(log_file, "r") as f:
lines = f.readlines()
output += f"Last update time: {time.ctime(c_time)}\n"
for line in lines[-n:]:
if "%" in line: # crude tqdm-like filter
output = output + line.strip()
except FileNotFoundError:
output += f"Log file {log_file} not found.\n"
pass
except Exception as e:
output += f"Error reading log file {log_file}: {e}\n"
return output
def get_stats(log_file) -> str:
try:
with open(log_file, "r") as f:
lines = f.readlines()
# find num_gpus
num_gpus = 0
for line in lines:
if "Training on " in line:
match = re.search(r"Training on (\d+) GPUs", line)
num_gpus = int(match.group(1))
break
# find num_gpu lines to get models
gpu_models = {}
searched = False
for line in reversed(lines):
# pattern is "Rank X: Current device: $device_name on local rank $rank"
if "Rank " in line and "Current device: " in line:
match = re.search(r"Rank (\d+): Current device: (.+) on local rank \d+", line)
if match:
rank = int(match.group(1))
if rank == 0 and searched:
# if rank 0 is already searched, we can stop searching
break
elif rank == 0:
searched = True
if rank not in gpu_models:
gpu_models[rank] = dict()
device_name = match.group(2)
gpu_models[rank]["device_name"] = device_name
# get batch size, pattern: "Rank 0: Estimated batch size: 64 for GPU xxx"
for rank in gpu_models.keys():
for line in reversed(lines):
if f"Rank {rank}: Estimated batch size: " in line:
match = re.search(r"Rank (\d+): Estimated batch size: (\d+) for GPU", line)
if match:
gpu_models[rank]["batch_size"] = int(match.group(2))
break
# find run_id
run_id = None
for line in lines:
if "Run ID: " in line:
match = re.search(r"Run ID: (.+)", line)
if match:
run_id = match.group(1)
break
# find number of trainings
num_trainings = 0
for line in lines:
if "Rank 0: Number of trainings: " in line:
match = re.search(r"Rank 0: Number of trainings: (\d+)", line)
if match:
num_trainings = int(match.group(1))
break
# get current model training name
current_model_training_name = None
for line in reversed(lines):
# pattern is "Rank 0: running $model_name"
if "Rank 0: Running: " in line:
match = re.search(r"Rank 0: Running: (.+)", line)
if match:
current_model_training_name = match.group(1)
break
# get model parameters, pattern: " Creating model with parameters: {'embed_dim': 128, 'num_enc_layers': 4, 'num_heads': 4, 'seq_len': 240}"
model_parameters = None
for line in reversed(lines):
if "Creating model with parameters: " in line:
match = re.search(r"Creating model with parameters: (.+)", line)
if match:
model_parameters = match.group(1)
break
# get input parameters pattern is "Input_parameters: {'item_limit': 1000, 'num_dataloader_workers': 4}"
input_parameters = None
for line in reversed(lines):
if "Input_parameters: " in line:
match = re.search(r"Input_parameters: (.+)", line)
if match:
input_parameters = match.group(1)
break
# find current training number
current_training_number = 0
for line in reversed(lines):
if "Rank 0: Running: " in line:
current_training_number += 1
if num_trainings == 0:
num_trainings = f"{current_training_number}+"
# find last epoch
current_epoch = None
num_epochs = None
for line in reversed(lines):
# pattern is "Rank 0: Epoch: 1/10"
if "Rank 0: Epoch" in line:
match = re.search(r"Rank 0: Epoch (\d+)/(\d+)", line)
if match:
current_epoch = int(match.group(1))
num_epochs = int(match.group(2))
break
# get last validation loss
val_losses = list()
if current_epoch > 1:
for line in lines:
# patttern is "Rank 0: Overall val loss: 0.4439"
if "Rank 0: Overall val loss" in line:
match = re.search(r"Rank 0: Overall val loss: ([\d.]+)", line)
if match:
val_losses.append(float(match.group(1)))
if val_losses:
last_val_loss = val_losses[-1]
best_val_loss = min(val_losses)
else:
last_val_loss = None
best_val_loss = None
else:
last_val_loss = None
best_val_loss = None
is_validation = False
is_training = False
is_evaluation = False
for line in reversed(lines):
if "Evaluating model" in line:
is_evaluation = True
break
if "Validation" in line:
is_validation = True
break
if "Rank 0: Epoch " in line:
is_training = True
break
if is_validation:
status = "Validation"
elif is_training:
status = "Training"
elif is_evaluation:
status = "Evaluation"
else:
status = "Unknown"
# create output
output = ""
output += f"Run ID: {run_id}, Current training: {current_training_number} of {num_trainings}\n"
output += f"Number of GPUs: {num_gpus}\n"
output += f"GPUs used: {gpu_models}\n"
output += f"Current model training name: {current_model_training_name}\n"
output += f"Model parameters: {model_parameters}\n"
output += f"Input parameters: {input_parameters}\n"
output += f"Current epoch: {current_epoch} of {num_epochs}\n"
output += f"Last validation loss: {last_val_loss}, Best validation loss: {best_val_loss}\n"
output += f"Status: {status}\n"
except FileNotFoundError:
output = f"Log file {log_file} not found."
pass
except Exception as e:
output = f"Error reading log file {log_file}: {e}"
return output
def monitor_jobs(logdir, interval=5):
try:
while True:
job_data = get_job_data()
outputs = list()
for job_id, job in job_data.items():
# add header with job data
outputs.append(f"\n--- Job {job_id} ---")
outputs.append(f"Partition: {job['partition']}, Job Name: {job['job_name']}, State: {job['state']}, "
f"Runtime: {job['runtime']}, Total Time: {job['total_time']}\n")
err_file_found = True
out_file_found = True
err_file = os.path.join(logdir, f"{job['job_name']}_{job_id}.err")
if not os.path.exists(err_file):
# try alternative with job_id only
err_file = os.path.join(logdir, f"{job_id}.err")
if not os.path.exists(err_file):
err_file_found = False
out_file = os.path.join(logdir, f"{job['job_name']}_{job_id}.out")
if not os.path.exists(out_file):
# try alternative with job_id only
out_file = os.path.join(logdir, f"{job_id}.out")
if not os.path.exists(out_file):
out_file_found = False
if err_file_found and out_file_found:
stat_output = get_stats(out_file)
progress_output = tail_progress(err_file)
outputs.append(stat_output)
outputs.append(progress_output)
else:
if not err_file_found:
outputs.append(f"Error log file: {err_file} not found.\n")
if not out_file_found:
outputs.append(f"Output log file: {out_file} not found.\n")
outputs.append("Job might not have started yet.\n")
os.system('clear')
print(f"Monitoring training logs in: {logdir}")
print("\n".join(outputs))
time.sleep(interval)
except KeyboardInterrupt:
print("Stopped monitoring.")
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Monitor SLURM job progress from log files.")
parser.add_argument("--logdir", required=True, help="Directory containing SLURM .err log files")
parser.add_argument("--interval", type=int, default=5, help="Update interval in seconds")
args = parser.parse_args()
monitor_jobs(args.logdir, args.interval)