278 lines
8.8 KiB
Python
278 lines
8.8 KiB
Python
"""Regression tests for streamed Modal progress forwarding."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import importlib.util
|
|
import sys
|
|
import types
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parents[1]
|
|
TEST_PACKAGE_NAME = "codex_modal_sync_testpkg"
|
|
|
|
|
|
def _clear_test_package_modules() -> None:
|
|
"""Remove any transient test-package modules from the interpreter cache."""
|
|
module_names = [
|
|
module_name
|
|
for module_name in sys.modules
|
|
if module_name == TEST_PACKAGE_NAME or module_name.startswith(f"{TEST_PACKAGE_NAME}.")
|
|
]
|
|
for module_name in module_names:
|
|
del sys.modules[module_name]
|
|
|
|
|
|
def _load_repo_module(module_name: str, file_path: Path) -> types.ModuleType:
|
|
"""Load one repository module under the transient test package."""
|
|
spec = importlib.util.spec_from_file_location(module_name, file_path)
|
|
if spec is None or spec.loader is None:
|
|
raise RuntimeError(f"Unable to create an import spec for {file_path}.")
|
|
module = importlib.util.module_from_spec(spec)
|
|
sys.modules[module_name] = module
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
def _load_modal_app_module() -> types.ModuleType:
|
|
"""Load `remote.modal_app` with package-relative imports resolved locally."""
|
|
_clear_test_package_modules()
|
|
|
|
package_module = types.ModuleType(TEST_PACKAGE_NAME)
|
|
package_module.__path__ = [str(REPO_ROOT)] # type: ignore[attr-defined]
|
|
sys.modules[TEST_PACKAGE_NAME] = package_module
|
|
|
|
remote_package_name = f"{TEST_PACKAGE_NAME}.remote"
|
|
remote_package = types.ModuleType(remote_package_name)
|
|
remote_package.__path__ = [str(REPO_ROOT / "remote")] # type: ignore[attr-defined]
|
|
sys.modules[remote_package_name] = remote_package
|
|
|
|
return _load_repo_module(
|
|
f"{remote_package_name}.modal_app",
|
|
REPO_ROOT / "remote" / "modal_app.py",
|
|
)
|
|
|
|
|
|
def test_mapped_stream_progress_preserves_real_node_id(monkeypatch: Any) -> None:
|
|
"""Mapped lane progress should target the real executing node, not the representative."""
|
|
modal_app = _load_modal_app_module()
|
|
emitted_progress: list[dict[str, Any]] = []
|
|
|
|
def capture_progress(**kwargs: Any) -> None:
|
|
"""Record forwarded local progress events for assertions."""
|
|
emitted_progress.append(kwargs)
|
|
|
|
monkeypatch.setattr(modal_app, "_emit_local_modal_progress", capture_progress)
|
|
|
|
payload = {
|
|
"prompt_id": "prompt-1",
|
|
"component_id": "component-1",
|
|
"component_node_ids": ["component-1", "node-a", "node-b"],
|
|
"extra_data": {"client_id": "client-1"},
|
|
"mapped_progress_lane_id": "7",
|
|
"map_item_index": 3,
|
|
}
|
|
stream_events = iter(
|
|
[
|
|
{
|
|
"kind": "progress",
|
|
"event_type": "node_progress",
|
|
"node_id": "component-1",
|
|
"display_node_id": "component-1",
|
|
"real_node_id": "node-b",
|
|
"value": 4,
|
|
"max": 9,
|
|
},
|
|
{
|
|
"kind": "result",
|
|
"outputs": [1],
|
|
},
|
|
]
|
|
)
|
|
|
|
with modal_app._MAPPED_PROGRESS_NODE_IDS_LOCK:
|
|
modal_app._MAPPED_PROGRESS_NODE_IDS.clear()
|
|
|
|
modal_app._consume_remote_payload_stream(payload, stream_events)
|
|
|
|
assert emitted_progress == [
|
|
{
|
|
"prompt_id": "prompt-1",
|
|
"client_id": "client-1",
|
|
"node_id": "component-1",
|
|
"value": 4.0,
|
|
"max_value": 9.0,
|
|
"display_node_id": "component-1",
|
|
"real_node_id": "node-b",
|
|
"lane_id": "7",
|
|
"clear": False,
|
|
"item_index": 3,
|
|
"aggregate_only": False,
|
|
}
|
|
]
|
|
with modal_app._MAPPED_PROGRESS_NODE_IDS_LOCK:
|
|
assert modal_app._MAPPED_PROGRESS_NODE_IDS[
|
|
("prompt-1", "component-1", "7")
|
|
] == "node-b"
|
|
|
|
emitted_progress.clear()
|
|
modal_app._clear_local_mapped_lane_progress(
|
|
payload,
|
|
lane_index=7,
|
|
item_index=3,
|
|
)
|
|
assert emitted_progress == [
|
|
{
|
|
"prompt_id": "prompt-1",
|
|
"client_id": "client-1",
|
|
"node_id": "node-b",
|
|
"value": 0.0,
|
|
"max_value": 1.0,
|
|
"display_node_id": "node-b",
|
|
"lane_id": "7",
|
|
"clear": True,
|
|
"item_index": 3,
|
|
}
|
|
]
|
|
|
|
|
|
def test_llm_stream_progress_preserves_stage_and_token_metrics(
|
|
monkeypatch: Any,
|
|
) -> None:
|
|
"""LLM labels, TTFT, and token rate should reach the local websocket."""
|
|
modal_app = _load_modal_app_module()
|
|
emitted_progress: list[dict[str, Any]] = []
|
|
monkeypatch.setattr(
|
|
modal_app,
|
|
"_emit_local_modal_progress",
|
|
lambda **kwargs: emitted_progress.append(kwargs),
|
|
)
|
|
payload = {
|
|
"prompt_id": "prompt-llm",
|
|
"component_id": "llm-node",
|
|
"component_node_ids": ["llm-node"],
|
|
"extra_data": {"client_id": "client-1"},
|
|
}
|
|
|
|
modal_app._consume_remote_payload_stream(
|
|
payload,
|
|
iter(
|
|
[
|
|
{
|
|
"kind": "progress",
|
|
"event_type": "node_progress",
|
|
"node_id": "llm-node",
|
|
"value": 12,
|
|
"max": 64,
|
|
"stage": "generating",
|
|
"message": "Generating",
|
|
"unit": "tokens",
|
|
"time_to_first_token_seconds": 2.5,
|
|
"tokens_per_second": 9.75,
|
|
},
|
|
{"kind": "result", "outputs": ["done"]},
|
|
]
|
|
),
|
|
)
|
|
|
|
assert emitted_progress == [
|
|
{
|
|
"prompt_id": "prompt-llm",
|
|
"client_id": "client-1",
|
|
"node_id": "llm-node",
|
|
"value": 12.0,
|
|
"max_value": 64.0,
|
|
"display_node_id": "llm-node",
|
|
"real_node_id": None,
|
|
"lane_id": None,
|
|
"clear": False,
|
|
"item_index": None,
|
|
"aggregate_only": False,
|
|
"stage": "generating",
|
|
"message": "Generating",
|
|
"unit": "tokens",
|
|
"time_to_first_token_seconds": 2.5,
|
|
"tokens_per_second": 9.75,
|
|
}
|
|
]
|
|
|
|
|
|
def test_stream_progress_reports_exact_modal_container_identity(monkeypatch: Any) -> None:
|
|
"""A Modal task id should follow its component progress into the node badge."""
|
|
modal_app = _load_modal_app_module()
|
|
emitted_progress: list[dict[str, Any]] = []
|
|
monkeypatch.setattr(
|
|
modal_app,
|
|
"_emit_local_modal_progress",
|
|
lambda **kwargs: emitted_progress.append(kwargs),
|
|
)
|
|
payload = {
|
|
"prompt_id": "prompt-location",
|
|
"component_id": "node-162",
|
|
"component_node_ids": ["node-162"],
|
|
"execution_provider": "modal",
|
|
"execution_environment_id": "modal:B300",
|
|
"extra_data": {"client_id": "client-1"},
|
|
}
|
|
|
|
modal_app._consume_remote_payload_stream(
|
|
payload,
|
|
iter(
|
|
[
|
|
{"kind": "remote_logs", "task_id": "ta-01K3MODAL"},
|
|
{
|
|
"kind": "progress",
|
|
"event_type": "node_progress",
|
|
"node_id": "node-162",
|
|
"value": 1,
|
|
"max": 8,
|
|
},
|
|
{"kind": "result", "outputs": ["done"]},
|
|
]
|
|
),
|
|
)
|
|
|
|
assert emitted_progress[0]["execution_provider"] == "modal"
|
|
assert emitted_progress[0]["execution_environment_id"] == "modal:B300"
|
|
assert emitted_progress[0]["execution_location"] == "ta-01K3MODAL"
|
|
|
|
|
|
def test_stream_status_reports_exact_modal_container_identity(monkeypatch: Any) -> None:
|
|
"""Node status should expose a Modal task id before numeric progress begins."""
|
|
modal_app = _load_modal_app_module()
|
|
emitted_status: list[dict[str, Any]] = []
|
|
monkeypatch.setattr(
|
|
modal_app,
|
|
"_emit_local_modal_status",
|
|
lambda **kwargs: emitted_status.append(kwargs),
|
|
)
|
|
payload = {
|
|
"prompt_id": "prompt-location",
|
|
"component_id": "node-162",
|
|
"component_node_ids": ["node-162"],
|
|
"execution_provider": "modal",
|
|
"execution_environment_id": "modal:B300",
|
|
"extra_data": {"client_id": "client-1"},
|
|
}
|
|
|
|
modal_app._consume_remote_payload_stream(
|
|
payload,
|
|
iter(
|
|
[
|
|
{"kind": "remote_logs", "task_id": "ta-01K3MODAL"},
|
|
{
|
|
"kind": "progress",
|
|
"event_type": "status",
|
|
"phase": "executing",
|
|
"active_node_id": "node-162",
|
|
},
|
|
{"kind": "result", "outputs": ["done"]},
|
|
]
|
|
),
|
|
)
|
|
|
|
assert emitted_status[0]["execution_provider"] == "modal"
|
|
assert emitted_status[0]["execution_environment_id"] == "modal:B300"
|
|
assert emitted_status[0]["execution_location"] == "ta-01K3MODAL"
|