403 lines
14 KiB
Python
403 lines
14 KiB
Python
import importlib.util
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import types
|
|
import unittest
|
|
from collections import deque
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
|
|
class _FakeResponse:
|
|
def __init__(self, payload, status=200):
|
|
self.payload = payload
|
|
self.status = status
|
|
|
|
|
|
class _FakeRequest:
|
|
def __init__(self, payload=None, match_info=None, query=None):
|
|
self._payload = payload
|
|
self.match_info = match_info or {}
|
|
self.query = query or {}
|
|
|
|
async def json(self):
|
|
return self._payload
|
|
|
|
|
|
class _FakeHTTPClientResponse:
|
|
def __init__(self, payload, status=200):
|
|
self._payload = payload
|
|
self.status = status
|
|
|
|
async def __aenter__(self):
|
|
return self
|
|
|
|
async def __aexit__(self, _exc_type, _exc, _tb):
|
|
return False
|
|
|
|
async def json(self):
|
|
return self._payload
|
|
|
|
async def text(self):
|
|
return str(self._payload)
|
|
|
|
|
|
class _FakeHTTPClientSession:
|
|
def __init__(self, payload, status=200):
|
|
self._payload = payload
|
|
self._status = status
|
|
self.calls = []
|
|
|
|
def get(self, url, params=None, timeout=None):
|
|
self.calls.append({"url": url, "params": params, "timeout": timeout})
|
|
return _FakeHTTPClientResponse(self._payload, status=self._status)
|
|
|
|
|
|
class _DummyWorkerManager:
|
|
def __init__(self):
|
|
self.processes = {}
|
|
|
|
def launch_worker(self, worker):
|
|
worker_id = str(worker["id"])
|
|
self.processes[worker_id] = {
|
|
"pid": 12345,
|
|
"log_file": f"/tmp/distributed_worker_{worker_id}.log",
|
|
"process": None,
|
|
}
|
|
return 12345
|
|
|
|
def _is_process_running(self, _pid):
|
|
return False
|
|
|
|
def save_processes(self):
|
|
return None
|
|
|
|
def stop_worker(self, _worker_id):
|
|
return True, "Stopped"
|
|
|
|
def get_managed_workers(self):
|
|
return []
|
|
|
|
|
|
class _ImmediateLoop:
|
|
async def run_in_executor(self, _executor, func, *args):
|
|
return func(*args)
|
|
|
|
|
|
def _load_worker_routes_module():
|
|
module_path = Path(__file__).resolve().parents[2] / "api" / "worker_routes.py"
|
|
package_name = "dist_api_worker_testpkg"
|
|
|
|
for mod_name in list(sys.modules):
|
|
if mod_name == package_name or mod_name.startswith(f"{package_name}."):
|
|
del sys.modules[mod_name]
|
|
|
|
root_pkg = types.ModuleType(package_name)
|
|
root_pkg.__path__ = []
|
|
sys.modules[package_name] = root_pkg
|
|
|
|
api_pkg = types.ModuleType(f"{package_name}.api")
|
|
api_pkg.__path__ = []
|
|
sys.modules[f"{package_name}.api"] = api_pkg
|
|
|
|
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
|
utils_pkg.__path__ = []
|
|
sys.modules[f"{package_name}.utils"] = utils_pkg
|
|
|
|
workers_pkg = types.ModuleType(f"{package_name}.workers")
|
|
workers_pkg.__path__ = []
|
|
workers_pkg.get_worker_manager = lambda: _DummyWorkerManager()
|
|
sys.modules[f"{package_name}.workers"] = workers_pkg
|
|
|
|
detection_module = types.ModuleType(f"{package_name}.workers.detection")
|
|
detection_module.is_local_worker = lambda *_args, **_kwargs: True
|
|
detection_module.is_same_physical_host = lambda *_args, **_kwargs: True
|
|
detection_module.get_machine_id = lambda: "machine-id"
|
|
detection_module.is_docker_environment = lambda: False
|
|
detection_module.is_runpod_environment = lambda: False
|
|
detection_module.get_comms_channel = lambda *_args, **_kwargs: "lan"
|
|
sys.modules[f"{package_name}.workers.detection"] = detection_module
|
|
|
|
created_aiohttp_stub = False
|
|
if "aiohttp" not in sys.modules:
|
|
created_aiohttp_stub = True
|
|
aiohttp_module = types.ModuleType("aiohttp")
|
|
|
|
class _ClientTimeout:
|
|
def __init__(self, total=None):
|
|
self.total = total
|
|
|
|
class _WSMsgType:
|
|
TEXT = "TEXT"
|
|
ERROR = "ERROR"
|
|
CLOSED = "CLOSED"
|
|
|
|
class _WebSocketResponse:
|
|
def __init__(self, *args, **kwargs):
|
|
self.args = args
|
|
self.kwargs = kwargs
|
|
|
|
async def prepare(self, _request):
|
|
return None
|
|
|
|
async def send_json(self, _payload):
|
|
return None
|
|
|
|
def __aiter__(self):
|
|
async def _empty():
|
|
if False:
|
|
yield None
|
|
return _empty()
|
|
|
|
aiohttp_module.ClientTimeout = _ClientTimeout
|
|
aiohttp_module.WSMsgType = _WSMsgType
|
|
aiohttp_module.web = types.SimpleNamespace(
|
|
json_response=lambda payload, status=200: _FakeResponse(payload, status=status),
|
|
WebSocketResponse=_WebSocketResponse,
|
|
)
|
|
sys.modules["aiohttp"] = aiohttp_module
|
|
|
|
class _Routes:
|
|
def get(self, _path):
|
|
def _decorator(fn):
|
|
return fn
|
|
return _decorator
|
|
|
|
def post(self, _path):
|
|
def _decorator(fn):
|
|
return fn
|
|
return _decorator
|
|
|
|
server_module = types.ModuleType("server")
|
|
server_module.PromptServer = types.SimpleNamespace(instance=types.SimpleNamespace(routes=_Routes()))
|
|
sys.modules["server"] = server_module
|
|
|
|
created_torch_stub = False
|
|
if "torch" not in sys.modules:
|
|
created_torch_stub = True
|
|
torch_module = types.ModuleType("torch")
|
|
torch_module.cuda = types.SimpleNamespace(
|
|
is_available=lambda: False,
|
|
empty_cache=lambda: None,
|
|
ipc_collect=lambda: None,
|
|
current_device=lambda: 0,
|
|
device_count=lambda: 0,
|
|
)
|
|
sys.modules["torch"] = torch_module
|
|
|
|
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
|
logging_module.debug_log = lambda *_args, **_kwargs: None
|
|
logging_module.log = lambda *_args, **_kwargs: None
|
|
sys.modules[f"{package_name}.utils.logging"] = logging_module
|
|
|
|
config_module = types.ModuleType(f"{package_name}.utils.config")
|
|
config_module.load_config = lambda: {"workers": []}
|
|
sys.modules[f"{package_name}.utils.config"] = config_module
|
|
|
|
network_module = types.ModuleType(f"{package_name}.utils.network")
|
|
|
|
async def _handle_api_error(_request, error, status=500):
|
|
return _FakeResponse({"status": "error", "message": str(error)}, status=status)
|
|
|
|
network_module.handle_api_error = _handle_api_error
|
|
network_module.normalize_host = lambda value: value
|
|
network_module.build_worker_url = lambda worker, endpoint="": f"http://localhost:{worker.get('port', 8188)}{endpoint}"
|
|
|
|
async def _probe_worker(*_args, **_kwargs):
|
|
return None
|
|
|
|
network_module.probe_worker = _probe_worker
|
|
|
|
async def _get_client_session():
|
|
raise RuntimeError("not used in these tests")
|
|
|
|
network_module.get_client_session = _get_client_session
|
|
sys.modules[f"{package_name}.utils.network"] = network_module
|
|
|
|
constants_module = types.ModuleType(f"{package_name}.utils.constants")
|
|
constants_module.CHUNK_SIZE = 8192
|
|
sys.modules[f"{package_name}.utils.constants"] = constants_module
|
|
|
|
async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
|
|
|
|
async def _queue_prompt_payload(*_args, **_kwargs):
|
|
return "prompt-id"
|
|
|
|
async_helpers_module.queue_prompt_payload = _queue_prompt_payload
|
|
|
|
class _PromptValidationError(RuntimeError):
|
|
def __init__(self, message="invalid prompt", validation_error=None, node_errors=None):
|
|
super().__init__(message)
|
|
self.validation_error = validation_error if isinstance(validation_error, dict) else {}
|
|
self.node_errors = node_errors if isinstance(node_errors, dict) else {}
|
|
|
|
async_helpers_module.PromptValidationError = _PromptValidationError
|
|
sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
|
|
|
|
schemas_module = types.ModuleType(f"{package_name}.api.schemas")
|
|
|
|
def _require_fields(data, *fields):
|
|
missing = []
|
|
for field in fields:
|
|
value = data.get(field) if isinstance(data, dict) else None
|
|
if value is None or (isinstance(value, str) and not value.strip()):
|
|
missing.append(field)
|
|
return missing
|
|
|
|
def _validate_worker_id(worker_id, config):
|
|
return any(str(worker.get("id")) == str(worker_id) for worker in config.get("workers", []))
|
|
|
|
schemas_module.require_fields = _require_fields
|
|
schemas_module.validate_worker_id = _validate_worker_id
|
|
sys.modules[f"{package_name}.api.schemas"] = schemas_module
|
|
|
|
spec = importlib.util.spec_from_file_location(f"{package_name}.api.worker_routes", module_path)
|
|
module = importlib.util.module_from_spec(spec)
|
|
assert spec is not None and spec.loader is not None
|
|
spec.loader.exec_module(module)
|
|
|
|
if created_aiohttp_stub:
|
|
sys.modules.pop("aiohttp", None)
|
|
if created_torch_stub:
|
|
sys.modules.pop("torch", None)
|
|
|
|
return module
|
|
|
|
|
|
worker_routes = _load_worker_routes_module()
|
|
|
|
|
|
class WorkerRoutesTests(unittest.IsolatedAsyncioTestCase):
|
|
async def test_launch_worker_valid_id_returns_200(self):
|
|
manager = _DummyWorkerManager()
|
|
config = {"workers": [{"id": "worker-a", "name": "Worker A", "port": 8188}]}
|
|
request = _FakeRequest({"worker_id": "worker-a"})
|
|
|
|
with patch.object(worker_routes, "get_worker_manager", return_value=manager), patch.object(
|
|
worker_routes, "load_config", return_value=config
|
|
), patch.object(
|
|
worker_routes.asyncio, "get_running_loop", return_value=_ImmediateLoop()
|
|
):
|
|
response = await worker_routes.launch_worker_endpoint(request)
|
|
|
|
self.assertEqual(response.status, 200)
|
|
self.assertEqual(response.payload.get("status"), "success")
|
|
self.assertEqual(response.payload.get("pid"), 12345)
|
|
|
|
async def test_launch_worker_unknown_id_returns_404(self):
|
|
manager = _DummyWorkerManager()
|
|
config = {"workers": [{"id": "worker-a", "name": "Worker A", "port": 8188}]}
|
|
request = _FakeRequest({"worker_id": "missing-worker"})
|
|
|
|
with patch.object(worker_routes, "get_worker_manager", return_value=manager), patch.object(
|
|
worker_routes, "load_config", return_value=config
|
|
):
|
|
response = await worker_routes.launch_worker_endpoint(request)
|
|
|
|
self.assertEqual(response.status, 404)
|
|
self.assertIn("not found", response.payload.get("message", "").lower())
|
|
|
|
async def test_worker_log_returns_content_json(self):
|
|
manager = _DummyWorkerManager()
|
|
with tempfile.NamedTemporaryFile("w", delete=False, encoding="utf-8") as handle:
|
|
handle.write("line-1\nline-2\nline-3\n")
|
|
log_path = handle.name
|
|
|
|
manager.processes["worker-a"] = {
|
|
"pid": 9999,
|
|
"log_file": log_path,
|
|
"process": None,
|
|
}
|
|
|
|
request = _FakeRequest(match_info={"worker_id": "worker-a"}, query={"lines": "2"})
|
|
try:
|
|
with patch.object(worker_routes, "get_worker_manager", return_value=manager), patch.object(
|
|
worker_routes.asyncio, "get_running_loop", return_value=_ImmediateLoop()
|
|
):
|
|
response = await worker_routes.get_worker_log_endpoint(request)
|
|
finally:
|
|
if os.path.exists(log_path):
|
|
os.remove(log_path)
|
|
|
|
self.assertEqual(response.status, 200)
|
|
self.assertEqual(response.payload.get("status"), "success")
|
|
self.assertIn("content", response.payload)
|
|
self.assertIn("line-3", response.payload["content"])
|
|
|
|
async def test_local_log_reads_memory_buffer(self):
|
|
request = _FakeRequest(query={"lines": "2"})
|
|
fake_logs = deque(
|
|
[
|
|
{"m": "line-1\n"},
|
|
{"m": "line-2\n"},
|
|
{"m": "line-3\n"},
|
|
],
|
|
maxlen=300,
|
|
)
|
|
app_module = types.ModuleType("app")
|
|
app_module.__path__ = []
|
|
logger_module = types.ModuleType("app.logger")
|
|
logger_module.get_logs = lambda: fake_logs
|
|
app_module.logger = logger_module
|
|
|
|
with patch.dict(sys.modules, {"app": app_module, "app.logger": logger_module}):
|
|
response = await worker_routes.get_local_log_endpoint(request)
|
|
|
|
self.assertEqual(response.status, 200)
|
|
self.assertEqual(response.payload.get("status"), "success")
|
|
self.assertEqual(response.payload.get("source"), "memory")
|
|
self.assertEqual(response.payload.get("entries"), 2)
|
|
self.assertIn("line-3", response.payload.get("content", ""))
|
|
|
|
async def test_remote_worker_log_proxies_to_worker_local_log_endpoint(self):
|
|
config = {
|
|
"workers": [
|
|
{
|
|
"id": "worker-remote",
|
|
"name": "Remote Worker",
|
|
"host": "worker.example.com",
|
|
"port": 8188,
|
|
"type": "remote",
|
|
}
|
|
]
|
|
}
|
|
request = _FakeRequest(match_info={"worker_id": "worker-remote"}, query={"lines": "120"})
|
|
proxied_payload = {
|
|
"status": "success",
|
|
"content": "remote-log-content\n",
|
|
"entries": 1,
|
|
"source": "memory",
|
|
"truncated": False,
|
|
"lines_shown": 1,
|
|
}
|
|
fake_session = _FakeHTTPClientSession(proxied_payload)
|
|
|
|
async def _fake_get_client_session():
|
|
return fake_session
|
|
|
|
with patch.object(worker_routes, "load_config", return_value=config), patch.object(
|
|
worker_routes, "get_client_session", side_effect=_fake_get_client_session
|
|
):
|
|
response = await worker_routes.get_remote_worker_log_endpoint(request)
|
|
|
|
self.assertEqual(response.status, 200)
|
|
self.assertEqual(response.payload.get("content"), "remote-log-content\n")
|
|
self.assertEqual(len(fake_session.calls), 1)
|
|
self.assertEqual(fake_session.calls[0]["params"], {"lines": "120"})
|
|
self.assertTrue(fake_session.calls[0]["url"].endswith("/distributed/local_log"))
|
|
|
|
async def test_remote_worker_log_rejects_local_workers(self):
|
|
config = {"workers": [{"id": "worker-local", "name": "Local Worker", "port": 8188}]}
|
|
request = _FakeRequest(match_info={"worker_id": "worker-local"})
|
|
|
|
with patch.object(worker_routes, "load_config", return_value=config):
|
|
response = await worker_routes.get_remote_worker_log_endpoint(request)
|
|
|
|
self.assertEqual(response.status, 400)
|
|
self.assertIn("local", response.payload.get("message", "").lower())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|