483 lines
17 KiB
Python
483 lines
17 KiB
Python
"""
|
|
@package pmsco.schedule
|
|
Job schedule interface
|
|
|
|
This module defines common infrastructure to submit a PMSCO calculation job to a job scheduler such as Slurm.
|
|
|
|
The schedule can be defined as part of the run-file (see pmsco module).
|
|
Users may derive sub-classes in a separate module to adapt to their own computing cluster.
|
|
|
|
The basic call sequence is:
|
|
1. Create a schedule object.
|
|
2. Initialize its properties with job parameters.
|
|
3. Validate()
|
|
4. Submit()
|
|
|
|
@author Matthias Muntwiler, matthias.muntwiler@psi.ch
|
|
|
|
@copyright (c) 2015-23 by Paul Scherrer Institut @n
|
|
Licensed under the Apache License, Version 2.0 (the "License"); @n
|
|
you may not use this file except in compliance with the License.
|
|
You may obtain a copy of the License at
|
|
http://www.apache.org/licenses/LICENSE-2.0
|
|
"""
|
|
|
|
import collections.abc
|
|
import datetime
|
|
import logging
|
|
import os
|
|
from pathlib import Path
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
from typing import Any, Callable, Dict, Generator, Iterable, Iterator, List, Mapping, Optional, Sequence, Set, Tuple, Union
|
|
|
|
try:
|
|
import commentjson as json
|
|
except ImportError:
|
|
import json
|
|
|
|
import pmsco.config
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class JobSchedule(pmsco.config.ConfigurableObject):
|
|
"""
|
|
Base class for job schedule
|
|
|
|
Usage:
|
|
1. Create object, assigning a project instance.
|
|
2. Set attributes.
|
|
3. Assign run_dict.
|
|
4. Call validate.
|
|
5. Call submit.
|
|
|
|
An example can be seen in pmsco.schedule_project.
|
|
|
|
This class defines the abstract interface, common actions and some utilities.
|
|
Derived classes may override any method, but should call the inherited method.
|
|
"""
|
|
|
|
## @var enabled (bool)
|
|
#
|
|
# This parameter signals whether pmsco should schedule a job or run the calculation.
|
|
# It is not directly used by the schedule classes but by the pmsco module.
|
|
# It must be defined in the run file and set to true to submit the job to a scheduler.
|
|
# It is set to false in the run file copied to the job directory so that the job script starts the calculation.
|
|
|
|
def __init__(self, project):
|
|
super(JobSchedule, self).__init__()
|
|
self.project = project
|
|
|
|
# Specifies whether a job script is created (True) or not (False, default).
|
|
self.enabled = False
|
|
|
|
# Specifies whether the job script is submitted manually by the user (True, default)
|
|
# or as a part of this class' execution (False).
|
|
self.manual = True
|
|
|
|
# Specifies whether an existing job directory is overwritten (True).
|
|
# Otherwise, an exception is raised if it exists (False, default).
|
|
self.overwrite_job_dir = False
|
|
|
|
# Content of the run-file.
|
|
# This must be set by the caller before the object is executed.
|
|
self.run_dict = {}
|
|
|
|
# Path to the destination run-file.
|
|
# This value is set by self.validate.
|
|
self.run_file = Path()
|
|
|
|
# Job work directory.
|
|
# This value is set by self.validate.
|
|
self.job_dir = Path()
|
|
|
|
# Dest path of the shell script to be submitted to the queue.
|
|
# This value is set by self.validate.
|
|
self.job_file = Path()
|
|
|
|
# Directory that contains the pmsco code.
|
|
# This value is set by self.validate.
|
|
self.pmsco_root_dir = Path(__file__).parent.parent
|
|
|
|
# Directory that contains the project module
|
|
# This value is set by self.validate.
|
|
self.project_dir = Path()
|
|
|
|
# Project files to be copied to the job_dir.
|
|
# Paths should be absolute or relative to project_dir.
|
|
# The project module is appended by self.validate.
|
|
# Other files have to be added by the caller.
|
|
self.project_files = []
|
|
|
|
# Name of the conda environment to activate (`source activate` command).
|
|
# Either conda_env or virtual_env should be specified.
|
|
# If both are empty, the environment is detected from the environment of the current process.
|
|
self.conda_env = ""
|
|
|
|
# Path of the virtual environment to activate (must contain the `activate` command).
|
|
# Either conda_env or virtual_env should be specified.
|
|
# If both are empty, the environment is detected from the environment of the current process.
|
|
self.virtual_env = ""
|
|
|
|
def validate(self):
|
|
"""
|
|
validate the job parameters.
|
|
|
|
make sure all object attributes are correct for submission.
|
|
|
|
@return: None
|
|
"""
|
|
|
|
assert self.run_dict, "run_dict not set"
|
|
|
|
self.pmsco_root_dir = Path(self.project.directories['pmsco']).parent
|
|
self.project_dir = Path(self.project.directories['project'])
|
|
output_dir = Path(self.project.directories['output'])
|
|
|
|
assert self.pmsco_root_dir.is_dir()
|
|
assert (self.pmsco_root_dir / "pmsco").is_dir(), "can't find pmsco directory (source code)"
|
|
assert self.project_dir.is_dir(), "can't find project directory"
|
|
assert output_dir.is_dir(), "can't find output directory"
|
|
assert output_dir.is_absolute(), "output directory must be an absolute path"
|
|
assert self.project.job_name, "job_name is undefined"
|
|
|
|
if output_dir.name == self.project.job_name:
|
|
self.job_dir = output_dir
|
|
else:
|
|
self.job_dir = output_dir / self.project.job_name
|
|
try:
|
|
self.job_dir.mkdir(parents=True, exist_ok=self.overwrite_job_dir)
|
|
except FileExistsError:
|
|
logger.error("job directory exists - check job name or clean up manually")
|
|
raise
|
|
|
|
self.job_file = (self.job_dir / self.project.job_name).with_suffix(".sh")
|
|
self.run_file = (self.job_dir / self.project.job_name).with_suffix(".json")
|
|
|
|
self.project_files.append(sys.modules[self.project.__module__].__file__)
|
|
|
|
def submit(self):
|
|
"""
|
|
submit the job to the scheduler.
|
|
|
|
as of this class, the method does to following:
|
|
|
|
1. copy source files
|
|
2. copy a patched version of the run file.
|
|
3. write the job file (_write_job_file must be implemented by a derived class).
|
|
|
|
@return: None
|
|
"""
|
|
self._copy_source()
|
|
self._fix_run_file()
|
|
self._write_run_file()
|
|
self._write_job_file()
|
|
|
|
def _copy_source(self):
|
|
"""
|
|
copy the source files to the job directory.
|
|
|
|
the files to copy must be listed explicitly in the project_files attribute.
|
|
the files are copied to the job_dir directory.
|
|
the directory must exist.
|
|
|
|
@return: None
|
|
"""
|
|
|
|
files = set((self.project_dir.joinpath(pf) for pf in self.project_files))
|
|
for f in files:
|
|
shutil.copy2(f, self.job_dir)
|
|
|
|
def _fix_run_file(self):
|
|
"""
|
|
fix the run file.
|
|
|
|
patch some entries of self.run_dict so that it can be used as run file.
|
|
the following changes are made:
|
|
1. set schedule.enabled to false so that the calculation is run.
|
|
2. set the output directory to the job directory.
|
|
3. set the log file to the job directory.
|
|
|
|
@return: None
|
|
"""
|
|
self.run_dict['schedule']['enabled'] = False
|
|
self.run_dict['project']['directories']['output'] = str(self.job_dir)
|
|
self.run_dict['project']['job_name'] = self.project.job_name
|
|
self.run_dict['project']['log_file'] = str((self.job_dir / self.project.job_name).with_suffix(".log"))
|
|
|
|
def _write_run_file(self):
|
|
"""
|
|
copy the run file.
|
|
|
|
this is a JSON dump of self.run_dict to the self.run_file file.
|
|
|
|
@return: None
|
|
"""
|
|
with open(self.run_file, "wt") as f:
|
|
json.dump(self.run_dict, f, indent=2)
|
|
|
|
def _write_job_file(self):
|
|
"""
|
|
create the job script.
|
|
|
|
this method must be implemented by a derived class.
|
|
the script must be written to the self.job_file file.
|
|
don't forget to make the file executable.
|
|
|
|
@return: None
|
|
"""
|
|
pass
|
|
|
|
@staticmethod
|
|
def detect_env() -> Dict[str, os.PathLike]:
|
|
"""
|
|
detect the python environment
|
|
|
|
determines the current python environment.
|
|
|
|
examples:
|
|
- /das/work/p17/p17274/conda/envs/pmsco310/bin
|
|
- /home/user/envs/pmsco-uv
|
|
|
|
@return: dictionary type -> path containing one or zero items.
|
|
type is either 'conda', 'venv' or 'system';
|
|
path is the bin directory containing python.
|
|
"""
|
|
|
|
pp = Path(sys.executable).parent
|
|
if (pp / 'activate').is_file():
|
|
return {"venv": pp}
|
|
for parent in pp.parents:
|
|
if (parent / "condabin").is_dir():
|
|
return {"conda": pp}
|
|
else:
|
|
return {"system": pp}
|
|
|
|
|
|
class SlurmSchedule(JobSchedule):
|
|
"""
|
|
job schedule for a slurm scheduler.
|
|
|
|
this class implements commonly used features of the slurm scheduler.
|
|
host-specific features and the creation of the job file should be done in a derived class.
|
|
derived classes must, in particular, implement the _write_job_file method.
|
|
they can override other methods, too, but should call the inherited method first.
|
|
|
|
1. copy the source trees (pmsco and projects) to the job directory
|
|
2. copy a patched version of the run file.
|
|
3. call the submission command
|
|
|
|
the public properties of this class should be assigned from the run file.
|
|
"""
|
|
def __init__(self, project):
|
|
super(SlurmSchedule, self).__init__(project)
|
|
self.host = ""
|
|
self.nodes = 1
|
|
self.tasks = 8
|
|
self.wall_time = datetime.timedelta(hours=1)
|
|
self.signal_time = 600
|
|
|
|
@staticmethod
|
|
def parse_timedelta(td):
|
|
"""
|
|
parse time delta input formats
|
|
|
|
converts a string or dictionary from run-file into datetime.timedelta.
|
|
|
|
@param td:
|
|
str: [days-]hours[:minutes[:seconds]]
|
|
dict: days, hours, minutes, seconds - at least one needs to be defined. values must be numeric.
|
|
datetime.timedelta - native type
|
|
@return: datetime.timedelta
|
|
"""
|
|
if isinstance(td, str):
|
|
dt = {}
|
|
d = td.split("-")
|
|
if len(d) > 1:
|
|
dt['days'] = float(d.pop(0))
|
|
t = d[0].split(":")
|
|
try:
|
|
dt['hours'] = float(t.pop(0))
|
|
dt['minutes'] = float(t.pop(0))
|
|
dt['seconds'] = float(t.pop(0))
|
|
except (IndexError, ValueError):
|
|
pass
|
|
td = datetime.timedelta(**dt)
|
|
elif isinstance(td, collections.abc.Mapping):
|
|
td = {k: float(v) for k, v in td.items()}
|
|
td = datetime.timedelta(**td)
|
|
return td
|
|
|
|
@property
|
|
def tasks_per_node(self):
|
|
return self.tasks // self.nodes
|
|
|
|
@tasks_per_node.setter
|
|
def tasks_per_node(self, value):
|
|
self.tasks = value * self.nodes
|
|
|
|
def validate(self):
|
|
super(SlurmSchedule, self).validate()
|
|
self.wall_time = self.parse_timedelta(self.wall_time)
|
|
|
|
def submit(self):
|
|
"""
|
|
call the sbatch command
|
|
|
|
if manual is true, the job files are generated but the job is not submitted.
|
|
|
|
@return: None
|
|
"""
|
|
super(SlurmSchedule, self).submit()
|
|
args = ['sbatch', str(self.job_file)]
|
|
print(" ".join(args))
|
|
if self.manual:
|
|
print("manual run - job files created but not submitted")
|
|
else:
|
|
cp = subprocess.run(args)
|
|
cp.check_returncode()
|
|
|
|
|
|
class PsiRaSchedule(SlurmSchedule):
|
|
"""
|
|
job shedule for the Ra cluster at PSI.
|
|
|
|
this class selects specific features of the Ra cluster,
|
|
such as the partition and node type.
|
|
it also implements the _write_job_file method.
|
|
|
|
for information about the Ra cluster, see
|
|
https://www.psi.ch/en/photon-science-data-services/offline-computing-facility-for-sls-and-swissfel-data-analysis
|
|
|
|
COMPUTE NODES
|
|
|
|
| NodeName | Weight | CPUs | RealMemory | MemSpecLimit | Sockets | CoresPerSocket | ThreadsPerCore |
|
|
| --- | :---: | :---: | :---: | :---: | :---: | :---: | :---: |
|
|
| ra-c-[033-048] | 100 | 36 | 256000 | 21502 | 2 | 18 | 1 |
|
|
| ra-c-[049-072] | 100 | 40 | 386000 | 21502 | 2 | 20 | 1 |
|
|
| ra-c-[073-084] | | 52 | | | 2 | 26 | 1 |
|
|
| ra-c-[085-096] | | 56 | | | 2 | 28 | 1 |
|
|
|
|
PARTITIONS
|
|
|
|
| PartitionName | Nodes | MaxTime | DefaultTime | Priority |
|
|
| --- | :---: | :---: | :---: | :---: | :---: |
|
|
| hour | | 0-01:00:00 | 0-01:00:00 | 4 |
|
|
| day | | 1-00:00:00 | 0-08:00:00 | 2 |
|
|
| week | | 8-00:00:00 | 2-00:00:00 | 1 |
|
|
|
|
The option --ntasks-per-node is meant to be used with the --nodes option.
|
|
(For the --ntasks option, the default is one task per node, use the --cpus-per-task option to change this default.)
|
|
"""
|
|
|
|
## @var partition (str)
|
|
#
|
|
# the partition is selected based on wall time and number of tasks by the validate() method.
|
|
# it should not be listed in the run file.
|
|
|
|
## @var modules (list of str)
|
|
#
|
|
# names of the software modules to load (`module load` command)
|
|
|
|
def __init__(self, project):
|
|
super(PsiRaSchedule, self).__init__(project)
|
|
self.partition = ""
|
|
self.modules = []
|
|
self.default_env = {}
|
|
|
|
def validate(self):
|
|
"""
|
|
check validity of parameters and detect the environment
|
|
|
|
the validity is not checked in detail - just intercept some common or severe mistakes.
|
|
|
|
detect whether we are in a conda or virtual environment
|
|
so that we can add the necessary activation to the job script.
|
|
|
|
detect which modules are active. depending on the python environment, LOADEDMODULES is:
|
|
- venv: openssl/3.4.1:TclTk/8.6.16:xz/5.8.0:Python/3.12.9:gcc/9.3.0:libfabric/1.18.0:cuda/11.1.0:openmpi/4.0.5_slurm
|
|
- conda: miniforge/2025-03-25:gcc/9.3.0:libfabric/1.18.0:cuda/11.1.0:openmpi/4.0.5_slurm
|
|
|
|
other potentially useful but currently unused environment variables are:
|
|
- PMODULES_LOADED_COMPILER: openmpi/4.0.5_slurm
|
|
- PMODULES_LOADED_TOOLS: openssl/3.4.1:xz/5.8.0
|
|
- PMODULES_LOADED_LIBRARIES: libfabric/1.18.0
|
|
- PMODULES_LOADED_PROGRAMMING: TclTk/8.6.16:Python/3.12.9:gcc/9.3.0:cuda/11.1.0
|
|
- PMODULES_LOADED_TOMCAT: miniforge/2025-03-25
|
|
|
|
note: openssl, TclTk, xz and cuda were not requested explicitly.
|
|
"""
|
|
|
|
super(PsiRaSchedule, self).validate()
|
|
|
|
# check that the submission is sane - the values are not firm
|
|
assert self.nodes <= 2
|
|
assert self.tasks <= 64
|
|
assert 30 * 60 <= self.wall_time.total_seconds() <= 8 * 24 * 60 * 60
|
|
|
|
if self.wall_time.total_seconds() > 24 * 60 * 60:
|
|
self.partition = "week"
|
|
else:
|
|
self.partition = "day"
|
|
|
|
assert self.partition in ["day", "week"]
|
|
|
|
self.default_env = self.detect_env()
|
|
|
|
if len(self.modules) == 0:
|
|
self.modules = os.environ["LOADEDMODULES"].split(":")
|
|
|
|
def _write_job_file(self):
|
|
"""
|
|
write a job file for the ra cluster
|
|
|
|
@return: None
|
|
"""
|
|
lines = []
|
|
|
|
lines.append('#!/bin/bash')
|
|
lines.append('#SBATCH --export=NONE')
|
|
lines.append(f'#SBATCH --job-name="{self.project.job_name}"')
|
|
lines.append(f'#SBATCH --partition={self.partition}')
|
|
lines.append(f'#SBATCH --time={int(self.wall_time.total_seconds() / 60)}')
|
|
lines.append(f'#SBATCH --nodes={self.nodes}')
|
|
lines.append(f'#SBATCH --ntasks-per-node={self.tasks_per_node}')
|
|
# 0 - 65535 seconds
|
|
# currently, PMSCO does not react to signals properly
|
|
# lines.append(f'#SBATCH --signal=TERM@{self.signal_time}')
|
|
lines.append(f'#SBATCH --output="{self.project.job_name}.o.%j"')
|
|
lines.append(f'#SBATCH --error="{self.project.job_name}.e.%j"')
|
|
|
|
# environment
|
|
if self.modules:
|
|
lines.append('module use unstable || true')
|
|
lines.append('module use Libraries || true')
|
|
for module in self.modules:
|
|
lines.append(f'module load {module}')
|
|
|
|
conda_env = self.conda_env or self.default_env.get("conda")
|
|
virtual_env = self.virtual_env or self.default_env.get("venv")
|
|
if conda_env:
|
|
lines.append(f'source activate {conda_env}')
|
|
elif virtual_env:
|
|
activate_script = Path(virtual_env) / "activate"
|
|
if activate_script.is_file():
|
|
lines.append(f'source {activate_script}')
|
|
|
|
lines.append('env')
|
|
lines.append('')
|
|
|
|
# run
|
|
lines.append(f'cd "{self.job_dir}"')
|
|
lines.append(f'mpirun python -m pmsco -r {self.run_file}')
|
|
lines.append('')
|
|
|
|
# clean up
|
|
lines.append(f'cd "{self.job_dir}"')
|
|
lines.append('exit 0')
|
|
|
|
self.job_file.write_text("\n".join(lines))
|
|
self.job_file.chmod(0o755)
|