Slurm

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
 

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

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)

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)

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"