import json import zipfile from pathlib import Path from src import service from src.jobs import JobPaths def _job(tmp_path: Path) -> JobPaths: root = tmp_path / "job" job = JobPaths("a" * 32, root, root / "input", root / "work", root / "output") for path in (job.input_dir, job.work_dir, job.output_dir, job.logs_dir, job.bundle_dir, job.config_dir): path.mkdir(parents=True, exist_ok=True) return job def _config(job: JobPaths, prepared: list[Path]) -> dict: return { "schema": "sesa-preparation-v21", "reproducibility_schema": "sesa-reproducibility-v21", "status": "PREPARED", "job_id": job.job_id, "preparation_wall_seconds": 1.0, "input_total_seconds": 30.0, "inputs": [ { "display_name": f"source-{index}.wav", "prepared_path": path.relative_to(job.root).as_posix(), "prepared_sha256": f"sha-{index}", "prepared_duration_seconds": 15.0, } for index, path in enumerate(prepared) ], "batch": { "schema": "sesa-batch-plan-v1", "enabled": True, "input_count": 2, "execution_order": "sequential-inputs-shared-loaded-model", "model_load_scope": "once-per-job", "continue_on_item_error": True, "output_mapping": "explicit-per-separate-call", "status": "PREPARED", }, "models": { "selected_filenames": ["model.ckpt"], "custom_records": [], "ensemble_algorithm": "avg_wave", }, "output": { "format": "FLAC", "bitrate": "Auto", "sample_rate": 44100, "normalization_threshold": 0.9, "amplification_threshold": 0.0, "single_stem": "All stems", "bundle_layout": "Group by input", }, "parameters": { "pitch_shift": 0, "chunk_duration": 0, "allow_cpu_fallback": False, "log_level": "INFO", }, "duration": {"seconds": 180}, } def _patch_runtime(monkeypatch, job: JobPaths, config: dict, fake_separator): monkeypatch.setattr(service, "parse_prepared_state", lambda value: (job, config)) monkeypatch.setattr(service, "_custom_records_from_config", lambda value: []) monkeypatch.setattr(service.torch.cuda, "is_available", lambda: True) monkeypatch.setattr(service.torch.cuda, "get_device_name", lambda index: "Fake CUDA") monkeypatch.setattr(service, "create_separator", lambda **kwargs: fake_separator) monkeypatch.setattr(service, "release_accelerators", lambda separator=None: None) def timed(function, use_cuda): return function(), 0.1, None monkeypatch.setattr(service, "synchronized_wall_time", timed) def test_batch_runs_each_item_sequentially_with_one_model_load(monkeypatch, tmp_path): job = _job(tmp_path) prepared = [] for index in range(2): path = job.work_dir / f"{index:02d}_mix.wav" path.write_bytes(f"input-{index}".encode()) prepared.append(path) config = _config(job, prepared) class FakeSeparator: def __init__(self): self.load_calls = 0 self.separate_calls = [] def load_model(self, value): self.load_calls += 1 def separate(self, values): assert len(values) == 1 source = Path(values[0]) self.separate_calls.append(source.name) output = job.output_dir / f"{source.stem}_(Vocals).flac" output.write_bytes(source.name.encode()) return [str(output)] separator = FakeSeparator() _patch_runtime(monkeypatch, job, config, separator) result = service.run_prepared_job("state") status, _, _, outputs, archive, _, _, repro_path = result assert status.startswith("### Completed") assert separator.load_calls == 1 assert separator.separate_calls == [path.name for path in prepared] assert len(outputs) == 3 stem_outputs = [Path(value) for value in outputs if Path(value).parent == job.output_dir] assert len(stem_outputs) == 2 assert Path(outputs[-1]).name == "batch_manifest.json" repro = json.loads(Path(repro_path).read_text(encoding="utf-8")) batch = repro["execution"]["batch"] assert batch["status"] == "COMPLETED" assert batch["completed_count"] == 2 assert batch["failed_count"] == 0 assert batch["model_load_count"] == 1 assert sorted(item["input_index"] for item in repro["outputs"]) == [0, 1] assert all("_input_" in path.name for path in stem_outputs) with zipfile.ZipFile(archive) as bundle: names = bundle.namelist() assert "diagnostics/batch_manifest.json" in names assert any(name.startswith("outputs/input_00/") for name in names) assert any(name.startswith("outputs/input_01/") for name in names) def test_batch_partial_failure_keeps_successful_item_and_records_error(monkeypatch, tmp_path): job = _job(tmp_path) prepared = [] for index in range(2): path = job.work_dir / f"{index:02d}_mix.wav" path.write_bytes(f"input-{index}".encode()) prepared.append(path) config = _config(job, prepared) class FakeSeparator: def load_model(self, value): pass def separate(self, values): source = Path(values[0]) if source.name.startswith("01_"): raise RuntimeError("forced second-item failure") output = job.output_dir / f"{source.stem}_(Vocals).flac" output.write_bytes(b"success") return [str(output)] _patch_runtime(monkeypatch, job, config, FakeSeparator()) result = service.run_prepared_job("state") status, _, _, outputs, archive, _, _, repro_path = result assert status.startswith("### Completed with item failures") assert len(outputs) == 2 stem_outputs = [Path(value) for value in outputs if Path(value).parent == job.output_dir] assert len(stem_outputs) == 1 assert Path(outputs[-1]).name == "batch_manifest.json" repro = json.loads(Path(repro_path).read_text(encoding="utf-8")) assert repro["status"] == "COMPLETED_WITH_ITEM_FAILURES" batch = repro["execution"]["batch"] assert batch["completed_count"] == 1 assert batch["failed_count"] == 1 assert batch["items"][1]["status"] == "FAILED" assert "forced second-item failure" in batch["items"][1]["error"]["message"] with zipfile.ZipFile(archive) as bundle: assert any("batch_item_01_error.txt" in name for name in bundle.namelist()) def test_retry_runs_only_failed_item_and_keeps_prior_outputs(monkeypatch, tmp_path): job = _job(tmp_path) prepared = [] for index in range(3): path = job.work_dir / f"{index:02d}_mix.wav" path.write_bytes(f"input-{index}".encode()) prepared.append(path) config = _config(job, prepared) config["inputs"] = [ { "display_name": f"source-{index}.wav", "prepared_path": path.relative_to(job.root).as_posix(), "prepared_sha256": f"sha-{index}", "prepared_duration_seconds": 10.0, } for index, path in enumerate(prepared) ] config["batch"]["input_count"] = 3 class FakeSeparator: def __init__(self): self.load_calls = 0 self.separate_calls = [] def load_model(self, value): self.load_calls += 1 def separate(self, values): source = Path(values[0]) self.separate_calls.append(source.name) if getattr(self, "fail_item_one", False) and source.name.startswith("01_"): raise RuntimeError("forced retryable item failure") output = job.output_dir / f"{source.stem}_(Vocals).flac" output.write_bytes(source.name.encode()) return [str(output)] first_separator = FakeSeparator() first_separator.fail_item_one = True _patch_runtime(monkeypatch, job, config, first_separator) first = service.run_prepared_job("state") first_manifest = json.loads((job.config_dir / "batch_manifest.json").read_text(encoding="utf-8")) assert first[0].startswith("### Completed with item failures") assert first_manifest["retryable_input_indexes"] == [1] assert first_manifest["completed_count"] == 2 assert first_separator.separate_calls == [path.name for path in prepared] second_separator = FakeSeparator() monkeypatch.setattr(service, "create_separator", lambda **kwargs: second_separator) second = service.run_prepared_job( "state", execution_mode=service.EXECUTION_MODE_RETRY_INCOMPLETE ) final_manifest = json.loads((job.config_dir / "batch_manifest.json").read_text(encoding="utf-8")) assert second[0].startswith("### Completed") assert second_separator.separate_calls == [prepared[1].name] assert final_manifest["status"] == "COMPLETED" assert final_manifest["completed_count"] == 3 assert final_manifest["retryable_input_indexes"] == [] assert len(final_manifest["attempts"]) == 2 assert final_manifest["attempts"][1]["target_input_indexes"] == [1] assert final_manifest["total_model_load_count"] == 2 assert sorted(item["input_index"] for item in final_manifest["outputs"]) == [0, 1, 2] def test_cooperative_stop_marks_remaining_items_retryable(monkeypatch, tmp_path): job = _job(tmp_path) prepared = [] for index in range(3): path = job.work_dir / f"{index:02d}_mix.wav" path.write_bytes(f"input-{index}".encode()) prepared.append(path) config = _config(job, prepared) config["inputs"] = [ { "display_name": f"source-{index}.wav", "prepared_path": path.relative_to(job.root).as_posix(), "prepared_sha256": f"sha-{index}", "prepared_duration_seconds": 10.0, } for index, path in enumerate(prepared) ] config["batch"]["input_count"] = 3 class FakeSeparator: def load_model(self, value): pass def __init__(self): self.calls = 0 def separate(self, values): source = Path(values[0]) output = job.output_dir / f"{source.stem}_(Vocals).flac" output.write_bytes(source.name.encode()) self.calls += 1 if self.calls == 1: service._cancel_request_path(job).write_text("{}\n", encoding="utf-8") return [str(output)] _patch_runtime(monkeypatch, job, config, FakeSeparator()) result = service.run_prepared_job("state") manifest = json.loads((job.config_dir / "batch_manifest.json").read_text(encoding="utf-8")) assert result[0].startswith("### Stopped with partial results") assert manifest["status"] == "CANCELED_WITH_PARTIAL_RESULTS" assert manifest["completed_count"] == 1 assert manifest["canceled_count"] == 2 assert manifest["retryable_input_indexes"] == [1, 2] assert manifest["attempts"][0]["attempted_input_indexes"] == [0] assert manifest["attempts"][0]["cancel_observed"] is True