282 lines
10 KiB
Python
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)
|