Split project export() method

src/osekit/public/project.py

@@ -26,11 +26,10 @@
 from osekit.core.spectro_dataset import SpectroDataset
 from osekit.public.transform import OutputType, Transform
 from osekit.utils.core import (
-    file_indexes_per_batch,
     get_umask,
     locked,
 )
-from osekit.utils.path import move_tree, ensure_within_base
+from osekit.utils.path import ensure_within_base, move_tree
 
 if TYPE_CHECKING:
     from collections.abc import Iterable
@@ -522,7 +521,7 @@ def _get_audio_dataset_subpath(
     def export(
         self,
         output_type: OutputType,
-        ads: AudioDataset | None = None,
+        ads: AudioDataset,
         sds: SpectroDataset | LTASDataset | None = None,
         subtype: str | None = None,
         spectrum_folder_name: str = "spectrum",
@@ -603,43 +602,79 @@ def export(
             )
             return
 
-        batch_indexes = file_indexes_per_batch(
-            total_nb_files=len(ads.data),
-            nb_batches=nb_jobs,
+        ads_json, sds_json = Project.get_json_paths(
+            audio_dataset=ads,
+            spectro_dataset=sds,
+            output_type=output_type,
+        )
+
+        script_args = {
+            "output-type": output_type.value,
+            "ads-json": ads_json,
+            "sds-json": sds_json,
+            "subtype": subtype,
+            "spectrum-folder-path": spectrum_folder_path,
+            "spectrogram-folder-path": spectrogram_folder_path,
+            "welch-folder-path": welch_folder_path,
+            "downsampling-quality": resample_quality_settings["downsample"],
+            "upsampling-quality": resample_quality_settings["upsample"],
+            "umask": get_umask(),
+            "multiprocessing": config.multiprocessing["is_active"],
+            "nb-processes": config.multiprocessing["nb_processes"],
+            "use-logging-setup": True,
+            "dataset-json-path": self.folder / "project.json",
+        }
+
+        self.job_builder.create_jobs(
+            nb_tasks=len(ads.data),
+            script_path=Path(export_transform.__file__),
+            script_args=script_args,
+            job_name=name,
+            output_folder=self.folder / self.SUBFOLDERS["log"],
+            nb_jobs=nb_jobs,
         )
 
+        self.job_builder.submit_pbs()
+
+    @staticmethod
+    def get_json_paths(
+        audio_dataset: AudioDataset,
+        spectro_dataset: SpectroDataset | None,
+        output_type: OutputType,
+    ) -> tuple[Path | str, Path | str]:
+        """Return the paths of the audio and spectro output JSON files.
+
+        Parameters
+        ----------
+        audio_dataset: AudioDataset
+            The ``AudioDataset`` the transform is based on.
+        spectro_dataset: SpectroDataset | None
+            The ``SpectroDataset`` that is output by the transform.
+            ``None`` if the transform is audio-only.
+        output_type: OutputType
+            The ``OutputType`` of the transform.
+
+        Returns
+        -------
+        tuple[Path | str, Path | str]:
+            Paths of the audio and spectro output JSON files, respectively.
+            If there is no output dataset for the given ``OutputType``,
+            the corresponding path in the tuple is replaced with "None".
+
+        """
         ads_json = (
-            ads.folder / f"{ads.name}.json"
+            audio_dataset.folder / f"{audio_dataset.name}.json"
             if OutputType.AUDIO in output_type
             else "None"
         )
-        sds_json = sds.folder / f"{sds.name}.json" if sds is not None else "None"
-
-        for index, (start, stop) in enumerate(batch_indexes):
-            self.job_builder.create_job(
-                script_path=Path(export_transform.__file__),
-                script_args={
-                    "output-type": output_type.value,
-                    "ads-json": ads_json,
-                    "sds-json": sds_json,
-                    "subtype": subtype,
-                    "spectrum-folder-path": spectrum_folder_path,
-                    "spectrogram-folder-path": spectrogram_folder_path,
-                    "welch-folder-path": welch_folder_path,
-                    "first": start,
-                    "last": stop,
-                    "downsampling-quality": resample_quality_settings["downsample"],
-                    "upsampling-quality": resample_quality_settings["upsample"],
-                    "umask": get_umask(),
-                    "multiprocessing": config.multiprocessing["is_active"],
-                    "nb-processes": config.multiprocessing["nb_processes"],
-                    "use-logging-setup": True,
-                    "dataset-json-path": self.folder / "project.json",
-                },
-                name=name + (f"_{index}" if len(batch_indexes) > 1 else ""),
-                output_folder=self.folder / self.SUBFOLDERS["log"],
-            )
-        self.job_builder.submit_pbs()
+
+        sds_json = (
+            spectro_dataset.folder / f"{spectro_dataset.name}.json"
+            if spectro_dataset is not None
+            else "None"
+        )
+
+        return ads_json, sds_json
 
     def _add_spectro_dataset(
         self,

src/osekit/utils/job.py

@@ -14,6 +14,8 @@
 
 from pandas import Timedelta
 
+from osekit.utils.core import file_indexes_per_batch
+
 if TYPE_CHECKING:
     from pathlib import Path
 
@@ -97,7 +99,7 @@ def __init__(
         """
         config = JobConfig() if config is None else config
         self.script_path = script_path
-        self.script_args = script_args if script_args else {}
+        self.script_args = script_args or {}
         self.nb_nodes = config.nb_nodes
         self.ncpus = config.ncpus
         self.ngpus = config.ngpus
@@ -535,6 +537,48 @@ def __init__(self, config: JobConfig = JobConfig) -> None:
         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,

tests/test_job.py

@@ -703,3 +703,97 @@ def test_build_dependency_string_with_different_types(
 def test_job_walltime(walltime: str | Timedelta) -> None:
     job = Job(Path(), config=JobConfig(walltime=walltime))
     assert Timedelta(job.walltime_str) == Timedelta(walltime)
+
+
+@pytest.mark.parametrize(
+    (
+        "nb_tasks",
+        "script_path",
+        "script_args",
+        "output_folder",
+        "job_name",
+        "nb_jobs",
+        "expected_task_indexes",
+    ),
+    [
+        pytest.param(
+            10,
+            Path("path/to/script.py"),
+            {"int_arg": 1, "str_arg": "cool"},
+            Path("path/to/output"),
+            "cool_name",
+            1,
+            [(0, 10)],
+            id="one_job_covers_all_tasks",
+        ),
+        pytest.param(
+            10,
+            Path("path/to/script.py"),
+            {"int_arg": 1, "str_arg": "cool"},
+            Path("path/to/output"),
+            "cool_name",
+            5,
+            [(0, 2), (2, 4), (4, 6), (6, 8), (8, 10)],
+            id="tasks_are_equally_distributed",
+        ),
+    ],
+)
+def test_create_jobs(  # noqa: PLR0917
+    monkeypatch: pytest.MonkeyPatch,
+    nb_tasks: int,
+    script_path: Path,
+    script_args: dict,
+    output_folder: Path,
+    job_name: str,
+    nb_jobs: int,
+    expected_task_indexes: list[tuple[int, int]],
+) -> None:
+    created_jobs = {}
+
+    def patch_create_job(self: JobBuilder, **kwargs: str) -> None:
+        job_name = kwargs.pop("name")
+        created_jobs[job_name] = kwargs
+
+    monkeypatch.setattr(JobBuilder, "create_job", patch_create_job)
+
+    JobBuilder().create_jobs(
+        nb_tasks=nb_tasks,
+        script_path=script_path,
+        script_args=script_args,
+        output_folder=output_folder,
+        job_name=job_name,
+        nb_jobs=nb_jobs,
+    )
+
+    # Correct number of jobs
+    assert len(created_jobs) == nb_jobs
+
+    # Correct distribution across jobs
+    for job in created_jobs.values():
+        assert (
+            job["script_args"]["first"],
+            job["script_args"]["last"],
+        ) in expected_task_indexes
+
+    # Script path
+    assert all(job["script_path"] == script_path for job in created_jobs.values())
+
+    # Script args
+    for job in created_jobs.values():
+        for arg in script_args:
+            assert arg in job["script_args"]
+
+    # Output folder
+    assert all(job["output_folder"] == output_folder for job in created_jobs.values())
+
+    # Job names
+    if nb_jobs == 1:
+        assert np.array_equal(list(created_jobs.keys()), [job_name])
+    else:
+        for idx, job in enumerate(
+            sorted(
+                created_jobs.items(),
+                key=lambda kvp: kvp[1]["script_args"]["first"],
+            ),
+        ):
+            assert job[0] == f"{job_name}_{idx}"

tests/test_public_api.py

@@ -1840,3 +1840,64 @@ def test_run_transform_with_same_name_in_different_process(
     transform.output_type = OutputType.SPECTROGRAM
     with pytest.raises(FileExistsError, match="already exists"):
         project.run(transform=transform)
+
+
+@pytest.mark.parametrize(
+    ("ads_folder_and_name", "sds_folder_and_name", "output_type", "expected"),
+    [
+        pytest.param(
+            (Path("cool"), "cool_ads"),
+            None,
+            OutputType.SPECTROGRAM,
+            ("None", "None"),
+            id="no_audio_output_is_none",
+        ),
+        pytest.param(
+            (Path("cool"), "cool_ads"),
+            None,
+            OutputType.AUDIO,
+            (Path(r"cool/cool_ads.json"), "None"),
+            id="audio_json_only",
+        ),
+        pytest.param(
+            (Path("cool"), "cool_ads"),
+            (Path("fun"), "fun_sds"),
+            OutputType.SPECTROGRAM,
+            ("None", Path(r"fun/fun_sds.json")),
+            id="spectro_json_only",
+        ),
+        pytest.param(
+            (Path("cool"), "cool_ads"),
+            (Path("fun"), "fun_sds"),
+            OutputType.AUDIO | OutputType.SPECTROGRAM,
+            (Path(r"cool/cool_ads.json"), Path(r"fun/fun_sds.json")),
+            id="both_ads_and_sds_jsons",
+        ),
+    ],
+)
+def test_get_json_paths(
+    monkeypatch: pytest.MonkeyPatch,
+    ads_folder_and_name: tuple[Path, str],
+    sds_folder_and_name: tuple[Path, str] | None,
+    output_type: OutputType,
+    expected: tuple[Path | str, Path | str],
+) -> None:
+    class DummyDataset:
+        def __init__(self, folder: Path, name: str) -> None:
+            self.folder = folder
+            self.name = name
+
+    monkeypatch.setattr("osekit.public.project.AudioDataset", DummyDataset)
+    monkeypatch.setattr("osekit.public.project.SpectroDataset", DummyDataset)
+
+    ads = DummyDataset(*ads_folder_and_name) if ads_folder_and_name else None
+    sds = DummyDataset(*sds_folder_and_name) if sds_folder_and_name else None
+
+    assert (
+        Project.get_json_paths(
+            audio_dataset=ads,
+            spectro_dataset=sds,
+            output_type=output_type,
+        )
+        == expected
+    )