| import json |
| import threading |
| import time |
| from concurrent.futures import ThreadPoolExecutor |
| from pathlib import Path |
|
|
| from src import downloads, jobs, preparation, service |
|
|
|
|
| def _prepared_job(job, token): |
| prepared = job.input_dir / "00_input.wav" |
| prepared.write_bytes(job.job_id.encode("ascii")) |
| config = { |
| "schema": preparation.PREPARATION_SCHEMA, |
| "reproducibility_schema": preparation.REPRODUCIBILITY_SCHEMA, |
| "status": "PREPARED", |
| "job_id": job.job_id, |
| "access_token": token, |
| "preparation_wall_seconds": 0.1, |
| "inputs": [{"prepared_path": "input/00_input.wav"}], |
| "input_total_seconds": 5.0, |
| "models": { |
| "selected_filenames": ["UVR-De-Reverb-aufr33-jarredou.pth"], |
| "package_filenames": ["UVR-De-Reverb-aufr33-jarredou.pth"], |
| "ensemble_algorithm": "avg_wave", |
| "custom_records": [], |
| }, |
| "output": { |
| "format": "FLAC", "bitrate": "Auto", "sample_rate": 44100, |
| "normalization_threshold": 0.9, "amplification_threshold": 0.0, |
| "single_stem": "All stems", "bundle_layout": "Flat outputs", |
| }, |
| "parameters": {"pitch_shift": 0, "chunk_duration": 0, "allow_cpu_fallback": True, "log_level": "INFO"}, |
| "duration": {"seconds": 30}, |
| "providers": {}, |
| "outputs": [], |
| } |
| config["config_sha256"] = preparation._digest_for_private_config(config) |
| (job.config_dir / "preparation.json").write_text(json.dumps(config), encoding="utf-8") |
| (job.config_dir / "preflight.json").write_text("{}\n", encoding="utf-8") |
| (job.config_dir / ".prepared").write_text(config["config_sha256"] + "\n", encoding="utf-8") |
| (job.logs_dir / f"sesa_prepare_{job.job_id}.log").write_text("prepared\n", encoding="utf-8") |
| return preparation.prepared_state(job.job_id, token) |
|
|
|
|
| def test_overlapping_prepared_jobs_keep_outputs_and_reproducibility_isolated(monkeypatch, tmp_path): |
| monkeypatch.setattr(jobs, "JOB_ROOT", tmp_path / "jobs") |
| first = jobs.create_job() |
| second = jobs.create_job() |
| states = [_prepared_job(first, "token-a"), _prepared_job(second, "token-b")] |
| barrier = threading.Barrier(2) |
|
|
| class FakeSeparator: |
| def __init__(self, output_dir): |
| self.output_dir = Path(output_dir) |
|
|
| def load_model(self, value): |
| del value |
| barrier.wait(timeout=5) |
|
|
| def separate(self, values): |
| del values |
| output = self.output_dir / "00_input_(Vocals)_model.flac" |
| output.write_bytes(self.output_dir.parent.name.encode("ascii")) |
| return [str(output)] |
|
|
| monkeypatch.setattr(service, "create_separator", lambda **kwargs: FakeSeparator(kwargs["output_dir"])) |
| monkeypatch.setattr(service.torch.cuda, "is_available", lambda: False) |
| monkeypatch.setattr(service, "release_accelerators", lambda separator=None: None) |
|
|
| with ThreadPoolExecutor(max_workers=2) as pool: |
| results = list(pool.map(service.run_prepared_job, states)) |
|
|
| assert all(result[0].startswith("### Completed") for result in results) |
| output_paths = [Path(result[3][0]) for result in results] |
| assert output_paths[0].parent != output_paths[1].parent |
| assert first.job_id in output_paths[0].name or first.job_id in output_paths[1].name |
| assert second.job_id in output_paths[0].name or second.job_id in output_paths[1].name |
| for result, expected_job in zip(results, (first, second)): |
| output = Path(result[3][0]) |
| repro = json.loads(Path(result[7]).read_text(encoding="utf-8")) |
| assert output.parent == expected_job.output_dir |
| assert repro["job_id"] == expected_job.job_id |
| assert all(expected_job.job_id in record["name"] for record in repro["outputs"]) |
| other = second.job_id if expected_job is first else first.job_id |
| assert other not in json.dumps(repro) |
|
|
| preparation.cleanup_prepared_state(states[0]) |
| assert not first.root.exists() |
| assert second.root.exists() |
|
|
|
|
| def test_overlapping_package_prefetch_uses_one_shared_cache_download(monkeypatch, tmp_path): |
| model_root = tmp_path / "models" |
| download_root = tmp_path / "downloads" |
| monkeypatch.setattr(downloads, "MODEL_ROOT", model_root) |
| monkeypatch.setattr(downloads, "DOWNLOAD_ROOT", download_root) |
| monkeypatch.setattr(downloads.shutil, "which", lambda name: "/usr/bin/audio-separator") |
| calls = [] |
| calls_lock = threading.Lock() |
|
|
| def fake_run(command, **kwargs): |
| del command, kwargs |
| with calls_lock: |
| calls.append(time.monotonic()) |
| time.sleep(0.1) |
| model_root.mkdir(parents=True, exist_ok=True) |
| (model_root / "UVR_MDXNET_KARA_2.onnx").write_bytes(b"model") |
|
|
| monkeypatch.setattr(downloads.subprocess, "run", fake_run) |
| with ThreadPoolExecutor(max_workers=2) as pool: |
| records = list(pool.map( |
| lambda _: downloads.prefetch_package_models(["UVR_MDXNET_KARA_2.onnx"]), |
| range(2), |
| )) |
| assert len(calls) == 1 |
| assert all(result[0]["cache_after"] is True for result in records) |
| assert records[0][0]["files"][0]["sha256"] == records[1][0]["files"][0]["sha256"] |
|
|
