73 lines
2.5 KiB
Python
73 lines
2.5 KiB
Python
import asyncio
|
|
|
|
from app import jobs
|
|
from app.commands import CommandSpec
|
|
from app.jobs import default_conda_env_for_job
|
|
from app.main import _job_with_progress
|
|
from app.schemas import JobCreate
|
|
|
|
|
|
async def _stream_text(response) -> str:
|
|
chunks = []
|
|
async for chunk in response.body_iterator:
|
|
chunks.append(chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk)
|
|
return "".join(chunks)
|
|
|
|
|
|
def test_mmseg_jobs_use_mmseg_conda_env_by_default():
|
|
assert default_conda_env_for_job("mmseg.train") == "seg_mmcv"
|
|
assert default_conda_env_for_job("segmodel.train") == "seg_smp"
|
|
|
|
|
|
def test_build_task_sets_cuda_visible_devices(tmp_path, monkeypatch):
|
|
def fake_build_module_task(job_type, params, conda_env):
|
|
return CommandSpec(["python", "-c", "print('ok')"], tmp_path, "fake task")
|
|
|
|
monkeypatch.setattr(jobs, "build_module_task", fake_build_module_task)
|
|
spec = jobs._build_task(JobCreate(type="mock.echo", params={}, gpus=[2, 0]))
|
|
|
|
assert spec.env["CUDA_VISIBLE_DEVICES"] == "2,0"
|
|
|
|
|
|
def test_build_task_uses_requested_conda_env(tmp_path, monkeypatch):
|
|
captured = {}
|
|
|
|
def fake_build_module_task(job_type, params, conda_env):
|
|
captured["conda_env"] = conda_env
|
|
return CommandSpec(["python", "-c", "print('ok')"], tmp_path, "fake task")
|
|
|
|
monkeypatch.setattr(jobs, "build_module_task", fake_build_module_task)
|
|
jobs._build_task(JobCreate(type="segmodel.train", params={}, conda_env="seg_custom"))
|
|
|
|
assert captured["conda_env"] == "seg_custom"
|
|
|
|
|
|
def test_job_progress_reports_log_size(tmp_path):
|
|
log_path = tmp_path / "job.log"
|
|
log_path.write_text("line one\nline two\n", encoding="utf-8")
|
|
job = {"id": "job1", "status": "running", "log_path": str(log_path)}
|
|
|
|
enriched = _job_with_progress(job, include_log_tail=True)
|
|
|
|
assert enriched["log_size"] == log_path.stat().st_size
|
|
assert enriched["log_tail"] == "line one\nline two\n"
|
|
|
|
|
|
def test_job_events_respect_log_offset(tmp_path, monkeypatch):
|
|
from app import main
|
|
|
|
log_path = tmp_path / "job.log"
|
|
old_chunk = "old chunk\n"
|
|
log_path.write_text(old_chunk + "new chunk\n", encoding="utf-8")
|
|
|
|
def fake_get_job(job_id):
|
|
return {"id": job_id, "status": "success", "log_path": str(log_path)}
|
|
|
|
monkeypatch.setattr(main.db, "get_job", fake_get_job)
|
|
response = asyncio.run(main.api_job_events("job1", offset=len(old_chunk.encode("utf-8"))))
|
|
text = asyncio.run(_stream_text(response))
|
|
|
|
assert response.media_type == "text/event-stream"
|
|
assert "new chunk" in text
|
|
assert "old chunk" not in text
|