88 lines
2.8 KiB
Python
88 lines
2.8 KiB
Python
|
import json
|
||
|
|
import asyncio
|
||
|
|
import sys
|
||
|
|
import types
|
||
|
|
from dataclasses import dataclass
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
ROOT = Path(__file__).resolve().parents[2]
|
||
|
|
if str(ROOT) not in sys.path:
|
||
|
|
sys.path.insert(0, str(ROOT))
|
||
|
|
|
||
|
|
from gfa_core.mas.stage_catalog import ( # noqa: E402
|
||
|
|
clamp_stage_cpu,
|
||
|
|
get_stage_queue,
|
||
|
|
get_stage_resource_profile,
|
||
|
|
)
|
||
|
|
from gfa_core.mas.tools.stage_tools import StageTool, StageToolExecutor # noqa: E402
|
||
|
|
|
||
|
|
|
||
|
|
def test_stage_queue_and_cpu_profiles_are_resource_aware(monkeypatch, tmp_path):
|
||
|
|
monkeypatch.setenv("BBAV5_QUEUE_PREFIX", "copy_a")
|
||
|
|
monkeypatch.setenv("BBAV5_RESOURCE_STATE_DIR", str(tmp_path / "resource_governor"))
|
||
|
|
|
||
|
|
assert get_stage_queue("phylogeny") == "copy_a.heavy"
|
||
|
|
assert get_stage_queue("identify") == "copy_a.medium"
|
||
|
|
assert get_stage_queue("ppi") == "copy_a.light"
|
||
|
|
assert get_stage_queue("rnaseq_align") == "copy_a.rnaseq"
|
||
|
|
|
||
|
|
assert clamp_stage_cpu("phylogeny", 32) == 6
|
||
|
|
assert clamp_stage_cpu("identify", 32) == 4
|
||
|
|
assert clamp_stage_cpu("promoter", 8) == 4
|
||
|
|
assert clamp_stage_cpu("rnaseq_align", 32) == 20
|
||
|
|
|
||
|
|
profile = get_stage_resource_profile("interspecies_synteny")
|
||
|
|
assert profile["resource_class"] == "heavy"
|
||
|
|
assert profile["max_cpu"] == 6
|
||
|
|
|
||
|
|
|
||
|
|
def test_stage_tool_executor_clamps_cpu_before_real_run_stage(monkeypatch, tmp_path):
|
||
|
|
monkeypatch.setenv("BBAV5_RESOURCE_STATE_DIR", str(tmp_path / "resource_governor"))
|
||
|
|
captured = {}
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class FakeJobInfo:
|
||
|
|
job_id: str
|
||
|
|
input_dir: str
|
||
|
|
workspace_dir: str
|
||
|
|
logs_dir: str
|
||
|
|
created_at: str
|
||
|
|
|
||
|
|
def fake_run_stage(**kwargs):
|
||
|
|
captured.update(kwargs)
|
||
|
|
final_dir = Path(kwargs["job"].workspace_dir) / "final"
|
||
|
|
final_dir.mkdir(parents=True, exist_ok=True)
|
||
|
|
manifest = final_dir / f"{kwargs['prefix']}.{kwargs['stage']}.manifest.json"
|
||
|
|
manifest.write_text(json.dumps({"status": "success", "stage": kwargs["stage"]}), encoding="utf-8")
|
||
|
|
return "✅ fake stage completed", manifest
|
||
|
|
|
||
|
|
fake_module = types.ModuleType("apps.web.pipeline_backend")
|
||
|
|
fake_module.JobInfo = FakeJobInfo
|
||
|
|
fake_module.run_stage = fake_run_stage
|
||
|
|
monkeypatch.setitem(sys.modules, "apps.web.pipeline_backend", fake_module)
|
||
|
|
|
||
|
|
executor = StageToolExecutor(project_root=tmp_path)
|
||
|
|
stage_tool = StageTool(
|
||
|
|
id="phylogeny",
|
||
|
|
name="Phylogeny",
|
||
|
|
script="apps/cli/run_phylogeny.py",
|
||
|
|
category="analysis",
|
||
|
|
description="fake",
|
||
|
|
)
|
||
|
|
|
||
|
|
result = asyncio.run(
|
||
|
|
executor.execute(
|
||
|
|
stage_tool=stage_tool,
|
||
|
|
job_id="job_resource_test",
|
||
|
|
prefix="Ft.SPL",
|
||
|
|
extra_params={"cpu": 32},
|
||
|
|
)
|
||
|
|
)
|
||
|
|
|
||
|
|
assert result["status"] == "success"
|
||
|
|
assert result["requested_cpu"] == 32
|
||
|
|
assert result["effective_cpu"] == 6
|
||
|
|
assert captured["cpu"] == 6
|