diff --git a/src/osekit/public/project.py b/src/osekit/public/project.py index ba8e64d1..ff9fff0c 100644 --- a/src/osekit/public/project.py +++ b/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, diff --git a/src/osekit/utils/job.py b/src/osekit/utils/job.py index edb78931..63494ae8 100644 --- a/src/osekit/utils/job.py +++ b/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, diff --git a/tests/test_job.py b/tests/test_job.py index 76633663..f3b4c842 100644 --- a/tests/test_job.py +++ b/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}" diff --git a/tests/test_public_api.py b/tests/test_public_api.py index 842df70a..478dab16 100644 --- a/tests/test_public_api.py +++ b/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 + )