290 lines
9.1 KiB
Python
290 lines
9.1 KiB
Python
"""Tests for the direct Vast worker executor boundary."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import io
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from typing import Any
|
|
|
|
import pytest
|
|
|
|
|
|
class _CompletedRelayProcess:
|
|
"""Expose in-memory streams for one completed Vast relay command."""
|
|
|
|
def __init__(self, stdout: bytes, stderr: bytes, returncode: int) -> None:
|
|
"""Initialize deterministic subprocess-compatible state."""
|
|
self.stdin = io.BytesIO()
|
|
self.stdout = io.BytesIO(stdout)
|
|
self.stderr = io.BytesIO(stderr)
|
|
self.returncode = returncode
|
|
|
|
def poll(self) -> int:
|
|
"""Return the completed relay status."""
|
|
return self.returncode
|
|
|
|
def wait(self, timeout: float) -> int:
|
|
"""Return the completed relay status without blocking."""
|
|
del timeout
|
|
return self.returncode
|
|
|
|
def terminate(self) -> None:
|
|
"""Reject unexpected termination of an already-completed relay."""
|
|
raise AssertionError("completed relay should not be terminated")
|
|
|
|
def kill(self) -> None:
|
|
"""Reject unexpected killing of an already-completed relay."""
|
|
raise AssertionError("completed relay should not be killed")
|
|
|
|
|
|
def test_payload_instance_identity_accepts_concrete_environment(
|
|
vast_executor_module: Any,
|
|
) -> None:
|
|
"""Queue-time environment identities should route to exact leases."""
|
|
assert (
|
|
vast_executor_module._payload_instance_id(
|
|
{"execution_environment_id": "vast:node-17:9001"}
|
|
)
|
|
== 9001
|
|
)
|
|
assert vast_executor_module._payload_instance_id({"vast_instance_id": 9}) == 9
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"payload",
|
|
[
|
|
{},
|
|
{"execution_environment_id": "vast:profile"},
|
|
{"vast_instance_id": 0},
|
|
],
|
|
)
|
|
def test_payload_instance_identity_rejects_unresolved_profile(
|
|
vast_executor_module: Any,
|
|
payload: dict[str, Any],
|
|
) -> None:
|
|
"""Execution cannot begin before a profile resolves to a concrete lease."""
|
|
with pytest.raises(
|
|
vast_executor_module.VastRemoteInvocationError,
|
|
match="instance identity",
|
|
):
|
|
vast_executor_module._payload_instance_id(payload)
|
|
|
|
|
|
def test_payload_retention_defaults_to_twenty_four_hours(
|
|
vast_executor_module: Any,
|
|
) -> None:
|
|
"""Legacy or defensive payloads should retain the planned long cooldown."""
|
|
assert vast_executor_module._payload_retention_seconds({}) == 24 * 3600
|
|
assert (
|
|
vast_executor_module._payload_retention_seconds(
|
|
{"vast_idle_retention_seconds": 60}
|
|
)
|
|
== 60
|
|
)
|
|
|
|
|
|
def test_runtime_rejects_unregistered_or_draining_lease(
|
|
tmp_path: Path,
|
|
vast_executor_module: Any,
|
|
vast_leases_module: Any,
|
|
vast_runtime_module: Any,
|
|
) -> None:
|
|
"""The executor must never infer or reconnect to an unowned instance."""
|
|
|
|
class Activity:
|
|
"""Unused activity manager."""
|
|
|
|
client = vast_executor_module.VastExecutorClient(
|
|
registry=vast_leases_module.VastLeaseRegistry.for_user_directory(tmp_path),
|
|
activity_manager=Activity(),
|
|
runtime_configuration=vast_runtime_module.VastRuntimeConfiguration(
|
|
image="worker",
|
|
runtime_fingerprint="a" * 64,
|
|
),
|
|
user_directory=tmp_path,
|
|
)
|
|
|
|
with pytest.raises(
|
|
vast_executor_module.VastRemoteInvocationError,
|
|
match="not present in the local registry",
|
|
):
|
|
client._runtime({"vast_instance_id": 9001})
|
|
|
|
|
|
def test_vast_cancel_targets_owned_stager_before_worker_request(
|
|
vast_executor_module: Any,
|
|
monkeypatch: Any,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
"""A staging cancellation must terminate its transport and exact remote owner."""
|
|
client = vast_executor_module.VastExecutorClient(
|
|
registry=SimpleNamespace(),
|
|
activity_manager=SimpleNamespace(),
|
|
runtime_configuration=SimpleNamespace(),
|
|
user_directory=tmp_path,
|
|
)
|
|
invocation_id = "RIV_staging"
|
|
runner = SimpleNamespace()
|
|
process = SimpleNamespace()
|
|
terminated: list[Any] = []
|
|
targeted: list[tuple[Any, str]] = []
|
|
monkeypatch.setattr(
|
|
vast_executor_module,
|
|
"terminate_staging_transport",
|
|
terminated.append,
|
|
)
|
|
client._cancel_remote_stager = ( # type: ignore[method-assign]
|
|
lambda active_runner, owner_id, **_kwargs: (
|
|
targeted.append((active_runner, owner_id)) or True
|
|
)
|
|
)
|
|
with vast_executor_module._ACTIVE_VAST_STAGERS_LOCK:
|
|
vast_executor_module._ACTIVE_VAST_STAGERS[invocation_id] = (
|
|
runner,
|
|
"vast:RIV_staging:owner",
|
|
process,
|
|
)
|
|
try:
|
|
cancelled = client.cancel({}, invocation_id)
|
|
finally:
|
|
with vast_executor_module._ACTIVE_VAST_STAGERS_LOCK:
|
|
vast_executor_module._ACTIVE_VAST_STAGERS.pop(invocation_id, None)
|
|
|
|
assert cancelled is True
|
|
assert terminated == [process]
|
|
assert targeted == [(runner, "vast:RIV_staging:owner")]
|
|
|
|
|
|
def test_invoke_stream_surfaces_structured_worker_oom(
|
|
vast_executor_module: Any,
|
|
remote_protocol_module: Any,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
"""The controller should show cgroup OOM evidence as a terminal resource error."""
|
|
error_frame = remote_protocol_module.encode_json_frame(
|
|
remote_protocol_module.RemoteFrameKind.ERROR,
|
|
{
|
|
"error_type": "WorkerOutOfMemoryError",
|
|
"message": (
|
|
"Vast worker was killed because the instance ran out of host RAM "
|
|
"(container cgroup OOM; instance RAM limit 88.1 GiB)."
|
|
),
|
|
},
|
|
)
|
|
process = _CompletedRelayProcess(
|
|
error_frame,
|
|
b"Welcome to vast.ai.\nHave fun!.\n",
|
|
70,
|
|
)
|
|
runner = SimpleNamespace(popen=lambda _arguments: process)
|
|
client = vast_executor_module.VastExecutorClient(
|
|
registry=SimpleNamespace(),
|
|
activity_manager=SimpleNamespace(),
|
|
runtime_configuration=SimpleNamespace(),
|
|
user_directory=tmp_path,
|
|
)
|
|
|
|
with pytest.raises(
|
|
vast_executor_module.VastRemoteResourceError,
|
|
match="ran out of host RAM.*88.1 GiB",
|
|
):
|
|
list(
|
|
client._invoke_stream(
|
|
runner,
|
|
{"invocation_id": "RIV_oom"},
|
|
b"inputs",
|
|
)
|
|
)
|
|
|
|
|
|
def test_invoke_stream_discards_vast_login_greeting_without_postmortem(
|
|
vast_executor_module: Any,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
"""Vast's generic SSH greeting must never be presented as the crash cause."""
|
|
process = _CompletedRelayProcess(
|
|
b"",
|
|
(
|
|
b"Welcome to vast.ai. If authentication fails, try again after a few "
|
|
b"seconds, and double check your ssh key.\nHave fun!.\n"
|
|
),
|
|
0,
|
|
)
|
|
runner = SimpleNamespace(popen=lambda _arguments: process)
|
|
client = vast_executor_module.VastExecutorClient(
|
|
registry=SimpleNamespace(),
|
|
activity_manager=SimpleNamespace(),
|
|
runtime_configuration=SimpleNamespace(),
|
|
user_directory=tmp_path,
|
|
)
|
|
|
|
with pytest.raises(
|
|
vast_executor_module.VastRemoteTransportError,
|
|
match="No structured worker postmortem was available",
|
|
) as raised:
|
|
list(
|
|
client._invoke_stream(
|
|
runner,
|
|
{"invocation_id": "RIV_legacy"},
|
|
b"inputs",
|
|
)
|
|
)
|
|
|
|
assert "Welcome to vast.ai" not in str(raised.value)
|
|
assert "authentication fails" not in str(raised.value)
|
|
assert "SSH diagnostic" not in str(raised.value)
|
|
|
|
|
|
def test_consume_stream_does_not_retry_evidenced_oom(
|
|
vast_executor_module: Any,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
"""Restarting the same worker cannot recover from an undersized RAM limit."""
|
|
restart_count = 0
|
|
|
|
class Runtime:
|
|
"""Record worker lifecycle calls made by the executor."""
|
|
|
|
def ensure_worker(self) -> None:
|
|
"""Represent an already-ready worker."""
|
|
|
|
def restart_worker(self) -> None:
|
|
"""Record an unexpected resource-error recovery attempt."""
|
|
nonlocal restart_count
|
|
restart_count += 1
|
|
|
|
client = vast_executor_module.VastExecutorClient(
|
|
registry=SimpleNamespace(),
|
|
activity_manager=SimpleNamespace(),
|
|
runtime_configuration=SimpleNamespace(),
|
|
user_directory=tmp_path,
|
|
)
|
|
runtime = Runtime()
|
|
client._runtime = lambda _payload: ( # type: ignore[method-assign]
|
|
SimpleNamespace(instance_id=42),
|
|
SimpleNamespace(),
|
|
runtime,
|
|
)
|
|
client._ensure_llm_profiles_staged = ( # type: ignore[method-assign]
|
|
lambda _lease, _runner, _payload: None
|
|
)
|
|
|
|
def invoke(*_arguments: Any) -> Any:
|
|
"""Raise the terminal error when the stream is first consumed."""
|
|
raise vast_executor_module.VastRemoteResourceError("container cgroup OOM")
|
|
yield
|
|
|
|
client._invoke_stream = invoke # type: ignore[method-assign]
|
|
|
|
with pytest.raises(
|
|
vast_executor_module.VastRemoteResourceError,
|
|
match="container cgroup OOM",
|
|
):
|
|
client._consume_stream(
|
|
{"invocation_id": "RIV_oom", "vast_instance_id": 42},
|
|
b"inputs",
|
|
)
|
|
|
|
assert restart_count == 0
|