diff --git a/docs/source/aplose.rst b/docs/source/aplose.rst index 646f60183..eba45bb84 100644 --- a/docs/source/aplose.rst +++ b/docs/source/aplose.rst @@ -38,7 +38,7 @@ The :class:`osekit.core.detection.Detection` class inherits from the :class:`ose Plotting a detection ^^^^^^^^^^^^^^^^^^^^ -Detection boxes can be plotted on spectrograms thanks to the :method:`osekit.core.detection.Detection.to_rectangle` method: +Detection boxes can be plotted on spectrograms thanks to the :meth:`osekit.core.detection.Detection.to_rectangle` method: .. code-block:: python diff --git a/docs/source/job.rst b/docs/source/job.rst index 9cd6c3854..6a01191bf 100644 --- a/docs/source/job.rst +++ b/docs/source/job.rst @@ -3,5 +3,23 @@ Job --- -.. automodule:: osekit.utils.job - :members: JobConfig, JobBuilder, Job +.. automodule:: osekit.job.job + :members: Job + +Job Config +---------- + +.. automodule:: osekit.job.config + :members: JobConfig + +Job Scheduler +------------- + +.. automodule:: osekit.job.scheduler.pbs + :members: Pbs + +Job Builder +----------- + +.. automodule:: osekit.job.builder + :members: JobBuilder diff --git a/docs/source/jobs.rst b/docs/source/jobs.rst index ded25ee6f..2b5ef43c4 100644 --- a/docs/source/jobs.rst +++ b/docs/source/jobs.rst @@ -2,51 +2,65 @@ Working with jobs ----------------- **OSEkit** can be set to send transform instructions to be computed on a remote server -through the PBS queuing system. +through queuing systems. This feature has mainly be thought for the Public API, but it can nonetheless be used for any Core API operation. -The job module is located at :mod:`osekit.utils.job`. +This is done thanks to the job package, located in :mod:`osekit.job`. Public API ^^^^^^^^^^ -Running Public API Analyses through PBS jobs only requires adding a :class:`osekit.utils.job.JobBuilder` -instance to the :attr:`osekit.public.project.Project.job_builder` attribute: +Running Public API Analyses through jobs only requires adding a :class:`osekit.job.builder.JobBuilder` +instance to the :attr:`osekit.public.project.Project.job_builder` attribute. + +The :class:`osekit.job.builder.JobBuilder` + +Here is an example for running a transform on a PBS queue: .. code-block:: python - from osekit.utils.job import JobConfig, JobBuilder + import os + + from pandas import Timedelta + + from osekit.job.builder import JobBuilder + from osekit.job.config import JobConfig + from osekit.job.scheduler.pbs import Pbs from osekit.public.project import Project - project = Project(...) # See the Project documentation + project = Project(...) # See the Project documentation job_config = JobConfig( - nb_nodes=1, # Number of nodes on which the job runs - ncpus=28, # Number of total cores used per node - ngpus=1, # Number of total GPU used per node - mem="60gb", # Maximum amount of physical memory used by the job - walltime=Timedelta(hours=5), # Maximum amount of real itime during which the job can be running - venv_name=os.environ["CONDA_DEFAULT_ENV"], # Works only for conda venvs - queue="omp" # Queue in which the job will be submitted + nb_nodes=1, # Number of nodes on which the job runs + ncpus=28, # Number of total cores used per node + ngpus=1, # Number of total GPU used per node + mem="60gb", # Maximum amount of physical memory used by the job + walltime=Timedelta( + hours=5 + ), # Maximum amount of real itime during which the job can be running + venv_name=os.environ["CONDA_DEFAULT_ENV"], # Works only for conda venvs ) + scheduler = Pbs(queue="omp") # Scheduler in which the job is submitted + project.job_builder = JobBuilder( config=job_config, + scheduler=scheduler, ) # Now the dataset has a non-None job_builder attribute, # running a transform will write a PBS file in the logs directory - # and submit it to the requested queue. + # and submit it through the selected scheduler. - project.run(...) # See the Transform documentation + project.run(...) # See the Transform documentation Core API ^^^^^^^^ -Exporting Core API datasets with jobs is doable by explicitly instantiating a :class:`osekit.utils.job.Job` object. +Exporting Core API datasets with jobs is doable by explicitly instantiating a :class:`osekit.job.job.Job` object. The export parameters are specified in the ``script_args`` parameter of the ``Job`` constructor, and follow the console arguments of the :mod:`osekit.public.export` script. @@ -60,11 +74,13 @@ and follow the console arguments of the :mod:`osekit.public.export` script. from osekit.core.audio_dataset import AudioDataset from osekit.core.spectro_dataset import SpectroDataset + from osekit.job.config import JobConfig + from osekit.job.job import Job + from osekit.job.scheduler.pbs import Pbs from osekit.public import export_transform # Some Public API imports are required from osekit.public.transform import OutputType - from osekit.utils.job import Job, JobConfig ads = AudioDataset(...) # See the AudioDataset doc sds = SpectroDataset(...) # See the SpectroDataset doc @@ -109,9 +125,11 @@ and follow the console arguments of the :mod:`osekit.public.export` script. mem="60gb", walltime=Timedelta(hours=1), venv_name=os.environ["CONDA_DEFAULT_ENV"], - queue="omp", ) + # Scheduler configuration + scheduler = Pbs(queue="omp") + job = Job( script_path=Path(export_transform.__file__), script_args=args, @@ -120,12 +138,12 @@ and follow the console arguments of the :mod:`osekit.public.export` script. output_folder=Path(...), # Path in which the .out and .err files are written ) - # Write the PBS file and submit the job - job.write_pbs(Path(...) / f"{job.name}.pbs") - job.submit_pbs() + # Write the job file and submit it through the scheduler + scheduler.write(job=job, path=Path(...) / f"{job.name}.pbs") + scheduler.submit(job=job) -You can then follow the status of the submitted job: +You can then follow the status of the submitted job through the scheduler: .. code-block:: python - job.update_status() + scheduler.update_status(job=job) diff --git a/src/osekit/job/__init__.py b/src/osekit/job/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/osekit/job/builder.py b/src/osekit/job/builder.py new file mode 100644 index 000000000..2e083286e --- /dev/null +++ b/src/osekit/job/builder.py @@ -0,0 +1,153 @@ +from pathlib import Path + +from osekit.job.config import JobConfig +from osekit.job.job import Job, JobStatus +from osekit.job.scheduler.pbs import Pbs +from osekit.job.scheduler.scheduler import Scheduler +from osekit.utils.core import file_indexes_per_batch + + +class JobBuilder: + """Class that should be attached to a Public API ``Project`` for working with jobs. + + If a ``Project`` has a ``JobBuilder``, it will run its transforms through jobs + using the specified scheduler. + + """ + + def __init__( + self, + config: JobConfig | None = None, + scheduler: Scheduler | None = None, + ) -> None: + """Initialize a ``JobBuilder`` instance. + + Parameters + ---------- + config: JobConfig + Config of the jobs built by this job builder. + scheduler: Scheduler + Scheduler used to format, write and submit jobs. + + """ + self.config = config or JobConfig() + self.scheduler = scheduler or Pbs() + self.jobs = [] + + @property + def jobs(self) -> list[Job]: + """Return the jobs created by this job builder.""" + return self._jobs + + @jobs.setter + def jobs(self, jobs: list[Job]) -> None: + self._jobs = jobs + + def create_jobs( + self, + nb_tasks: int, + script_path: Path, + script_args: dict, + output_folder: Path, + job_name: str = "osekit_transform", + nb_jobs: int = 1, + ) -> None: + """Create the jobs corresponding to each batch. + + Parameters + ---------- + nb_tasks: + The number of tasks that are distributed across ``nb_jobs`` jobs. + script_path: Path + Path to the export script. + script_args: dict + Arguments passed to the export script. + job_name: str + Name of the job. + If there are multiple batches, each batch will be suffixed + with "_{index}". + output_folder: Path + Folder in which the job output log files are saved. + nb_jobs: int + Number of batches used to run the transform. + Each batch will run in a separate job. + + """ + batch_indexes = file_indexes_per_batch( + total_nb_files=nb_tasks, + nb_batches=nb_jobs, + ) + for index, (start, stop) in enumerate(batch_indexes): + self.create_job( + script_path=script_path, + script_args=script_args | {"first": start, "last": stop}, + name=job_name + (f"_{index}" if len(batch_indexes) > 1 else ""), + output_folder=output_folder, + ) + + def create_job( + self, + script_path: Path, + script_args: dict | None = None, + name: str = "osekit_transform", + output_folder: Path | None = None, + ) -> None: + """Create a new ``Job`` instance. + + Parameters + ---------- + script_path: Path + Path to the script file the job must run. + script_args: dict | None + Additional arguments to pass to the script file. + name: str + Name of the job. + output_folder: Path | None + Folder in which the output files (``.out`` and ``.err``) will be written. + + """ + job = Job( + script_path=script_path, + script_args=script_args, + name=name, + output_folder=output_folder, + config=self.config, + ) + self.scheduler.write( + job=job, + path=output_folder / f"{name}.{self.scheduler.JOB_FILE_EXTENSION}", + ) + self.jobs.append(job) + + def submit( + self, + dependencies: dict[Job, dict[str, str | Job | list[str | Job]]] | None = None, + ) -> None: + """Submit all prepared jobs to the scheduler system. + + Parameters + ---------- + dependencies: dict[Job, dict[str, str | Job | list[str | Job]]] | None + Optional mapping of the jobs dependencies. + + For each key (which is a job), the value is a + dictionary explaining its dependencies. + + Such dictionary follows the following format: + The keys of the dictionary are the dependency types, + that are proper to the scheduler. + The values are the other jobs (or their ID) ``job`` depends on + with the given dependency type. + + If ``None``, the jobs are submitted without any dependency. + + """ + dependencies = dependencies or {} + for job in self.jobs: + if self.scheduler.update_status(job=job) is not JobStatus.PREPARED: + continue + + # Check if this job has dependencies + job_dependencies = dependencies.get(job, None) + + self.scheduler.submit(job=job, dependencies=job_dependencies) diff --git a/src/osekit/job/config.py b/src/osekit/job/config.py new file mode 100644 index 000000000..e7552bfa6 --- /dev/null +++ b/src/osekit/job/config.py @@ -0,0 +1,32 @@ +from dataclasses import dataclass + +from pandas import Timedelta + + +@dataclass +class JobConfig: + """Configuration of the computing resources allowed for a job. + + Parameters + ---------- + nb_nodes: int + Number of nodes on which the job runs. + ncpus: int + Number of total cores used per node. + ngpus: int | None + Number of total GPU used per node. + mem: str + Maximum amount of physical memory used by the job. + walltime: str | Timedelta + Maximum amount of real time during which the job can be running. + venv_name: str + Name (or path) of the conda virtual environment in which the job is running. + + """ + + nb_nodes: int = 1 + ncpus: int = 2 + ngpus: int | None = None + mem: str = "8gb" + walltime: str | Timedelta = "01:00:00" + venv_name: str = "osekit" diff --git a/src/osekit/job/job.py b/src/osekit/job/job.py new file mode 100644 index 000000000..a55cbbe73 --- /dev/null +++ b/src/osekit/job/job.py @@ -0,0 +1,237 @@ +"""The job module provides classes that run transforms on a remote server. + +If a ``JobBuilder`` is attached to a Public API ``Project``, +the transforms will run through jobs, with writting/submitting of ``pbs`` files. + +""" + +from __future__ import annotations + +from enum import Enum +from typing import TYPE_CHECKING + +from pandas import Timedelta + +from osekit.job.config import JobConfig + +if TYPE_CHECKING: + from pathlib import Path + + +class JobStatus(Enum): + """Status of the job. + + ``UNPREPARED``: The job file hasn't been written yet. + ``PREPARED``: The job file has been written but not submitted. + ``QUEUED``: The job has been queued. + ``RUNNING``: The job is currently running. + ``SUSPENDED``: The job has been suspended or is held. + ``COMPLETED``: The job is exiting or has been completed. + + """ + + UNPREPARED = 1 + PREPARED = 2 + QUEUED = 3 + RUNNING = 4 + SUSPENDED = 5 + COMPLETED = 6 + + +class Job: + """Job that concerns a specific transform.""" + + def __init__( + self, + script_path: Path, + script_args: dict | None = None, + config: JobConfig | None = None, + name: str = "osekit_transform", + output_folder: Path | None = None, + ) -> None: + """Initialize a Job. + + Parameters + ---------- + script_path: Path + Path to the script file the job must run. + script_args: dict | None + Additional arguments to pass to the script file. + config: JobConfig | None + Optional configuration to pass to the server request. + name: str + Name of the job. + output_folder: Path | None + Folder in which the output files (``.out`` and ``.err``) will be written. + + """ + config = JobConfig() if config is None else config + self.script_path = script_path + self.script_args = script_args or {} + self.nb_nodes = config.nb_nodes + self.ncpus = config.ncpus + self.ngpus = config.ngpus + self.mem = config.mem + self.walltime = config.walltime + self.venv_name = config.venv_name + self.name = name + self.output_folder = output_folder + self.info = {} + self._status = JobStatus.UNPREPARED + self._path = None + self._id = None + + @property + def script_path(self) -> Path: + """Path to the script file the job must run.""" + return self._script_path + + @script_path.setter + def script_path(self, path: Path) -> None: + self._script_path = path + + @property + def script_args(self) -> dict: + """Additional arguments to pass to the script file.""" + return self._script_args + + @script_args.setter + def script_args(self, args: dict) -> None: + self._script_args = args + + @property + def nb_nodes(self) -> int: + """Number of nodes on which the job runs.""" + return self._chunks + + @nb_nodes.setter + def nb_nodes(self, chunks: int) -> None: + self._chunks = chunks + + @property + def ncpus(self) -> int: + """Number of total cores used per node.""" + return self._ncpus + + @ncpus.setter + def ncpus(self, ncpus: int) -> None: + self._ncpus = ncpus + + @property + def ngpus(self) -> int | None: + """Number of total GPU used per node.""" + return self._ngpus + + @ngpus.setter + def ngpus(self, ngpus: int) -> None: + self._ngpus = ngpus + + @property + def mem(self) -> str: + """Maximum amount of physical memory used by the job.""" + return self._mem + + @mem.setter + def mem(self, mem: str) -> None: + self._mem = mem + + @property + def walltime(self) -> Timedelta: + """Maximum amount of real time during which the job can be running.""" + return self._walltime + + @property + def walltime_str(self) -> str: + """String representation of the ``walltime``.""" + total_seconds = self.walltime.total_seconds() + hours, remainder = divmod(total_seconds, 3600) + minutes, seconds = divmod(remainder, 60) + return ":".join(f"{t:02}" for t in map(int, (hours, minutes, seconds))) + + @walltime.setter + def walltime(self, walltime: str | Timedelta) -> None: + self._walltime = ( + walltime if type(walltime) is Timedelta else Timedelta(walltime) + ) + + @property + def venv_name(self) -> str: + """Name of the conda virtual environment in which the job is running.""" + return self._venv_name + + @venv_name.setter + def venv_name(self, venv_name: str) -> None: + self._venv_name = venv_name + + @property + def name(self) -> str: + """Name of the job.""" + return self._name + + @name.setter + def name(self, name: str) -> None: + self._name = name + + @property + def status(self) -> JobStatus: + """Status of the job. + + ``UNPREPARED``: The job file hasn't been written yet. + ``PREPARED``: The job file has been written but not submitted. + ``QUEUED``: The job has been queued. + ``RUNNING``: The job is currently running. + ``SUSPENDED``: The job has been suspended or is held. + ``COMPLETED``: The job is exiting or has been completed. + + """ + return self._status + + @status.setter + def status(self, status: JobStatus) -> None: + self._status = status + + @property + def path(self) -> Path | None: + """Path of the job file.""" + return self._path + + @path.setter + def path(self, path: Path) -> None: + self._path = path + + @property + def output_folder(self) -> Path | None: + """Folder in which the output files (``.out`` and ``.err``) will be written.""" + return self._output_folder + + @output_folder.setter + def output_folder(self, output_folder: Path | None) -> None: + self._output_folder = output_folder + + @property + def job_id(self) -> str | None: + """Job ID.""" + return self._id + + @job_id.setter + def job_id(self, job_id: str | None) -> None: + self._id = job_id + + @property + def info(self) -> dict: + """Information about the job.""" + return self._info + + @info.setter + def info(self, info: dict) -> None: + self._info = info + + def get_arg_string(self) -> str: + """Build a string representation of the job's arguments.""" + arg_list = [] + for key, value in self.script_args.items(): + if isinstance(value, bool): + arg_list.append(f"--{'no-' if not value else ''}{key}") + else: + arg_list.append(f"--{key} {value}") + return " ".join(arg_list) diff --git a/src/osekit/job/scheduler/__init__.py b/src/osekit/job/scheduler/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/osekit/job/scheduler/pbs.py b/src/osekit/job/scheduler/pbs.py new file mode 100644 index 000000000..64a85c350 --- /dev/null +++ b/src/osekit/job/scheduler/pbs.py @@ -0,0 +1,209 @@ +import subprocess +import typing +from typing import Literal + +from osekit.job.job import Job, JobStatus +from osekit.job.scheduler.scheduler import Scheduler + + +class Pbs(Scheduler): + """Abstract class representing a PBS job scheduler.""" + + _VALID_DEPENDENCY_TYPES: typing.ClassVar = frozenset( + { + "after", + "afterok", + "afternotok", + "afterany", + "before", + "beforeok", + "beforenotok", + "beforeany", + "on", + "runone", + }, + ) + JOB_FILE_EXTENSION: typing.ClassVar = "pbs" + + JOB_STATUS_CODES: typing.ClassVar = { + "Q": JobStatus.QUEUED, + "R": JobStatus.RUNNING, + "S": JobStatus.SUSPENDED, + "H": JobStatus.SUSPENDED, + "E": JobStatus.COMPLETED, + "F": JobStatus.COMPLETED, + } + + SUBMIT_CMD: typing.ClassVar = "qsub" + + def __init__(self, queue: Literal["omp", "mpi"] = "omp") -> None: + """Initialize the PBS scheduler.""" + self.queue = queue + + @property + def queue(self) -> str: + """Queue in which the job will be submitted.""" + return self._queue + + @queue.setter + def queue(self, queue: Literal["omp", "mpi"]) -> None: + self._queue = queue + + def _build_job_specification(self, job: Job) -> str: + """Build the job specification string. + + Parameters + ---------- + job: Job + The job for which to build the specifications. + + Returns + ------- + str: + Job specification string. + It includes the name of the job, the requested resources, + output log directories, etc. + + """ + select_parts = { + "select": job.nb_nodes, + "ncpus": job.ncpus, + "mem": job.mem, + } + if job.ngpus is not None: + select_parts["ngpus"] = job.ngpus + select_str = ":".join(f"{k}={v}" for k, v in select_parts.items()) + + request = { + "-N": job.name, + "-q": self.queue, + "-l": [ + select_str, + f"walltime={job.walltime_str}", + ], + "-o": f"{job.output_folder / job.name}.out" if job.output_folder else None, + "-e": f"{job.output_folder / job.name}.err" if job.output_folder else None, + } + return "\n".join( + f"#PBS {key} {value}" + if type(value) is not list + else "\n".join(f"#PBS {key} {value_part}" for value_part in value) + for key, value in request.items() + if value + ) + + @staticmethod + def _parse_job_id(submit_output: str) -> str: + """Parse the output of the submit command.""" + return submit_output.split(".", maxsplit=1)[0].strip() + + @staticmethod + def _get_info(job: Job) -> str: + """Request information about a job. + + Parameters + ---------- + job: Job + Job for which the information is requested. + + Returns + ------- + str: + The information string, as returned by the ``qstat -f`` command. + + """ + try: + request = subprocess.run( + ["qstat", "-f", str(job.job_id)], + capture_output=True, + text=True, + check=False, + ) + if request.stdout: + return request.stdout + + except subprocess.CalledProcessError as e: + msg = f"qstat failed with exit code {e.returncode}" + raise RuntimeError(msg) from e + + err = request.stderr + if err: + msg = f"{job.job_id}: {err}" + raise ValueError(msg) + + return "" + + @classmethod + def _parse_info_str(cls, job: Job, info: str) -> None: + """Parse the info from the requested qstat info string.""" + kvp_lines = (line for line in info.splitlines() if "=" in line) + kvp_lines = dict(line.split("=", maxsplit=1) for line in kvp_lines) + kvp_lines = {k.strip(): v.strip() for k, v in kvp_lines.items()} + + job.info["user"] = kvp_lines["Job_Owner"].split("@", maxsplit=1)[0] + job.info["time"] = kvp_lines.get("resources_used.walltime", "00:00:00") + job.info["queue"] = kvp_lines["queue"] + + if status := cls.JOB_STATUS_CODES.get(kvp_lines["job_state"], False): + job.status = status + + @staticmethod + def _build_venv_string(job: Job) -> str: + """Bash script used for activating the conda virtual environment.""" + return ( + f". /appli/anaconda/latest/etc/profile.d/conda.sh\n" + f"conda activate {job.venv_name}" + ) + + @classmethod + def _build_dependency_string( + cls, + dependencies: dict[str, Job | str | list[Job | str]], + ) -> str: + """Build a PBS dependency string. + + Parameters + ---------- + dependencies: dict[str, Job | str | list[Job|str]] + The dependencies of the submitted job. + The keys of the dictionary are the dependency types, + see https://help.altair.com/2022.1.0/PBS%20Professional/PBSReferenceGuide2022.1.pdf#page=151 + for the list of supported values. + The values are the other jobs (or their ID) ``job`` depends on + with the given dependency type. + + Returns + ------- + str + PBS dependency string. + + Examples + -------- + >>> Pbs._build_dependency_string({"afterok": "1234567"}) + '-W depend=afterok:1234567' + >>> Pbs._build_dependency_string({"afterok": ["1234567","4567891"]}) + '-W depend=afterok:1234567:4567891' + >>> from pathlib import Path + >>> job = Job(Path()) + >>> job._id = "7894561" + >>> Pbs._build_dependency_string({"afterany":job}) + '-W depend=afterany:7894561' + >>> from pathlib import Path + >>> job1 = Job(Path()) + >>> job1._id = "7894561" + >>> job2 = Job(Path()) + >>> job2._id = "4839572" + >>> Pbs._build_dependency_string({"afterany":[job1,job2]}) + '-W depend=afterany:7894561:4839572' + + """ + # Check that types are valid before submitting + for dependency_type in dependencies: + cls._validate_dependency_type(dependency_type=dependency_type) + + id_str = cls._parse_job_ids(dependencies=dependencies) + + return "-W depend=" + ",".join( + f"{dependency_type}:{':'.join(ids)}" + for dependency_type, ids in id_str.items() + ) diff --git a/src/osekit/job/scheduler/scheduler.py b/src/osekit/job/scheduler/scheduler.py new file mode 100644 index 000000000..d60c491a1 --- /dev/null +++ b/src/osekit/job/scheduler/scheduler.py @@ -0,0 +1,238 @@ +import subprocess +import typing +from abc import ABC, abstractmethod +from pathlib import Path + +from osekit.job.job import Job, JobStatus + + +class Scheduler(ABC): + """Abstract class representing a job scheduler.""" + + _VALID_DEPENDENCY_TYPES: typing.ClassVar = frozenset() + JOB_FILE_EXTENSION: typing.ClassVar = "job" + SUBMIT_CMD: typing.ClassVar = "" + INFO_CMD: typing.ClassVar = [] + + def write(self, job: Job, path: Path) -> None: + """Write a job script to file. + + Parameters + ---------- + job: Job + Job of which to write the script. + path: Path + Path of the file in which the job script is written. + + """ + preamble = "#!/bin/bash" + + request_str = self._build_job_specification(job=job) + venv_str = self._build_venv_string(job=job) + python_script = f"python {job.script_path} {job.get_arg_string()}" + + script = f"{preamble}\n\n{request_str}\n\n{venv_str}\n\n{python_script}" + + with path.open("w") as file: + file.write(script) + + job.path = path + job.status = JobStatus.PREPARED + + @abstractmethod + def _build_job_specification(self, job: Job) -> str: + """Build the job specification string. + + Parameters + ---------- + job: Job + The job for which to build the specifications. + + Returns + ------- + str: + Job specification string. + It includes the name of the job, the requested resources, + output log directories, etc. + + """ + ... + + def submit( + self, + job: Job, + dependencies: dict[str, Job | str | list[Job | str]] | None = None, + ) -> None: + """Submit the job to the scheduler. + + Parameters + ---------- + job: Job + Job to submit to the scheduler. + dependencies: dict[str, Job | str | list[Job|str]] + The dependencies of the submitted job. + The keys of the dictionary are the dependency types, + that are proper to the scheduler. + The values are the other jobs (or their ID) ``job`` depends on + with the given dependency type. + If ``None``, the job is submitted without any dependency. + + """ + if self.update_status(job=job) is not JobStatus.PREPARED: + msg = "Job should be written before being submitted." + raise ValueError(msg) + + cmd = [self.SUBMIT_CMD] + + if dependencies: + cmd.extend(self._build_dependency_string(dependencies=dependencies).split()) + + cmd.append(str(job.path)) + + try: + request = subprocess.run( + cmd, + capture_output=True, + text=True, + check=False, + ) + except subprocess.CalledProcessError as e: + msg = f"Submission failed with exit code {e.returncode}" + raise RuntimeError(msg) from e + + job.job_id = self._parse_job_id(submit_output=request.stdout) + self.update_status(job=job) + + @staticmethod + @abstractmethod + def _parse_job_id(submit_output: str) -> str: + """Parse the job id from the submit cmd output.""" + ... + + def update_info(self, job: Job) -> None: + """Request info about the job and update it. + + Parameters + ---------- + job: Job + Job for which to update the info. + + """ + if job.job_id is None: + return + + info = self._get_info(job=job) + if not info: + return + + self._parse_info_str(job=job, info=info) + + @staticmethod + @abstractmethod + def _get_info(job: Job) -> str: + """Request information about a job. + + Parameters + ---------- + job: Job + Job for which the information is requested. + + Returns + ------- + str: + The information string. + Depends on the scheduler. + + """ + ... + + @classmethod + @abstractmethod + def _parse_info_str(cls, job: Job, info: str) -> None: + """Parse the info string to update a job info. + + Parameters + ---------- + job: Job + Job of which the info is updated + info: str + Info string as returned by a request. + + """ + ... + + def update_status(self, job: Job) -> JobStatus: + """Request info about the job and update its status. + + Returns + ------- + JobStatus: + The updated status of the job. + + """ + if job.job_id is None: + job.status = ( + JobStatus.PREPARED + if job.path and job.path.exists() + else JobStatus.UNPREPARED + ) + return job.status + + self.update_info(job=job) + return job.status + + @staticmethod + @abstractmethod + def _build_venv_string(job: Job) -> str: ... + + @classmethod + def _validate_dependency_type(cls, dependency_type: str) -> None: + """Validate dependency types in the dependencies instruction.""" + if dependency_type not in cls._VALID_DEPENDENCY_TYPES: + msg = ( + f"Unsupported dependency type '{dependency_type}'.\n" + f"Expected one of:\n\t{'\n\t'.join(sorted(cls._VALID_DEPENDENCY_TYPES))}." + ) + raise ValueError(msg) + + @staticmethod + def _parse_job_ids( + dependencies: dict[str, Job | str | list[Job | str]], + ) -> dict[str, list[str]]: + """Replace all ``Job`` instances by their ID string.""" + parsed_dependencies = {} + for key, value in dependencies.items(): + parsed_values = value if isinstance(value, list) else [value] + parsed_values = [ + parsed_value.job_id if isinstance(parsed_value, Job) else parsed_value + for parsed_value in parsed_values + ] + parsed_dependencies[key] = parsed_values + + return parsed_dependencies + + @classmethod + @abstractmethod + def _build_dependency_string( + cls, + dependencies: dict[str, Job | str | list[Job | str]], + ) -> str: + """Build a job dependency string. + + Parameters + ---------- + dependencies: dict[str, Job | str | list[Job|str]] + The dependencies of the submitted job. + The keys of the dictionary are the dependency types, + that are proper to the scheduler. + The values are the other jobs (or their ID) ``job`` depends on + with the given dependency type. + If ``None``, the job is submitted without any dependency. + + Returns + ------- + str + Job dependency string. + + """ + ... diff --git a/src/osekit/job/scheduler/slurm.py b/src/osekit/job/scheduler/slurm.py new file mode 100644 index 000000000..dfb65229f --- /dev/null +++ b/src/osekit/job/scheduler/slurm.py @@ -0,0 +1,272 @@ +import subprocess +import typing +from typing import Literal + +from osekit.job.job import Job, JobStatus +from osekit.job.scheduler.scheduler import Scheduler + + +class Slurm(Scheduler): + """Abstract class representing a job scheduler.""" + + _VALID_DEPENDENCY_TYPES: typing.ClassVar = frozenset( + { + "after", + "afterany", + "afterburstbuffer", + "aftercorr", + "afternotok", + "afterok", + "singleton", + }, + ) + JOB_FILE_EXTENSION: typing.ClassVar = "slurm" + SUBMIT_CMD: typing.ClassVar = "sbatch" + JOB_STATUS_CODES: typing.ClassVar = { + "PENDING": JobStatus.QUEUED, + "RUNNING": JobStatus.RUNNING, + "SUSPENDED": JobStatus.SUSPENDED, + "COMPLETING": JobStatus.RUNNING, + "COMPLETED": JobStatus.COMPLETED, + } + + def __init__(self, partition: Literal["cpu", "gpu", "ops"] = "cpu") -> None: + """Initialize the SLURM scheduler.""" + self.partition = partition + + @property + def partition(self) -> str: + """Partition in which the job will be submitted.""" + return self._partition + + @partition.setter + def partition(self, partition: Literal["omp", "mpi"]) -> None: + self._partition = partition + + def _build_job_specification(self, job: Job) -> str: + """Build the job specification string. + + Parameters + ---------- + job: Job + The job for which to build the specifications. + + Returns + ------- + str: + Job specification string. + It includes the name of the job, the requested resources, + output log directories, etc. + + """ + specifications = { + "nodes": job.nb_nodes, + "cpus-per-task": job.ncpus, + "mem": job.mem, + "job-name": job.name, + "partition": self.partition, + "time": job.walltime_str, + "output": f"{job.output_folder / job.name}.out" + if job.output_folder + else None, + "error": f"{job.output_folder / job.name}.err" + if job.output_folder + else None, + } + + if job.ngpus is not None: + specifications["gpus"] = job.ngpus + + return "\n".join( + f"#SBATCH --{key}={value}" for key, value in specifications.items() if value + ) + + @staticmethod + def _parse_job_id(submit_output: str) -> str: + r"""Parse the output of the submit command. + + Parameters + ---------- + submit_output: str + stdout after a successful sbatch cmd. + + Returns + ------- + str: + ID of the submitted job. + + Examples + -------- + >>> Slurm._parse_job_id(submit_output="Submitted batch job 3647090\n") + '3647090' + + """ + return submit_output.removeprefix("Submitted batch job ").strip("\n") + + @classmethod + def _get_info(cls, job: Job) -> str: + """Request information about a job. + + Parameters + ---------- + job: Job + Job for which the information is requested. + + Returns + ------- + str: + The information string, as returned by the ``squeue`` command. + If the job is complete and doesn't appear in the queue, + the info is requested through the ``sacct`` command, parsed + by a specific parser, and an empty string is returned. + + """ + request = subprocess.run( + [ + "/usr/bin/squeue", + "--jobs", + str(job.job_id), + "--noheader", + ( + "--format=" + '"%i|' # ID + "%P|" # Partition + "%j|" # Job name + "%u|" # User name + "%T|" # Status + "%M|" # Time + "%D" # Nodes nb + '|%R"' + ), # Nodes List (Reason) + ], + capture_output=True, + text=True, + check=False, + ) + + if info := request.stdout.strip(): + return info + + # The job doesn't appear in squeue anymore + request = subprocess.run( + [ + "/usr/bin/sacct", + "--jobs", + str(job.job_id), + "--allocations", + "--noheader", + "--parsable2", + "--format=JobID,Partition,JobName,User,State,Elapsed,NNodes,NodeList,Reason", + ], + capture_output=True, + text=True, + check=False, + ) + + if acct_info := request.stdout.strip(): + return acct_info + + if error := request.stderr.strip(): + msg = f"{job.job_id}: {error}" + raise ValueError(msg) + + return "" + + @classmethod + def _parse_info_str(cls, job: Job, info: str) -> None: + """Parse the info from the requested squeue info string.""" + values = info.strip().split("|", maxsplit=7) + + ( + _job_id, + partition, + _name, + user, + status, + time, + nodes, + node_list, + ) = values + + job.info["user"] = user + job.info["time"] = time + job.info["partition"] = partition + + # If info is fetched by sacct, node_list could contain + # both the node name and the reason, sill separated by a "|" + if "|" in node_list: + node, reason = node_list.split("|", maxsplit=1) + node_list = f"{node} ({reason})" + job.info["node_list"] = node_list + + if status := cls.JOB_STATUS_CODES.get(status, False): + job.status = status + + @staticmethod + def _build_venv_string(job: Job) -> str: + """Bash script used for activating the conda virtual environment.""" + return f"module load conda\nconda activate {job.venv_name}" + + @classmethod + def _build_dependency_string( + cls, + dependencies: dict[str, Job | str | list[Job | str]], + *, + instructions_or: bool = False, + ) -> str: + """Build a job dependency string. + + Parameters + ---------- + dependencies: dict[str, Job | str | list[Job|str]] + The dependencies of the submitted job. + The keys of the dictionary are the dependency types, + see https://slurm.schedmd.com/sbatch.html + The values are the other jobs (or their ID) ``job`` depends on + with the given dependency type. + If ``None``, the job is submitted without any dependency. + instructions_or: bool, optional + If ``True``, the instructions in the ``dependencies`` list + are joined with a logical OR (``?`` character in Slurm). + If ``False``, the instructions are joined with a logical + AND (``,`` character in Slurm). + + Returns + ------- + str + Job dependency string. + + Examples + -------- + >>> Slurm._build_dependency_string({"afterok": "1234567"}) + '-d afterok:1234567' + >>> Slurm._build_dependency_string({"afterok": ["1234567","4567891"]}) + '-d afterok:1234567:4567891' + >>> Slurm._build_dependency_string({"afterok": ["1234567","4567891"], "afterany":"7654321"}, instructions_or=True) + '-d afterok:1234567:4567891?afterany:7654321' + >>> from pathlib import Path + >>> job = Job(Path()) + >>> job._id = "7894561" + >>> Slurm._build_dependency_string({"afterany":job}) + '-d afterany:7894561' + >>> from pathlib import Path + >>> job1 = Job(Path()) + >>> job1._id = "7894561" + >>> job2 = Job(Path()) + >>> job2._id = "4839572" + >>> Slurm._build_dependency_string({"afterany":[job1,job2]}) + '-d afterany:7894561:4839572' + + """ # noqa: E501 + # Check that types are valid before submitting + for dependency_type in dependencies: + cls._validate_dependency_type(dependency_type=dependency_type) + + id_str = cls._parse_job_ids(dependencies=dependencies) + + logical_join_character = "?" if instructions_or else "," + + return "-d " + logical_join_character.join( + f"{dependency_type}:{':'.join(ids)}" + for dependency_type, ids in id_str.items() + ) diff --git a/src/osekit/public/project.py b/src/osekit/public/project.py index ff9fff0ce..aef0db689 100644 --- a/src/osekit/public/project.py +++ b/src/osekit/public/project.py @@ -38,7 +38,7 @@ from pandas import Timestamp from osekit.core.audio_file import AudioFile - from osekit.utils.job import JobBuilder + from osekit.job.builder import JobBuilder class Project: @@ -634,7 +634,7 @@ def export( nb_jobs=nb_jobs, ) - self.job_builder.submit_pbs() + self.job_builder.submit() @staticmethod def get_json_paths( diff --git a/src/osekit/utils/job.py b/src/osekit/utils/job.py deleted file mode 100644 index 63494ae85..000000000 --- a/src/osekit/utils/job.py +++ /dev/null @@ -1,635 +0,0 @@ -"""The job module provides classes that run transforms on a remote server. - -If a ``JobBuilder`` is attached to a Public API ``Project``, -the transforms will run through jobs, with writting/submitting of ``pbs`` files. - -""" - -from __future__ import annotations - -import subprocess -from dataclasses import dataclass -from enum import Enum -from typing import TYPE_CHECKING, Literal - -from pandas import Timedelta - -from osekit.utils.core import file_indexes_per_batch - -if TYPE_CHECKING: - from pathlib import Path - - -class JobStatus(Enum): - """Status of the job. - - ``UNPREPARED``: The job file hasn't been written yet. - ``PREPARED``: The job file has been written but not submitted. - ``QUEUED``: The job has been queued. - ``RUNNING``: The job is currently running. - ``COMPLETED``: The job has been completed. - - """ - - UNPREPARED = 1 - PREPARED = 2 - QUEUED = 3 - RUNNING = 4 - COMPLETED = 5 - - -@dataclass -class JobConfig: - """Config used for creating a job. - - Parameters - ---------- - nb_nodes: int - Number of nodes on which the job runs. - ncpus: int - Number of total cores used per node. - ngpus: int | None - Number of total GPU used per node. - mem: str - Maximum amount of physical memory used by the job. - walltime: str | Timedelta - Maximum amount of real time during which the job can be running. - venv_name: str - Name (or path) of the conda virtual environment in which the job is running. - queue: Literal["omp", "mpi"] - Queue in which the job will be submitted. - - """ - - nb_nodes: int = 1 - ncpus: int = 2 - ngpus: int | None = None - mem: str = "8gb" - walltime: str | Timedelta = "01:00:00" - venv_name: str = "osekit" - queue: Literal["omp", "mpi"] = "omp" - - -class Job: - """Job that concerns a specific transform.""" - - def __init__( - self, - script_path: Path, - script_args: dict | None = None, - config: JobConfig | None = None, - name: str = "osekit_transform", - output_folder: Path | None = None, - ) -> None: - """Initialize a Job. - - Parameters - ---------- - script_path: Path - Path to the script file the job must run. - script_args: dict | None - Additional arguments to pass to the script file. - config: JobConfig | None - Optional configuration to pass to the server request. - name: str - Name of the job. - output_folder: Path | None - Folder in which the output files (``.out`` and ``.err``) will be written. - - """ - config = JobConfig() if config is None else config - self.script_path = script_path - self.script_args = script_args or {} - self.nb_nodes = config.nb_nodes - self.ncpus = config.ncpus - self.ngpus = config.ngpus - self.mem = config.mem - self.walltime = config.walltime - self.venv_name = config.venv_name - self.queue = config.queue - self.name = name - self.output_folder = output_folder - self._status = JobStatus.UNPREPARED - self._path = None - self._id = None - self._info = None - - @property - def script_path(self) -> Path: - """Path to the script file the job must run.""" - return self._script_path - - @script_path.setter - def script_path(self, path: Path) -> None: - self._script_path = path - - @property - def script_args(self) -> dict: - """Additional arguments to pass to the script file.""" - return self._script_args - - @script_args.setter - def script_args(self, args: dict) -> None: - self._script_args = args - - @property - def nb_nodes(self) -> int: - """Number of nodes on which the job runs.""" - return self._chunks - - @nb_nodes.setter - def nb_nodes(self, chunks: int) -> None: - self._chunks = chunks - - @property - def ncpus(self) -> int: - """Number of total cores used per node.""" - return self._ncpus - - @ncpus.setter - def ncpus(self, ncpus: int) -> None: - self._ncpus = ncpus - - @property - def ngpus(self) -> int: - """Number of total GPU used per node.""" - return self._ngpus - - @ngpus.setter - def ngpus(self, ngpus: int) -> None: - self._ngpus = ngpus - - @property - def mem(self) -> str: - """Maximum amount of physical memory used by the job.""" - return self._mem - - @mem.setter - def mem(self, mem: str) -> None: - self._mem = mem - - @property - def walltime(self) -> Timedelta: - """Maximum amount of real time during which the job can be running.""" - return self._walltime - - @property - def walltime_str(self) -> str: - """String representation of the ``walltime``.""" - total_seconds = self.walltime.total_seconds() - hours, remainder = divmod(total_seconds, 3600) - minutes, seconds = divmod(remainder, 60) - return ":".join(f"{t:02}" for t in map(int, (hours, minutes, seconds))) - - @walltime.setter - def walltime(self, walltime: str | Timedelta) -> None: - self._walltime = ( - walltime if type(walltime) is Timedelta else Timedelta(walltime) - ) - - @property - def venv_name(self) -> str: - """Name of the conda virtual environment in which the job is running.""" - return self._venv_name - - @venv_name.setter - def venv_name(self, venv_name: str) -> None: - self._venv_name = venv_name - - @property - def venv_activate_script(self) -> str: - """Bash script used for activating the conda virtual environment.""" - return f". /appli/anaconda/latest/etc/profile.d/conda.sh; conda activate {self.venv_name}" - - @property - def queue(self) -> Literal["omp", "mpi"]: - """Queue in which the job will be submitted.""" - return self._queue - - @queue.setter - def queue(self, queue: Literal["omp", "mpi"]) -> None: - self._queue = queue - - @property - def name(self) -> str: - """Name of the job.""" - return self._name - - @name.setter - def name(self, name: str) -> None: - self._name = name - - @property - def status(self) -> JobStatus: - """Status of the job. - - ``UNPREPARED``: The job file hasn't been written yet. - ``PREPARED``: The job file has been written but not submitted. - ``QUEUED``: The job has been queued. - ``RUNNING``: The job is currently running. - ``COMPLETED``: The job has been completed. - - """ - return self._status - - @status.setter - def status(self, status: JobStatus) -> None: - self._status = status - - @property - def path(self) -> Path | None: - """Path of the job file.""" - return self._path - - @path.setter - def path(self, path: Path) -> None: - self._path = path - - @property - def output_folder(self) -> Path | None: - """Folder in which the output files (``.out`` and ``.err``) will be written.""" - return self._output_folder - - @output_folder.setter - def output_folder(self, output_folder: Path | None) -> None: - self._output_folder = output_folder - - @property - def job_id(self) -> str: - """Job ID.""" - return self._id - - @job_id.setter - def job_id(self, job_id: str) -> None: - self._id = job_id - - @property - def job_info(self) -> dict | None: - """Information about the job as returned by a qstat request.""" - return self._info - - @job_info.setter - def job_info(self, info: dict) -> None: - self._info = info - - def progress(self) -> None: - """Bring the job to the next state.""" - if self.status == JobStatus.COMPLETED: - return - self._status = JobStatus(self._status.value + 1) - - def _build_arg_string(self) -> str: - """Build a string representation of the job's arguments.""" - arg_list = [] - for key, value in self.script_args.items(): - if isinstance(value, bool): - arg_list.append(f"--{'no-' if not value else ''}{key}") - else: - arg_list.append(f"--{key} {value}") - return " ".join(arg_list) - - def write_pbs(self, path: Path) -> None: - """Write a ``pbs`` file matching the job. - - Parameters - ---------- - path: Path - Path of the ``pbs`` file to write. - - """ - preamble = "#!/bin/bash" - - select_parts = { - "select": self.nb_nodes, - "ncpus": self.ncpus, - "mem": self.mem, - } - if self.ngpus is not None: - select_parts["ngpus"] = self.ngpus - select_str = ":".join(f"{k}={v}" for k, v in select_parts.items()) - - request = { - "-N": self.name, - "-q": self.queue, - "-l": [ - select_str, - f"walltime={self.walltime_str}", - ], - "-o": f"{self.output_folder}/{self.name}.out" - if self.output_folder - else None, - "-e": f"{self.output_folder}/{self.name}.err" - if self.output_folder - else None, - } - request_str = "\n".join( - f"#PBS {key} {value}" - if type(value) is not list - else "\n".join(f"#PBS {key} {value_part}" for value_part in value) - for key, value in request.items() - if value - ) - - script = f"python {self.script_path} {self._build_arg_string()}" - - pbs = f"{preamble}\n{request_str}\n{self.venv_activate_script}\n{script}" - with path.open("w") as file: - file.write(pbs) - - self.path = path - self.progress() - - def submit_pbs( - self, - dependency: Job | list[Job] | str | list[str] | None = None, - ) -> None: - """Submit the ``pbs`` file of the job to a ``pbs`` queueing system. - - Parameters - ---------- - dependency: Job | list[Job] | str | None - Job dependency. Can be: - - A ``Job`` instance: will wait for that job to complete successfully - - A ``list[Job]``: will wait for all jobs to complete successfully - - A ``str``: job ID (e.g., ``"12345.datarmor"``) or dependency specification - - ``None``: no dependency - - """ - if self.update_status() is not JobStatus.PREPARED: - msg = "Job should be written before being submitted." - raise ValueError(msg) - - cmd = ["qsub"] - - if dependency is not None: - dependency_str = self._build_dependency_string(dependency) - if dependency_str: - cmd.extend(["-W", f"depend={dependency_str}"]) - - cmd.append(str(self.path)) - - try: - request = subprocess.run( - cmd, - capture_output=True, - text=True, - check=False, - ) - except subprocess.CalledProcessError as e: - msg = f"Submission failed with exit code {e.returncode}" - raise RuntimeError(msg) from e - - self.job_id = request.stdout.split(".", maxsplit=1)[0].strip() - self.update_status() - - _VALID_DEPENDENCY_TYPES = {"afterok", "afterany", "afternotok", "after"} - - @staticmethod - def _validate_dependency_type(dependency_type: str) -> None: - if dependency_type not in Job._VALID_DEPENDENCY_TYPES: - msg = ( - f"Unsupported dependency type '{dependency_type}'. " - f"Expected one of {sorted(Job._VALID_DEPENDENCY_TYPES)}." - ) - raise ValueError(msg) - - @staticmethod - def _validate_dependency(dependency: list[str] | list[Job]) -> list[str]: - job_ids = [dep.job_id if isinstance(dep, Job) else dep for dep in dependency] - job_id_length = 7 - for job_id in job_ids: - if not job_id.isdigit() or len(job_id) != job_id_length: - msg = ( - f"Invalid job ID '{job_id}'. " - f"Job IDs must be {job_id_length} digits long." - ) - raise ValueError(msg) - return job_ids - - @staticmethod - def _build_dependency_string( - dependency: str | Job | list[str] | list[Job], - dependency_type: str = "afterok", - ) -> str: - """Build a PBS dependency string. - - Parameters - ---------- - dependency: Job | str - ``Job`` or job ID to depend on. - dependency_type: str - Type of dependency (``afterok``, ``afterany``, ``afternotok``, ``after``). - - Returns - ------- - str - PBS dependency string. - - Examples - -------- - >>> Job._build_dependency_string("1234567") - 'afterok:1234567' - >>> Job._build_dependency_string(["1234567", "4567891"]) - 'afterok:1234567:4567891' - >>> Job._build_dependency_string("7894561", dependency_type="afterany") - 'afterany:7894651' - - """ - dependency = dependency if isinstance(dependency, list) else [dependency] - id_str = Job._validate_dependency(dependency) - Job._validate_dependency_type(dependency_type) - - if unsubmitted_job := next( - ( - j - for j in dependency - if isinstance(j, Job) and j.status.value < JobStatus.QUEUED.value - ), - None, - ): - msg = f"Job '{unsubmitted_job.name}' has not been submitted yet." - raise ValueError(msg) - - return f"{dependency_type}:{':'.join(id_str)}" - - def update_info(self) -> None: - """Request info about the job and update it.""" - if self.job_id is None: - return - - try: - request = subprocess.run( - ["qstat", "-f", self.job_id], - capture_output=True, - text=True, - check=False, - ) - stdout = request.stdout - except subprocess.CalledProcessError as e: - msg = f"Qstat failed with exit code {e.returncode}" - raise RuntimeError(msg) from e - - if not stdout: - err = request.stderr - if "Job has finished" in err: - self.status = JobStatus.COMPLETED - self.job_info["job_state"] = "C" - if "Unknown Job Id" in err: - msg = f"Unknown Job Id {self.job_id}" - raise ValueError(msg) - return - - info = {} - for line in stdout.splitlines(): - if "=" not in line: - continue - key, value = line.split("=", 1) - info[key.strip()] = value.strip() - self.job_info = info - - def update_status(self) -> JobStatus: - """Request info about the job and update its status. - - Returns - ------- - JobStatus: - The updated status of the job. - - """ - if self.job_id is None: - self.status = ( - JobStatus.PREPARED - if self.path and self.path.exists() - else JobStatus.UNPREPARED - ) - return self.status - - self.update_info() - - if self.status == JobStatus.COMPLETED: - return self.status - - job_state = { - "Q": JobStatus.QUEUED, - "R": JobStatus.RUNNING, - } - if self.job_info["job_state"] in job_state: - self.status = job_state[self.job_info["job_state"]] - return self.status - - -class JobBuilder: - """Class that should be attached to a Public API ``Project`` for working with jobs. - - If a ``Project`` has a ``JobBuilder``, it will use it to run transforms through jobs. - - """ - - def __init__(self, config: JobConfig = JobConfig) -> None: - """Initialize a ``JobBuilder`` instance. - - Parameters - ---------- - config: JobConfig - Config of the jobs built by this job builder. - - """ - self.config = config - self.jobs = [] - - def create_jobs( - self, - nb_tasks: int, - script_path: Path, - script_args: dict, - output_folder: Path, - job_name: str = "osekit_transform", - nb_jobs: int = 1, - ) -> None: - """Create the jobs corresponding to each batch. - - Parameters - ---------- - nb_tasks: - The number of tasks that are distributed across ``nb_jobs`` jobs. - script_path: Path - Path to the export script. - script_args: dict - Arguments passed to the export script. - job_name: str - Name of the job. - If there are multiple batches, each batch will be suffixed - with "_{index}". - output_folder: Path - Folder in which the job output log files are saved. - nb_jobs: int - Number of batches used to run the transform. - Each batch will run in a separate job. - - """ - batch_indexes = file_indexes_per_batch( - total_nb_files=nb_tasks, - nb_batches=nb_jobs, - ) - for index, (start, stop) in enumerate(batch_indexes): - self.create_job( - script_path=script_path, - script_args=script_args | {"first": start, "last": stop}, - name=job_name + (f"_{index}" if len(batch_indexes) > 1 else ""), - output_folder=output_folder, - ) - - def create_job( - self, - script_path: Path, - script_args: dict | None = None, - name: str = "osekit_transform", - output_folder: Path | None = None, - ) -> None: - """Create a new ``Job`` instance. - - Parameters - ---------- - script_path: Path - Path to the script file the job must run. - script_args: dict | None - Additional arguments to pass to the script file. - name: str - Name of the job. - output_folder: Path | None - Folder in which the output files (``.out`` and ``.err``) will be written. - - """ - job = Job( - script_path=script_path, - script_args=script_args, - name=name, - output_folder=output_folder, - config=self.config, - ) - job.write_pbs(output_folder / f"{name}.pbs") - self.jobs.append(job) - - def submit_pbs( - self, - dependencies: dict[str, Job | list[Job]] | None = None, - ) -> None: - """Submit all prepared jobs to the ``pbs`` queueing system. - - Parameters - ---------- - dependencies: dict[str, Job | list[Job]] | None - Optional dictionary mapping job names to their dependencies. - Example: ``{"job2": job1, "job3": [job1, job2]}`` - - """ - for job in self.jobs: - if job.update_status() is not JobStatus.PREPARED: - continue - - # Check if this job has dependencies - depend_on = None - if dependencies and job.name in dependencies: - depend_on = dependencies[job.name] - - job.submit_pbs(dependency=depend_on) diff --git a/tests/_static/job_status_request_results/pbs.txt b/tests/_static/job_status_request_results/pbs.txt new file mode 100644 index 000000000..8be9ac450 --- /dev/null +++ b/tests/_static/job_status_request_results/pbs.txt @@ -0,0 +1,39 @@ +Job Id: 7137005.a24films0 + Job_Name = SwissArmyMan + Job_Owner = daniels@a24films.fr + job_state = R + queue = jetski + server = a24films0 + Checkpoint = u + ctime = Mon Sep 7 08:00:18 2026 + Error_Path = a24films0.a24films.fr:/home1/a24/filmls/SwissArmyMan.err + Hold_Types = n + Join_Path = oe + Keep_Files = n + Mail_Points = a + mtime = Mon Sep 7 08:00:18 2026 + Output_Path = a24films0.a24films.fr:/home1/a24/filmls/SwissArmyMan.out + Priority = 0 + qtime = Mon Sep 7 08:00:18 2026 + Rerunable = True + Resource_List.mem = 112gb + Resource_List.mpiprocs = 28 + Resource_List.ncpus = 24 + Resource_List.nodect = 1 + Resource_List.place = group=switch + Resource_List.select = 1:ncpus=24:mem=112g + Resource_List.walltime = 00:30:00 + resources_used.cpupercent = 1 + resources_used.cput = 01:01:25 + resources_used.mem = 40678964kb + resources_used.ncpus = 56 + resources_used.vmem = 3828570148kb + resources_used.walltime = 00:10:37 + substate = 10 + Variable_List = PBS_O_SYSTEM=Linux,PBS_O_SHELL=/bin/csh, + etime = Mon Sep 7 08:00:18 2026 + run_count = 15 + eligible_time = 00:00:00 + Submit_arguments = /home/a24films-exp/workspace/swiss/scheduled__2 + 026-09-07T07-00-00/army_man/run.sh + project = _pbs_project_default diff --git a/tests/_static/job_status_request_results/slurm_sacct.txt b/tests/_static/job_status_request_results/slurm_sacct.txt new file mode 100644 index 000000000..28f6bf2a2 --- /dev/null +++ b/tests/_static/job_status_request_results/slurm_sacct.txt @@ -0,0 +1 @@ +7137005|jetski|SwissArmyMan|daniels|COMPLETED|10:37|1|compute-114-9|Dependency diff --git a/tests/_static/job_status_request_results/slurm_squeue.txt b/tests/_static/job_status_request_results/slurm_squeue.txt new file mode 100644 index 000000000..525554a0c --- /dev/null +++ b/tests/_static/job_status_request_results/slurm_squeue.txt @@ -0,0 +1 @@ +7137005|jetski|SwissArmyMan|daniels|RUNNING|10:37|1|compute-114-9 diff --git a/tests/test_export_transform.py b/tests/test_export_transform.py index f2de2a370..0f979408b 100644 --- a/tests/test_export_transform.py +++ b/tests/test_export_transform.py @@ -10,7 +10,7 @@ from osekit.core.audio_dataset import AudioDataset from osekit.public import export_transform from osekit.public.export_transform import create_parser -from osekit.utils.job import Job +from osekit.job.job import Job def test_parser_factory() -> None: @@ -98,7 +98,7 @@ def script_arguments() -> dict: def test_specified_arguments(script_arguments: dict) -> None: parser = create_parser() - parsed_str = Job(Path(), script_arguments)._build_arg_string() + parsed_str = Job(Path(), script_arguments).get_arg_string() args = parser.parse_args(shlex.split(parsed_str)) diff --git a/tests/test_job.py b/tests/test_job.py index f3b4c842d..409d29d0f 100644 --- a/tests/test_job.py +++ b/tests/test_job.py @@ -1,52 +1,21 @@ from __future__ import annotations import subprocess -from contextlib import nullcontext from pathlib import Path import numpy as np import pytest from pandas import Timedelta -import osekit.utils.job as job_module -from osekit.utils.job import Job, JobBuilder, JobConfig, JobStatus +from osekit.job.builder import JobBuilder +from osekit.job.config import JobConfig +from osekit.job.job import Job, JobStatus +from osekit.job.scheduler.pbs import Pbs +from osekit.job.scheduler.scheduler import Scheduler +from osekit.job.scheduler.slurm import Slurm -@pytest.mark.parametrize( - ("initial_status", "expected_status"), - [ - pytest.param( - JobStatus.UNPREPARED, - JobStatus.PREPARED, - id="unprepared_becomes_prepared", - ), - pytest.param( - JobStatus.PREPARED, - JobStatus.QUEUED, - id="prepared_becomes_queued", - ), - pytest.param(JobStatus.QUEUED, JobStatus.RUNNING, id="queued_becomes_running"), - pytest.param( - JobStatus.RUNNING, - JobStatus.COMPLETED, - id="running_becomes_completed", - ), - pytest.param( - JobStatus.COMPLETED, - JobStatus.COMPLETED, - id="completed_remains_completed", - ), - ], -) -def test_job_progress(initial_status: JobStatus, expected_status: JobStatus) -> None: - job = Job(script_path=Path()) - assert job.status == JobStatus.UNPREPARED - job._status = initial_status - job.progress() - assert job.status == expected_status - - -def test_properties_and_venv_activation() -> None: +def test_properties() -> None: script = Path("myscript.py") nb_nodes = 2 ncpus = 28 @@ -56,7 +25,6 @@ def test_properties_and_venv_activation() -> None: mem="16gb", walltime=Timedelta(hours=2), venv_name="merriweather", - queue="mpi", ) job = Job( script_path=script, @@ -75,30 +43,8 @@ def test_properties_and_venv_activation() -> None: assert job.walltime == Timedelta(hours=2) assert job.venv_name == "merriweather" assert job.name == "post_pavillion" - assert job.queue == "mpi" assert job.output_folder == Path("output") - # venv activation - expected = ( - ". /appli/anaconda/latest/etc/profile.d/conda.sh; conda activate merriweather" - ) - assert job.venv_activate_script == expected - - -def test_progress_transitions() -> None: - job = Job(Path("strawberry.py")) - assert job.status == JobStatus.UNPREPARED - for expected in ( - JobStatus.PREPARED, - JobStatus.QUEUED, - JobStatus.RUNNING, - JobStatus.COMPLETED, - ): - job.progress() - assert job.status == expected - job.progress() - assert job.status == JobStatus.COMPLETED - def test_walltime_str_and_setter() -> None: job = Job(Path("bossanova.py")) @@ -108,6 +54,67 @@ def test_walltime_str_and_setter() -> None: assert job.walltime_str == "13:08:09" +def test_pbs_build_job_specifications() -> None: + job = Job( + script_path=Path(), + config=JobConfig( + nb_nodes=2, + ncpus=3, + ngpus=1, + mem="16gb", + walltime=Timedelta(hours=2), + venv_name="cool_env", + ), + output_folder=Path(r"cool/folder"), + name="cool_job", + ) + + specifications = Pbs(queue="mpi")._build_job_specification(job=job).splitlines() + + for expected_specification in ( + "#PBS -N cool_job", + "#PBS -q mpi", + "#PBS -l select=2:ncpus=3:mem=16gb:ngpus=1", + "#PBS -l walltime=02:00:00", + f"#PBS -o {Path('cool/folder') / 'cool_job.out'}", + f"#PBS -e {Path('cool/folder') / 'cool_job.err'}", + ): + assert expected_specification in specifications + + +def test_slurm_build_job_specifications() -> None: + job = Job( + script_path=Path(), + config=JobConfig( + nb_nodes=2, + ncpus=3, + ngpus=1, + mem="16gb", + walltime=Timedelta(hours=2), + venv_name="cool_env", + ), + output_folder=Path(r"cool/folder"), + name="cool_job", + ) + + specifications = ( + Slurm(partition="gpu")._build_job_specification(job=job).splitlines() + ) + + for expected_specification in ( + "#SBATCH --job-name=cool_job", + "#SBATCH --partition=gpu", + "#SBATCH --nodes=2", + "#SBATCH --cpus-per-task=3", + "#SBATCH --mem=16gb", + "#SBATCH --gpus=1", + "#SBATCH --time=02:00:00", + f"#SBATCH --output={Path('cool/folder') / 'cool_job.out'}", + f"#SBATCH --error={Path('cool/folder') / 'cool_job.err'}", + ): + assert expected_specification in specifications + + def test_write_pbs(tmp_path: Path) -> None: script = tmp_path / "shpouik_shpouik.py" script.write_text("print('edgar')") @@ -121,7 +128,9 @@ def test_write_pbs(tmp_path: Path) -> None: output_folder=output_dir, ) pbs_path = tmp_path / "lafayette.pbs" - job.write_pbs(pbs_path) + + pbs_scheduler = Pbs(queue="omp") + pbs_scheduler.write(job=job, path=pbs_path) content = pbs_path.read_text().splitlines() assert content[0] == "#!/bin/bash" @@ -130,18 +139,17 @@ def test_write_pbs(tmp_path: Path) -> None: assert any("select=1:ncpus=2:mem=8gb" in line for line in content) assert any("walltime=01:00:00" in line for line in content) assert any( - line.startswith(f"#PBS -o {job.output_folder}/{job.name}.out") + line.startswith(f"#PBS -o {job.output_folder / job.name}.out") for line in content ) assert any( - line.startswith(f"#PBS -e {job.output_folder}/{job.name}.err") + line.startswith(f"#PBS -e {job.output_folder / job.name}.err") for line in content ) - assert ( - ". /appli/anaconda/latest/etc/profile.d/conda.sh; conda activate osekit" - in content - ) + assert ". /appli/anaconda/latest/etc/profile.d/conda.sh" in content + assert "conda activate osekit" in content + last = content[-1] assert last.startswith(f"python {script}") assert "--vieille face" in last @@ -167,7 +175,8 @@ def test_write_pbs_job_with_gpu(tmp_path: Path) -> None: output_folder=output_dir, ) pbs_path = tmp_path / "patch.pbs" - job.write_pbs(pbs_path) + pbs_scheduler = Pbs(queue="omp") + pbs_scheduler.write(job=job, path=pbs_path) content = pbs_path.read_text().splitlines() assert any("select=1:ncpus=2:mem=8gb:ngpus=2" in line for line in content) @@ -178,11 +187,12 @@ def test_write_pbs_job_with_gpu(tmp_path: Path) -> None: def test_submit_pbs_without_write_raises() -> None: job = Job(Path("script.py")) + pbs = Pbs(queue="omp") with pytest.raises( ValueError, match=r"Job should be written before being submitted.", ): - job.submit_pbs() + pbs.submit(job=job) def test_submit_pbs_success(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: @@ -191,7 +201,8 @@ def test_submit_pbs_success(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> outdir = tmp_path job = Job(script, name="amobishoproden", output_folder=outdir) pbs_path = tmp_path / "amobishoproden.pbs" - job.write_pbs(pbs_path) + pbs = Pbs() + pbs.write(job=job, path=pbs_path) class Dummy: def __init__(self) -> None: @@ -207,14 +218,14 @@ def __init__(self) -> None: updated_jobs = [] - def mock_update_status(self: Job) -> JobStatus: + def mock_update_status(self, job: Job) -> JobStatus: updated_jobs.append(job) return JobStatus.PREPARED - monkeypatch.setattr(Job, "update_status", mock_update_status) + monkeypatch.setattr(Scheduler, "update_status", mock_update_status) assert job.status == JobStatus.PREPARED - job.submit_pbs() + pbs.submit(job=job) assert job.job_id == "35173" assert np.array_equal( @@ -224,12 +235,9 @@ def mock_update_status(self: Job) -> JobStatus: def test_submit_pbs_errors(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: - script = tmp_path / "boc.py" - script.write_text("") - outdir = tmp_path - job = Job(script, name="amobishoproden", output_folder=outdir) - pbs_path = tmp_path / "amobishoproden.pbs" - job.write_pbs(pbs_path) + job = Job(Path()) + pbs_scheduler = Pbs(queue="omp") + job.status = JobStatus.PREPARED class Dummy: def __init__(self) -> None: @@ -242,28 +250,34 @@ def __init__(self) -> None: lambda *args, **kwargs: Dummy(), ) - assert job.status == JobStatus.PREPARED + def mock_update_status(self: Pbs, job: Job) -> JobStatus: + return JobStatus.PREPARED + + monkeypatch.setattr(Scheduler, "update_status", mock_update_status) + + # Submit error should leave the job prepared: with pytest.raises(RuntimeError, match="Submission failed with exit code 5"): - job.submit_pbs() + pbs_scheduler.submit(job=job) assert job.status == JobStatus.PREPARED -def test_update_info_no_job_id() -> None: +def test_pbs_update_info_no_job_id() -> None: job = Job(Path("pixies.py")) + pbs_scheduler = Pbs() job.job_id = None - job.update_info() - assert job.job_info is None + pbs_scheduler.update_info(job=job) + assert not job.info -def test_update_info_parse_stdout(monkeypatch: pytest.MonkeyPatch) -> None: - job = Job(Path("fontaines.py")) - job.job_id = "43" - job.status = JobStatus.RUNNING - raw = " frankie = cosmos \navey=tare\nattic= abasement\nthis will be ignored" +def test_pbs_update_info_parse_stdout(monkeypatch: pytest.MonkeyPatch) -> None: + job = Job(script_path=Path("fontaines.py"), name="SwissArmyMan") + job.job_id = "7137005" class Dummy: - stdout = raw + stdout = ( + Path(__file__).parent / "_static/job_status_request_results/pbs.txt" + ).read_text() stderr = "" monkeypatch.setattr( @@ -271,31 +285,138 @@ class Dummy: "run", lambda *args, **kwargs: Dummy(), ) - job.update_info() - assert job.job_info == {"frankie": "cosmos", "avey": "tare", "attic": "abasement"} + scheduler = Pbs() + scheduler.update_info(job=job) + assert job.job_id == "7137005" + assert job.name == "SwissArmyMan" + assert job.status == JobStatus.RUNNING + assert job.info == { + "user": "daniels", + "time": "00:10:37", + "queue": "jetski", + } -def test_update_info_completed(monkeypatch: pytest.MonkeyPatch) -> None: - job = Job(Path("amok.py")) - job.job_id = "25022013" - job.job_info = {} +def test_slurm_update_info_parse_stdout(monkeypatch: pytest.MonkeyPatch) -> None: + job = Job(script_path=Path("fontaines.py"), name="SwissArmyMan") + job.job_id = "7137005" class Dummy: - stdout = "" - stderr = "Atoms\nJob has finished\nFor peace" + stdout = ( + Path(__file__).parent + / "_static/job_status_request_results/slurm_squeue.txt" + ).read_text() + stderr = "" monkeypatch.setattr( subprocess, "run", lambda *args, **kwargs: Dummy(), ) + scheduler = Slurm() + scheduler.update_info(job=job) + assert job.job_id == "7137005" + assert job.name == "SwissArmyMan" + assert job.status == JobStatus.RUNNING + assert job.info == { + "user": "daniels", + "time": "10:37", + "partition": "jetski", + "node_list": "compute-114-9", + } + + +def test_slurm_get_info_completed_job(monkeypatch: pytest.MonkeyPatch) -> None: + job = Job(script_path=Path("fontaines.py"), name="SwissArmyMan") + job.job_id = "7137005" + + class EmptyRequest: + def __init__(self) -> None: + self.stdout = "" + self.stderr = "" + + class DummySacctRequest: + def __init__(self) -> None: + self.stdout = ( + Path(__file__).parent + / "_static/job_status_request_results/slurm_sacct.txt" + ).read_text() + self.stderr = "" + + def mock_run(*args, **kwargs) -> EmptyRequest | DummySacctRequest: + cmd = args[0][0] + return EmptyRequest() if "squeue" in cmd else DummySacctRequest() - job.update_info() + monkeypatch.setattr(subprocess, "run", mock_run) + + scheduler = Slurm() + scheduler.update_info(job=job) + assert job.job_id == "7137005" + assert job.name == "SwissArmyMan" assert job.status == JobStatus.COMPLETED - assert job.job_info["job_state"] == "C" + assert job.info == { + "user": "daniels", + "time": "10:37", + "partition": "jetski", + "node_list": "compute-114-9 (Dependency)", + } + + +def test_slurm_get_info_error(monkeypatch: pytest.MonkeyPatch) -> None: + job = Job(script_path=Path("fontaines.py"), name="SwissArmyMan") + job.job_id = "7137005" + + class EmptyRequest: + def __init__(self) -> None: + self.stdout = "" + self.stderr = "" + + class DummySacctRequestError: + def __init__(self) -> None: + self.stdout = "" + self.stderr = "timeout" + + def mock_run(*args, **kwargs) -> EmptyRequest | DummySacctRequestError: + cmd = args[0][0] + return EmptyRequest() if "squeue" in cmd else DummySacctRequestError() + + monkeypatch.setattr(subprocess, "run", mock_run) + + scheduler = Slurm() + with pytest.raises(ValueError, match=r"7137005.*timeout"): + scheduler.update_info(job=job) + + +def test_slurm_get_info_no_output(monkeypatch: pytest.MonkeyPatch) -> None: + job = Job(script_path=Path("fontaines.py"), name="SwissArmyMan") + job.job_id = "7137005" + job.status = JobStatus.PREPARED + + class EmptyRequest: + def __init__(self) -> None: + self.stdout = "" + self.stderr = "" + + class DummySacctRequest: + def __init__(self) -> None: + self.stdout = "" + self.stderr = "" + + def mock_run(*args, **kwargs) -> EmptyRequest | DummySacctRequest: + cmd = args[0][0] + return EmptyRequest() if "squeue" in cmd else DummySacctRequest() + + monkeypatch.setattr(subprocess, "run", mock_run) + + scheduler = Slurm() + scheduler.update_info(job=job) # Shouldn't raise + # Information should not have changed + assert job.job_id == "7137005" + assert job.status == JobStatus.PREPARED + assert job.info == {} -def test_update_info_unknown_job_raises(monkeypatch: pytest.MonkeyPatch) -> None: +def test_pbs_update_info_unknown_job_raises(monkeypatch: pytest.MonkeyPatch) -> None: job = Job(Path("pompom.py")) job.job_id = "17112014" @@ -309,11 +430,34 @@ class Dummy: lambda *args, **kwargs: Dummy(), ) + scheduler = Pbs() with pytest.raises(ValueError, match="Unknown Job Id 17112014"): - job.update_info() + scheduler.update_info(job=job) + + +def test_pbs_update_info_without_qstat_output(monkeypatch: pytest.MonkeyPatch) -> None: + job = Job(Path("pompom.py")) + job.job_id = "17112014" + job.status = JobStatus.RUNNING + + class Dummy: + stdout = "" + stderr = "" + monkeypatch.setattr( + subprocess, + "run", + lambda *args, **kwargs: Dummy(), + ) + + Pbs().update_info(job=job) -def test_update_info_error(monkeypatch: pytest.MonkeyPatch) -> None: + # No qstat output shouldn't do anything + assert job.job_id == "17112014" + assert job.status == JobStatus.RUNNING + + +def test_pbs_update_info_error(monkeypatch: pytest.MonkeyPatch) -> None: job = Job(Path("pompom.py")) job.job_id = "17112014" @@ -327,40 +471,43 @@ def __init__(self) -> None: lambda *args, **kwargs: Dummy(), ) - with pytest.raises(RuntimeError, match="Qstat failed with exit code 5"): - job.update_info() + scheduler = Pbs() + with pytest.raises(RuntimeError, match=r"qstat failed.*code 5"): + scheduler.update_info(job=job) -def test_update_status(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: +def test_pbs_update_status(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: job = Job(Path("porticoquartet.py")) job.path = tmp_path / "pompidou.pbs" - assert job.update_status() == JobStatus.UNPREPARED + scheduler = Pbs() + + assert scheduler.update_status(job=job) == JobStatus.UNPREPARED job.path.write_text("prickly pear") - assert job.update_status() == JobStatus.PREPARED + assert scheduler.update_status(job=job) == JobStatus.PREPARED + + def mock_update_info( + self: Scheduler, + job: Job, + status: JobStatus, + *args: list, + **kwargs: dict, + ) -> None: + job.status = status monkeypatch.setattr( - job, + Scheduler, "update_info", - lambda: None, + lambda self, job: mock_update_info(self, job=job, status=JobStatus.QUEUED), ) - job.job_info = {"job_state": "Q"} job.job_id = "5129195" - assert job.update_status() == JobStatus.QUEUED + assert scheduler.update_status(job=job) == JobStatus.QUEUED assert job.status == JobStatus.QUEUED - job.job_info = {"job_state": "R"} - assert job.update_status() == JobStatus.RUNNING - assert job.status == JobStatus.RUNNING - - job.status = JobStatus.COMPLETED - assert job.update_status() == JobStatus.COMPLETED - assert job.status == JobStatus.COMPLETED - -def test_job_builder_write(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: +def test_pbs_job_builder_write(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: called = {} class DummyJob: @@ -369,12 +516,13 @@ def __init__(self, **kwargs: dict) -> None: self.path = None self.status = JobStatus.UNPREPARED - def write_pbs(self, path: Path) -> None: - called["write_pbs"] = path - self.path = path - self.status = JobStatus.PREPARED + def mock_write(self: Pbs, job: Job, path: Path) -> None: + called["write_pbs"] = path + self.path = path + job.status = JobStatus.PREPARED - monkeypatch.setattr(job_module, "Job", DummyJob) + monkeypatch.setattr("osekit.job.builder.Job", DummyJob) + monkeypatch.setattr(Scheduler, "write", mock_write) job_config = JobConfig( nb_nodes=2, @@ -382,10 +530,9 @@ def write_pbs(self, path: Path) -> None: mem="60gb", walltime=Timedelta(hours=2), venv_name="abyssinie", - queue="mpi", ) - job_builder = JobBuilder(config=job_config) + job_builder = JobBuilder(scheduler=Pbs(), config=job_config) assert job_builder.jobs == [] @@ -440,7 +587,7 @@ def test_build_arg_string_booleans(tmp_path: Path): ) job = next(iter(job_builder.jobs)) - arg_str = job._build_arg_string() + arg_str = job.get_arg_string() assert arg_str == "--no-danser --avec --le 0.3 --vent test" @@ -453,13 +600,19 @@ def __init__(self, name: str, status: JobStatus) -> None: self.name = name self.status = status - def submit_pbs(self, dependency=None) -> None: - submitted_jobs.append((self.name, dependency)) + def mock_submit( + self: Scheduler, + job: Job, + dependencies: Job | str | None = None, + ) -> None: + submitted_jobs.append((job, dependencies)) - def update_status(self) -> JobStatus: - return self.status + def mock_update_status(self: Scheduler, job: Job) -> JobStatus: + return job.status - monkeypatch.setattr(job_module, "Job", DummyJob) + monkeypatch.setattr("osekit.job.job.Job", DummyJob) + monkeypatch.setattr(Scheduler, "submit", mock_submit) + monkeypatch.setattr(Scheduler, "update_status", mock_update_status) jobs = [ DummyJob(name="unprepared", status=JobStatus.UNPREPARED), @@ -472,149 +625,172 @@ def update_status(self) -> JobStatus: job_builder = JobBuilder() job_builder.jobs = jobs - dependencies = {"prepared": jobs[0]} + unprepared_job = job_builder.jobs[0] + prepared_job = job_builder.jobs[1] + + dependencies = { + prepared_job: {"beforeok": unprepared_job}, + unprepared_job: {"afterany": prepared_job}, + } + + job_builder.submit(dependencies=dependencies) + + # Only the prepared job should be submitted + assert len(submitted_jobs) == 1 + submitted_job = submitted_jobs[0] + assert submitted_job[0] == prepared_job + + # Only the submitted job dependencies should be injected + assert submitted_job[1] == dependencies[prepared_job] + + +def test_pbs_build_dependencies_string_validates_type( + monkeypatch: pytest.MonkeyPatch, +) -> None: + validate_calls = [] + + def mock_validate(dependency_type: str) -> None: + validate_calls.append(dependency_type) + + monkeypatch.setattr(Scheduler, "_validate_dependency_type", mock_validate) + + dependencies = {"afterok": "1234567", "afterany": ["2345678", "3456789"]} + Pbs()._build_dependency_string( + dependencies=dependencies, + ) + + assert all(dependency_type in validate_calls for dependency_type in dependencies) + + +def test_pbs_validate_dependency_type() -> None: + pbs = Pbs() + + # Supported dependency type shouldn't raise + pbs._validate_dependency_type("afterok") + + # Unsupported dependency type should raise + with pytest.raises(ValueError) as e: + pbs._validate_dependency_type("afterdummy") - job_builder.submit_pbs(dependencies=dependencies) + assert e.match("Unsupported dependency type 'afterdummy'") + for supported in Pbs._VALID_DEPENDENCY_TYPES: + assert e.match(supported) - assert submitted_jobs == [("prepared", jobs[0])] + +def id_to_job(job_id: str | list[str]) -> Job | list[Job]: + """Convert a Job ID ``job_id`` to a Job object with an ID of ``job_id`` + + If ``job_id`` is a list, converts the list of job_ids to a list of jobs + with the given IDs.""" + if isinstance(job_id, str): + job = Job(Path()) + job._id = job_id + job.status = JobStatus.QUEUED + return job + output = [] + for j_id in job_id: + job = Job(Path()) + job._id = j_id + job.status = JobStatus.QUEUED + output.append(job) + return output @pytest.mark.parametrize( - ("dependency", "ids", "status", "expected"), + ("dependencies", "expected"), [ pytest.param( - ["1234567"], - [None], - [None], - nullcontext("afterok:1234567"), - id="single_job_id", - ), - pytest.param( - ["1234567", "4567891", "7891234"], - [None] * 3, - [None] * 3, - nullcontext("afterok:1234567:4567891:7891234"), - id="multiple_job_ids", + {"afterok": "1234567"}, + "-W depend=afterok:1234567", + id="one_type_one_job", ), pytest.param( - ["123"], - [None], - [None], - pytest.raises( - ValueError, - match=r"Invalid job ID '123'\. Job IDs must be 7 digits long\.", - ), - id="invalid_job_id_too_short", - ), - pytest.param( - [Job(script_path=Path("test.py"), name="job_1")], - ["12345678"], - [JobStatus.QUEUED], - pytest.raises( - ValueError, - match=r"Invalid job ID '12345678'\. Job IDs must be 7 digits long\.", - ), - id="invalid_job_id_too_long", - ), - pytest.param( - ["abcdefg"], - [None], - [None], - pytest.raises( - ValueError, - match=r"Invalid job ID 'abcdefg'\. Job IDs must be 7 digits long\.", - ), - id="invalid_job_id_non_numeric", + {"afterok": ["1234567", "2345678"]}, + "-W depend=afterok:1234567:2345678", + id="one_type_multiple_jobs", ), pytest.param( - ["1234567", "not_a_job_id"], - [None] * 2, - [None] * 2, - pytest.raises( - ValueError, - match=r"Invalid job ID 'not_a_job_id'\. Job IDs must be 7 digits long\.", - ), - id="multiple_job_id_one_invalid", + {"afterok": "1234567", "afterany": "2345678"}, + "-W depend=afterok:1234567,afterany:2345678", + id="multiple_types_one_job", ), pytest.param( - [Job(script_path=Path("test.py"), name="job_1")], - ["1234567"], - [JobStatus.QUEUED], - nullcontext("afterok:1234567"), - id="single_job_instance", + {"afterok": ["1234567", "2345678"], "afterany": ["3456789", "4567890"]}, + "-W depend=afterok:1234567:2345678,afterany:3456789:4567890", + id="multiple_types_multiple_jobs", ), + ], +) +def test_pbs_build_dependencies_string( + dependencies: dict[str, str | list[str]], + expected: str, +) -> None: + # %% Dependencies string from job IDs + assert Pbs()._build_dependency_string(dependencies=dependencies) == expected + + # %% Dependencies string from Job instances + dependencies = {key: id_to_job(value) for key, value in dependencies.items()} + assert Pbs()._build_dependency_string(dependencies=dependencies) == expected + + +@pytest.mark.parametrize( + ("dependencies", "instructions_or", "expected"), + [ pytest.param( - [ - Job(script_path=Path("horse_with.py"), name="job_1"), - Job(script_path=Path("no_name.py"), name="job_2"), - ], - ["1234567", "4567891"], - [JobStatus.QUEUED, JobStatus.QUEUED], - nullcontext("afterok:1234567:4567891"), - id="multiple_job_instance", + {"afterok": "1234567"}, + False, + "-d afterok:1234567", + id="one_type_one_job", ), pytest.param( - [ - Job(script_path=Path("king_crimson.py"), name="job_1"), - Job(script_path=Path("crimson_king.py"), name="job_2"), - ], - ["1234567", "not_an_id"], - [JobStatus.QUEUED, JobStatus.QUEUED], - pytest.raises( - ValueError, - match=r"Invalid job ID 'not_an_id'\. Job IDs must be 7 digits long\.", - ), - id="multiple_job_instance_invalid_one", + {"afterok": ["1234567", "2345678"]}, + False, + "-d afterok:1234567:2345678", + id="one_type_multiple_jobs", ), pytest.param( - [ - Job(script_path=Path("king_crimson.py"), name="job_1"), - "9876543", - ], - ["1234567", None], - [JobStatus.QUEUED, None], - nullcontext("afterok:1234567:9876543"), - id="job_and_string_input", + {"afterok": "1234567", "afterany": "2345678"}, + False, + "-d afterok:1234567,afterany:2345678", + id="multiple_types_one_job", ), pytest.param( - [Job(script_path=Path("test.py"), name="tornero")], - ["1234567"], - [JobStatus.UNPREPARED], - pytest.raises( - ValueError, - match="Job 'tornero' has not been submitted yet.", - ), - id="unprepared_job_instance", + {"afterok": ["1234567", "2345678"], "afterany": ["3456789", "4567890"]}, + False, + "-d afterok:1234567:2345678,afterany:3456789:4567890", + id="multiple_types_multiple_jobs", ), pytest.param( - [ - Job(script_path=Path("script.py"), name="dalida"), - Job(script_path=Path("script.py"), name="mourir_sur_scene"), - ], - ["1234567", "4567896"], - [JobStatus.QUEUED, JobStatus.PREPARED], - pytest.raises( - ValueError, - match="Job 'mourir_sur_scene' has not been submitted yet.", - ), - id="multiple_job_instance_one_not_submitted", + {"afterok": ["1234567", "2345678"], "afterany": ["3456789", "4567890"]}, + True, + "-d afterok:1234567:2345678?afterany:3456789:4567890", + id="logical_or", ), ], ) -def test_build_dependency_string_with_string_input( - dependency: list[str] | list[Job], - ids: list[str] | None, - status: list[JobStatus], - expected: str | None, +def test_slurm_build_dependencies_string( + dependencies: dict[str, str | list[str]], + instructions_or: bool, + expected: str, ) -> None: - """Test building dependency string from string and Job inputs.""" - for dep, id, st in zip(dependency, ids, status, strict=True): - if isinstance(dep, Job): - dep.status = st - dep.job_id = id + # %% Dependencies string from job IDs + assert ( + Slurm()._build_dependency_string( + dependencies=dependencies, + instructions_or=instructions_or, + ) + == expected + ) - with expected as e: - assert Job._build_dependency_string(dependency) == e + # %% Dependencies string from Job instances + dependencies = {key: id_to_job(value) for key, value in dependencies.items()} + assert ( + Slurm()._build_dependency_string( + dependencies=dependencies, + instructions_or=instructions_or, + ) + == expected + ) def test_submit_pbs_adds_dependency_flag( @@ -623,8 +799,9 @@ def test_submit_pbs_adds_dependency_flag( ) -> None: script = tmp_path / "script.py" script.write_text("") + scheduler = Pbs() job = Job(script, name="crazy_diamond", output_folder=tmp_path) - job.write_pbs(tmp_path / "wywh.pbs") + scheduler.write(job=job, path=tmp_path / "wywh.pbs") captured_cmd = {} @@ -636,41 +813,18 @@ def fake_run(cmd: list[str], *args: None, **kwargs: None) -> Dummy: captured_cmd["cmd"] = cmd return Dummy() + def mock_update_status(self: Pbs, job: Job) -> JobStatus: + return JobStatus.PREPARED + monkeypatch.setattr(subprocess, "run", fake_run) - monkeypatch.setattr(Job, "update_status", lambda _: JobStatus.PREPARED) + monkeypatch.setattr(Scheduler, "update_status", mock_update_status) - job.submit_pbs(dependency="1234567") + scheduler.submit(job=job, dependencies={"afterok": "1234567"}) assert "-W" in captured_cmd["cmd"] assert "depend=afterok:1234567" in captured_cmd["cmd"] -@pytest.mark.parametrize( - ("dependency_type", "expected"), - [ - pytest.param("afterok", nullcontext("afterok:1234567"), id="afterok"), - pytest.param("afterany", nullcontext("afterany:1234567"), id="afterany"), - pytest.param("afternotok", nullcontext("afternotok:1234567"), id="afternotok"), - pytest.param("after", nullcontext("after:1234567"), id="after"), - pytest.param( - "not_a_supported_type", - pytest.raises( - ValueError, - match=r"Unsupported dependency type 'not_a_supported_type'\. Expected one of \['after', 'afterany', 'afternotok', 'afterok'\]\.", - ), - id="invalid_dependency_type", - ), - ], -) -def test_build_dependency_string_with_different_types( - dependency_type: str, - expected: type[Exception], -) -> None: - """Test building dependency strings with different dependency types.""" - with expected as e: - assert Job._build_dependency_string("1234567", dependency_type) == e - - @pytest.mark.parametrize( "walltime", [ @@ -797,3 +951,18 @@ def patch_create_job(self: JobBuilder, **kwargs: str) -> None: ), ): assert job[0] == f"{job_name}_{idx}" + + +def test_slurm_venv_str() -> None: + slurm = Slurm() + job = Job(script_path=Path(), config=JobConfig(venv_name="cool_venv")) + venv_commands = slurm._build_venv_string(job=job).splitlines() + + assert venv_commands[0] == "module load conda" + assert venv_commands[1] == "conda activate cool_venv" + + +def test_slurm_parse_job_id() -> None: + submit_output = "Submitted batch job 3647090\n" + job_id = "3647090" + assert Slurm()._parse_job_id(submit_output=submit_output) == job_id