merge(runtime): integrate process and task state fixes
This commit is contained in:
commit
e9b821c0ec
14 files changed
+918
-197
No files matched your search
@@ -17,29 +17,7 @@ from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
from artemis.config import PAUSE_FILE
|
||||
|
||||
|
||||
def _pid_exists(pid: int) -> bool:
|
||||
"""Best-effort liveness probe; returns True when liveness cannot be determined.
|
||||
|
||||
psutil is preferred because on Windows ``os.kill(pid, 0)`` does not report
|
||||
dead pids reliably.
|
||||
"""
|
||||
try:
|
||||
import psutil
|
||||
|
||||
return psutil.pid_exists(pid)
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
import os
|
||||
|
||||
os.kill(pid, 0)
|
||||
return True
|
||||
except PermissionError:
|
||||
return True
|
||||
except (ProcessLookupError, OSError):
|
||||
return False
|
||||
from artemis.runtime.process_probe import pid_is_alive
|
||||
|
||||
|
||||
class ServerState:
|
||||
@@ -59,7 +37,12 @@ class ServerState:
|
||||
self.current_profile: str | None = None
|
||||
self.active_connections: dict[str, dict[str, Any]] = {}
|
||||
self.active_session_id: str | None = None
|
||||
self.was_stopped_manually: bool = False
|
||||
# Run keys (session id, or the synthetic run key of session-less runs)
|
||||
# that were stopped manually. Tracked per run so a manual stop of one
|
||||
# device's task never pollutes the terminal-status resolution of a
|
||||
# concurrently running task on another device. Entries are discarded by
|
||||
# each run's finalizer.
|
||||
self.manually_stopped_run_ids: set[str] = set()
|
||||
self.cancelled_session_ids: set[str] = set()
|
||||
self.startup_progress: dict[str, list[dict[str, Any]]] = {}
|
||||
|
||||
@@ -143,7 +126,7 @@ class ServerState:
|
||||
self.active_runs.pop(sid, None)
|
||||
continue
|
||||
pid = getattr(proc, "pid", None)
|
||||
if pid and not _pid_exists(pid):
|
||||
if pid and not pid_is_alive(pid):
|
||||
self.active_runs.pop(sid, None)
|
||||
|
||||
@property
|
||||
@@ -168,15 +151,8 @@ class ServerState:
|
||||
self.current_process = None
|
||||
else:
|
||||
pid = getattr(self.current_process, "pid", None)
|
||||
if pid:
|
||||
try:
|
||||
import os
|
||||
|
||||
os.kill(pid, 0)
|
||||
has_proc = True
|
||||
except (ProcessLookupError, PermissionError, OSError):
|
||||
has_proc = False
|
||||
self.current_process = None
|
||||
if pid and pid_is_alive(pid):
|
||||
has_proc = True
|
||||
else:
|
||||
has_proc = False
|
||||
self.current_process = None
|
||||
@@ -189,12 +165,9 @@ class ServerState:
|
||||
for sid, conn in list(self.active_connections.items()):
|
||||
c_pid = conn.get("pid")
|
||||
if c_pid:
|
||||
try:
|
||||
import os
|
||||
|
||||
os.kill(c_pid, 0)
|
||||
if pid_is_alive(c_pid):
|
||||
has_live_connection = True
|
||||
except (ProcessLookupError, PermissionError, OSError):
|
||||
else:
|
||||
self.active_connections.pop(sid, None)
|
||||
|
||||
if not has_proc and not has_running_item and not has_live_connection:
|
||||
|
||||
@@ -41,14 +41,19 @@ class SessionRepository:
|
||||
|
||||
@staticmethod
|
||||
def process_is_alive(pid: Any) -> bool:
|
||||
"""Return whether a session's worker PID is still alive."""
|
||||
try:
|
||||
import psutil
|
||||
"""Return whether a session's worker PID is still alive.
|
||||
|
||||
process = psutil.Process(int(pid))
|
||||
return process.is_running() and process.status() != psutil.STATUS_ZOMBIE
|
||||
except Exception:
|
||||
Delegates to the repository-wide probe: indeterminate liveness
|
||||
(e.g. AccessDenied) counts as alive so a running session is never
|
||||
reaped by mistake.
|
||||
"""
|
||||
try:
|
||||
pid_int = int(pid)
|
||||
except (TypeError, ValueError):
|
||||
return False
|
||||
from artemis.runtime.process_probe import pid_is_alive
|
||||
|
||||
return pid_is_alive(pid_int)
|
||||
|
||||
def get_all_sessions(self) -> list[dict[str, Any]]:
|
||||
with db_session(self.db_path) as conn:
|
||||
|
||||
@@ -415,10 +415,12 @@ class TaskQueueService:
|
||||
task_item["status"] = "running"
|
||||
task_item["start_time"] = time.time()
|
||||
|
||||
# A fresh launch clears a stale stop request from a previous task,
|
||||
# matching the historical serial worker's per-task reset. Per-session
|
||||
# stops are tracked in cancelled_session_ids and are unaffected.
|
||||
state.was_stopped_manually = False
|
||||
# A fresh launch clears a stale stop request left over for this run
|
||||
# key from a previous task. Manual stops are tracked per run in
|
||||
# manually_stopped_run_ids so stopping one device's task never
|
||||
# affects concurrent runs. Per-session stops are tracked in
|
||||
# cancelled_session_ids and are unaffected.
|
||||
state.manually_stopped_run_ids.discard(run_key)
|
||||
state.current_goal = goal
|
||||
state.current_profile = profile
|
||||
state.active_session_id = sess_id
|
||||
@@ -546,7 +548,7 @@ class TaskQueueService:
|
||||
|
||||
if sess_id and (
|
||||
str(sess_id) in getattr(state, "cancelled_session_ids", set())
|
||||
or state.was_stopped_manually
|
||||
or run_key in state.manually_stopped_run_ids
|
||||
):
|
||||
print(
|
||||
f"[QueueWorker] Task [{sess_id}] was cancelled during launch. Terminating."
|
||||
@@ -557,7 +559,7 @@ class TaskQueueService:
|
||||
returncode = await cls._wait_for_worker_process(proc)
|
||||
print(f"[QueueWorker] Task [{sess_id}] exited with returncode {returncode}")
|
||||
|
||||
manual_stop = state.was_stopped_manually or bool(
|
||||
manual_stop = run_key in state.manually_stopped_run_ids or bool(
|
||||
sess_id and str(sess_id) in state.cancelled_session_ids
|
||||
)
|
||||
|
||||
@@ -663,6 +665,7 @@ class TaskQueueService:
|
||||
cls._remove_task(sess_id)
|
||||
state.cancelled_session_ids.discard(str(sess_id))
|
||||
state.cancelled_session_ids.discard(run_key)
|
||||
state.manually_stopped_run_ids.discard(run_key)
|
||||
state.active_runs.pop(run_key, None)
|
||||
if proc is not None and state.current_process is proc:
|
||||
state.current_process = None
|
||||
@@ -670,7 +673,9 @@ class TaskQueueService:
|
||||
state.active_session_id = None
|
||||
state.current_goal = None
|
||||
state.current_profile = None
|
||||
state.was_stopped_manually = False
|
||||
# No runs left: any not-yet-consumed manual-stop markers are
|
||||
# stale and must not leak into future runs.
|
||||
state.manually_stopped_run_ids.clear()
|
||||
state.wake_event.set()
|
||||
|
||||
@classmethod
|
||||
@@ -892,11 +897,13 @@ class TaskQueueService:
|
||||
},
|
||||
)
|
||||
|
||||
if state.active_runs or state.current_process:
|
||||
state.was_stopped_manually = True
|
||||
for run_key, run in list(state.active_runs.items()):
|
||||
# Track the stop per run so each finalizer resolves its terminal
|
||||
# status as cancelled even if the global flag is reset meanwhile.
|
||||
# Mark each run individually: the global "stopped manually"
|
||||
# boolean was shared process-wide and polluted the terminal
|
||||
# status of unrelated concurrent runs.
|
||||
state.manually_stopped_run_ids.add(str(run_key))
|
||||
# Also track it per session so each finalizer resolves its
|
||||
# terminal status as cancelled.
|
||||
state.cancelled_session_ids.add(str(run_key))
|
||||
run_proc = run.get("process")
|
||||
if run_proc is not None and run_proc.returncode is None:
|
||||
@@ -969,20 +976,35 @@ class TaskQueueService:
|
||||
# Resolve the locally-managed run for this target (concurrent scheduling
|
||||
# keeps one subprocess per run in state.active_runs).
|
||||
local_run = None
|
||||
local_run_key: str | None = None
|
||||
if target_sid:
|
||||
local_run = state.active_runs.get(target_sid)
|
||||
if local_run is not None:
|
||||
local_run_key = target_sid
|
||||
elif target_device:
|
||||
local_run = next(
|
||||
local_run_key, local_run = next(
|
||||
(
|
||||
run
|
||||
for run in state.active_runs.values()
|
||||
(key, run)
|
||||
for key, run in state.active_runs.items()
|
||||
if run.get("device_id") and str(run["device_id"]) == target_device
|
||||
),
|
||||
(None, None),
|
||||
)
|
||||
elif len(state.active_runs) == 1:
|
||||
local_run_key, local_run = next(iter(state.active_runs.items()))
|
||||
local_proc = (local_run or {}).get("process") or state.current_process
|
||||
if local_run is None and local_proc is not None:
|
||||
# Legacy path: current_process without a resolved run entry. Find
|
||||
# the run that owns this process so a manual stop can be attributed
|
||||
# to it instead of to the whole scheduler.
|
||||
local_run_key = next(
|
||||
(
|
||||
str(key)
|
||||
for key, run in state.active_runs.items()
|
||||
if run.get("process") is local_proc
|
||||
),
|
||||
None,
|
||||
)
|
||||
elif len(state.active_runs) == 1:
|
||||
local_run = next(iter(state.active_runs.values()))
|
||||
local_proc = (local_run or {}).get("process") or state.current_process
|
||||
local_pid = getattr(local_proc, "pid", None)
|
||||
is_local_owner = bool(owner_pid and local_pid and owner_pid == local_pid)
|
||||
|
||||
@@ -1006,11 +1028,16 @@ class TaskQueueService:
|
||||
),
|
||||
None,
|
||||
)
|
||||
# active_session_id mirrors the most recently launched run, so with
|
||||
# concurrent runs the resolved local run key (== its session id) must
|
||||
# take precedence to avoid attributing the stop to an unrelated run.
|
||||
stopped_session_id = (
|
||||
owner.session_id
|
||||
if owner and owner.session_id
|
||||
else target_sid
|
||||
if target_sid
|
||||
else local_run_key
|
||||
if local_run_key is not None
|
||||
else state.active_session_id
|
||||
if is_local_owner
|
||||
else local_item.get("session_id")
|
||||
@@ -1038,8 +1065,8 @@ class TaskQueueService:
|
||||
# initialization window before Agent acquires the device lease.
|
||||
if target_sid:
|
||||
state.cancelled_session_ids.add(target_sid)
|
||||
else:
|
||||
state.was_stopped_manually = True
|
||||
elif local_run_key is not None:
|
||||
state.manually_stopped_run_ids.add(str(local_run_key))
|
||||
if local_pid:
|
||||
try:
|
||||
stopped = process_supervisor.terminate_tree(local_pid)
|
||||
@@ -1092,8 +1119,8 @@ class TaskQueueService:
|
||||
# Per-session cancellation: the run's own finalizer resolves the
|
||||
# terminal status without affecting other concurrent runs.
|
||||
state.cancelled_session_ids.add(str(stopped_session_id))
|
||||
else:
|
||||
state.was_stopped_manually = True
|
||||
elif local_run_key is not None:
|
||||
state.manually_stopped_run_ids.add(str(local_run_key))
|
||||
if local_proc is not None and state.current_process is local_proc:
|
||||
state.current_process = None
|
||||
|
||||
|
||||
@@ -28,6 +28,7 @@ from artemis.runtime.adb_endpoint import (
|
||||
current_adb_endpoint,
|
||||
)
|
||||
from artemis.runtime.device_pool import DevicePool, DeviceStatus, device_pool
|
||||
from artemis.runtime.process_probe import pid_is_alive
|
||||
from artemis.runtime.daemon_client import (
|
||||
ensure_daemon_running,
|
||||
get_daemon_session,
|
||||
@@ -78,6 +79,7 @@ __all__ = [
|
||||
"get_server_status",
|
||||
"is_daemon_running",
|
||||
"is_port_in_use",
|
||||
"pid_is_alive",
|
||||
"ProcessSupervisor",
|
||||
"process_supervisor",
|
||||
"read_server_info",
|
||||
|
||||
@@ -28,6 +28,7 @@ import time
|
||||
import uuid
|
||||
|
||||
from artemis.config.paths import get_temp_dir
|
||||
from artemis.runtime.process_probe import pid_is_alive
|
||||
from artemis.utils.logger import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
@@ -77,16 +78,11 @@ class ScreenAwakeLease:
|
||||
@staticmethod
|
||||
def _owner_is_alive(payload: dict) -> bool:
|
||||
try:
|
||||
import psutil
|
||||
|
||||
pid = int(payload["pid"])
|
||||
process = psutil.Process(pid)
|
||||
if not process.is_running():
|
||||
return False
|
||||
created_at = float(payload.get("process_created_at", 0.0))
|
||||
return created_at <= 0 or abs(process.create_time() - created_at) < 1.0
|
||||
except Exception:
|
||||
created_at = float(payload.get("process_created_at", 0.0) or 0.0)
|
||||
except (KeyError, TypeError, ValueError):
|
||||
return False
|
||||
return pid_is_alive(pid, created_at if created_at > 0 else None)
|
||||
|
||||
@staticmethod
|
||||
def _read_payload(path: Path) -> dict | None:
|
||||
|
||||
+131
-44
@@ -27,6 +27,7 @@ from typing import Any
|
||||
import uuid
|
||||
|
||||
from artemis.config.paths import get_temp_dir
|
||||
from artemis.runtime.process_probe import pid_is_alive
|
||||
|
||||
|
||||
class DeviceBusyError(RuntimeError):
|
||||
@@ -202,22 +203,22 @@ class DeviceExecutionLock:
|
||||
def _owner_is_alive(owner: DeviceLockOwner) -> bool:
|
||||
if owner.pid <= 0:
|
||||
return False
|
||||
try:
|
||||
import psutil
|
||||
if owner.process_created_at <= 0:
|
||||
# Legacy owner payloads without a create time cannot use PID-reuse
|
||||
# protection; require an Artemis-looking process before trusting
|
||||
# the record (a recycled PID must not keep a lock alive forever).
|
||||
try:
|
||||
import psutil
|
||||
|
||||
process = psutil.Process(owner.pid)
|
||||
if not process.is_running():
|
||||
return False
|
||||
if owner.process_created_at <= 0:
|
||||
try:
|
||||
name = process.name().lower()
|
||||
cmdline = " ".join(process.cmdline()).lower()
|
||||
return any(k in name or k in cmdline for k in ("python", "artemis", "pytest"))
|
||||
except Exception:
|
||||
process = psutil.Process(owner.pid)
|
||||
if not process.is_running():
|
||||
return False
|
||||
return abs(process.create_time() - owner.process_created_at) < 1.0
|
||||
except Exception:
|
||||
return False
|
||||
name = process.name().lower()
|
||||
cmdline = " ".join(process.cmdline()).lower()
|
||||
return any(k in name or k in cmdline for k in ("python", "artemis", "pytest"))
|
||||
except Exception:
|
||||
return False
|
||||
return pid_is_alive(owner.pid, owner.process_created_at)
|
||||
|
||||
@staticmethod
|
||||
def _read_owner(path: Path) -> DeviceLockOwner | None:
|
||||
@@ -268,6 +269,108 @@ class DeviceExecutionLock:
|
||||
payload["lock_scope"] = str(lock_scope)
|
||||
return payload
|
||||
|
||||
@classmethod
|
||||
def _ticket_targets_device(
|
||||
cls,
|
||||
ticket_owner: DeviceLockOwner,
|
||||
*,
|
||||
target_lock_id: str,
|
||||
target_scope: str | None,
|
||||
include_pending: bool,
|
||||
) -> bool:
|
||||
"""Return whether a queue ticket belongs to the given device's queue.
|
||||
|
||||
A concrete ticket matches only its own device. A ticket whose device id
|
||||
is still ``pending``/``any`` is a submission placeholder that no worker
|
||||
has claimed yet; with ``include_pending=True`` it matches any device
|
||||
with a compatible lock scope (used by listings, and by
|
||||
``_build_device_queue`` which then narrows the match down to a single
|
||||
parked device), while ``include_pending=False`` excludes it.
|
||||
"""
|
||||
ticket_dev = cls._normalize_device_id(ticket_owner.device_id)
|
||||
ticket_lock_id = cls._normalize_lock_id(
|
||||
ticket_owner.device_id,
|
||||
ticket_owner.lock_scope,
|
||||
)
|
||||
if ticket_lock_id == target_lock_id:
|
||||
return True
|
||||
if not include_pending:
|
||||
return False
|
||||
same_scope = (
|
||||
not ticket_owner.lock_scope
|
||||
or not target_scope
|
||||
or ticket_owner.lock_scope == target_scope
|
||||
)
|
||||
return ticket_dev in ("pending", "any") and same_scope
|
||||
|
||||
@classmethod
|
||||
def _build_device_queue(
|
||||
cls,
|
||||
wait_files: list[Path],
|
||||
*,
|
||||
target_lock_id: str,
|
||||
target_scope: str | None,
|
||||
) -> list[Path]:
|
||||
"""Return the FIFO wait-ticket queue that gates ``target_lock_id``.
|
||||
|
||||
Concrete tickets belong to exactly their own device's queue. An
|
||||
unclaimed ``pending``/``any`` ticket is *parked* on at most one device:
|
||||
pending tickets, taken in FIFO order, occupy the idle scope-compatible
|
||||
candidate devices in sorted lock-id order. Historically such a ticket
|
||||
was counted into every device's queue, so a single unclaimed
|
||||
head-of-queue ticket blocked all devices at once. Parking preserves
|
||||
submission order on the one device the ticket is presumed to claim
|
||||
while leaving every other idle device schedulable.
|
||||
|
||||
The pending submission itself cannot starve: its ticket file keeps its
|
||||
original timestamp, so once a worker claims a concrete device the
|
||||
ticket sorts ahead of every younger ticket for that device.
|
||||
"""
|
||||
entries: list[tuple[Path, str | None]] = []
|
||||
known_lock_ids: set[str] = set()
|
||||
for wait_path in wait_files:
|
||||
owner = cls._read_owner(wait_path)
|
||||
if owner is None:
|
||||
continue
|
||||
if cls._normalize_device_id(owner.device_id) in ("pending", "any"):
|
||||
if cls._ticket_targets_device(
|
||||
owner,
|
||||
target_lock_id=target_lock_id,
|
||||
target_scope=target_scope,
|
||||
include_pending=True,
|
||||
):
|
||||
entries.append((wait_path, None))
|
||||
continue
|
||||
same_scope = (
|
||||
not owner.lock_scope
|
||||
or not target_scope
|
||||
or owner.lock_scope == target_scope
|
||||
)
|
||||
if not same_scope:
|
||||
continue
|
||||
lock_id = cls._normalize_lock_id(owner.device_id, owner.lock_scope)
|
||||
known_lock_ids.add(lock_id)
|
||||
entries.append((wait_path, lock_id))
|
||||
|
||||
lock_dir = get_temp_dir("device-locks")
|
||||
idle_candidates = sorted(
|
||||
lock_id
|
||||
for lock_id in known_lock_ids
|
||||
if not (lock_dir / f"artemis-device-{lock_id}.lock").exists()
|
||||
)
|
||||
parked: dict[Path, str] = {}
|
||||
next_candidate = 0
|
||||
for wait_path, lock_id in entries:
|
||||
if lock_id is None and next_candidate < len(idle_candidates):
|
||||
parked[wait_path] = idle_candidates[next_candidate]
|
||||
next_candidate += 1
|
||||
|
||||
return [
|
||||
wait_path
|
||||
for wait_path, lock_id in entries
|
||||
if lock_id == target_lock_id or parked.get(wait_path) == target_lock_id
|
||||
]
|
||||
|
||||
@classmethod
|
||||
def reserve(
|
||||
cls,
|
||||
@@ -547,27 +650,15 @@ class DeviceExecutionLock:
|
||||
)
|
||||
is_eligible = is_reentrant or bool(all_wait_files and all_wait_files[0] == self._queue_path)
|
||||
else:
|
||||
# In multi-device parallel mode, per-device FIFO applies
|
||||
target_dev = self.clean_device_id
|
||||
device_queue: list[Path] = []
|
||||
for wait_path in all_wait_files:
|
||||
ticket_owner = self._read_owner(wait_path)
|
||||
if ticket_owner is None:
|
||||
continue
|
||||
ticket_dev = self._normalize_device_id(ticket_owner.device_id)
|
||||
ticket_lock_id = self._normalize_lock_id(
|
||||
ticket_owner.device_id,
|
||||
ticket_owner.lock_scope,
|
||||
)
|
||||
same_scope = (
|
||||
not ticket_owner.lock_scope
|
||||
or not self.lock_scope
|
||||
or ticket_owner.lock_scope == self.lock_scope
|
||||
)
|
||||
if ticket_lock_id == target_dev or (
|
||||
ticket_dev in ("pending", "any") and same_scope
|
||||
):
|
||||
device_queue.append(wait_path)
|
||||
# In multi-device parallel mode, per-device FIFO applies.
|
||||
# Unclaimed pending/any tickets are parked on at most one
|
||||
# idle device instead of gating the head of every device's
|
||||
# queue (see _build_device_queue).
|
||||
device_queue = self._build_device_queue(
|
||||
all_wait_files,
|
||||
target_lock_id=self.clean_device_id,
|
||||
target_scope=self.lock_scope,
|
||||
)
|
||||
is_eligible = bool(device_queue and device_queue[0] == self._queue_path)
|
||||
|
||||
if is_eligible and self._try_acquire_owner_lock():
|
||||
@@ -729,15 +820,11 @@ class DeviceExecutionLock:
|
||||
if owner.token in active_tokens:
|
||||
continue
|
||||
|
||||
target_dev = cls._normalize_device_id(owner.device_id)
|
||||
target_lock_id = cls._normalize_lock_id(owner.device_id, owner.lock_scope)
|
||||
same_scope = (
|
||||
not target_scope
|
||||
or not owner.lock_scope
|
||||
or owner.lock_scope == target_scope
|
||||
)
|
||||
if clean_target is not None and target_lock_id != clean_target and not (
|
||||
target_dev in ("pending", "any") and same_scope
|
||||
if clean_target is not None and not cls._ticket_targets_device(
|
||||
owner,
|
||||
target_lock_id=clean_target,
|
||||
target_scope=target_scope,
|
||||
include_pending=True,
|
||||
):
|
||||
continue
|
||||
|
||||
|
||||
@@ -0,0 +1,98 @@
|
||||
# Copyright 2026 Google LLC
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Single source of truth for PID liveness probing across the Artemis codebase.
|
||||
|
||||
Semantics (decided by the maintainers):
|
||||
|
||||
- psutil only. No bare ``os.kill(pid, 0)`` fallback: on Windows, CPython
|
||||
implements ``os.kill(pid, 0)`` for non-CTRL signals via
|
||||
``OpenProcess``/``TerminateProcess`` handle probing, which is both unreliable
|
||||
for dead pids and semantically surprising.
|
||||
- ``psutil.NoSuchProcess`` -> dead. Zombie processes -> dead.
|
||||
- ``psutil.AccessDenied`` / ``PermissionError`` -> **alive**. When liveness
|
||||
cannot be determined the probe defaults to "alive": it is always safer to
|
||||
skip reaping a resource than to reap one that is still owned.
|
||||
- ``created_at`` (a ``psutil``-style process create time) enables PID-reuse
|
||||
protection: when provided and the probed process's create time differs by
|
||||
one second or more, the original process is dead and its PID was recycled.
|
||||
- Any other unexpected error -> alive (default-alive), with a warning log.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def pid_is_alive(pid: int, created_at: float | None = None) -> bool:
|
||||
"""Return whether the process identified by ``pid`` is still alive.
|
||||
|
||||
Args:
|
||||
pid: The process id to probe. Non-positive or unparsable pids are dead.
|
||||
created_at: Optional expected process create time (as reported by
|
||||
``psutil.Process.create_time()``). When given and > 0, a create-time
|
||||
mismatch of >= 1 second means the PID was recycled -> dead.
|
||||
|
||||
Returns:
|
||||
True when the process is alive or liveness cannot be determined
|
||||
(default-alive); False when the process is definitively gone, a
|
||||
zombie, or its PID has been recycled by another process.
|
||||
"""
|
||||
try:
|
||||
pid = int(pid)
|
||||
except (TypeError, ValueError):
|
||||
return False
|
||||
if pid <= 0:
|
||||
return False
|
||||
|
||||
try:
|
||||
import psutil
|
||||
except Exception as exc: # pragma: no cover - psutil is a hard dependency
|
||||
logger.warning(f"psutil unavailable while probing pid {pid}: {exc}; assuming alive")
|
||||
return True
|
||||
|
||||
try:
|
||||
process = psutil.Process(pid)
|
||||
if not process.is_running():
|
||||
return False
|
||||
|
||||
try:
|
||||
if process.status() == psutil.STATUS_ZOMBIE:
|
||||
return False
|
||||
except psutil.NoSuchProcess:
|
||||
return False
|
||||
except (psutil.AccessDenied, PermissionError):
|
||||
# Cannot inspect the status; keep the default-alive verdict.
|
||||
pass
|
||||
|
||||
if created_at is not None and created_at > 0:
|
||||
try:
|
||||
if abs(process.create_time() - float(created_at)) >= 1.0:
|
||||
# Same PID, different birth time: the PID was recycled.
|
||||
return False
|
||||
except psutil.NoSuchProcess:
|
||||
return False
|
||||
except (psutil.AccessDenied, PermissionError):
|
||||
pass
|
||||
|
||||
return True
|
||||
except psutil.NoSuchProcess:
|
||||
return False
|
||||
except (psutil.AccessDenied, PermissionError):
|
||||
return True
|
||||
except Exception as exc:
|
||||
logger.warning(f"Could not determine liveness of pid {pid}: {exc}; assuming alive")
|
||||
return True
|
||||
@@ -168,31 +168,9 @@ def _reconcile_orphaned_sessions() -> int:
|
||||
|
||||
def _any_pid_alive(pids: list[int]) -> bool:
|
||||
"""Return whether any target process is still running (zombies excluded)."""
|
||||
if not pids:
|
||||
return False
|
||||
try:
|
||||
import psutil
|
||||
from artemis.runtime.process_probe import pid_is_alive
|
||||
|
||||
for pid in pids:
|
||||
try:
|
||||
process = psutil.Process(pid)
|
||||
if process.is_running() and process.status() != psutil.STATUS_ZOMBIE:
|
||||
return True
|
||||
except psutil.NoSuchProcess:
|
||||
continue
|
||||
except psutil.AccessDenied:
|
||||
return True
|
||||
return False
|
||||
except Exception:
|
||||
return any(_pid_exists_portably(pid) for pid in pids)
|
||||
|
||||
|
||||
def _pid_exists_portably(pid: int) -> bool:
|
||||
try:
|
||||
os.kill(pid, 0)
|
||||
return True
|
||||
except OSError:
|
||||
return False
|
||||
return any(pid_is_alive(pid) for pid in pids)
|
||||
|
||||
|
||||
def find_server_pids(port: int = 8000) -> list[int]:
|
||||
@@ -210,17 +188,10 @@ def find_server_pids(port: int = 8000) -> list[int]:
|
||||
if info and info.get("port") == port:
|
||||
saved_pid = info.get("pid")
|
||||
if saved_pid and isinstance(saved_pid, int):
|
||||
try:
|
||||
import psutil
|
||||
from artemis.runtime.process_probe import pid_is_alive
|
||||
|
||||
if psutil.pid_exists(saved_pid):
|
||||
discovered.add(saved_pid)
|
||||
except Exception:
|
||||
try:
|
||||
os.kill(saved_pid, 0)
|
||||
discovered.add(saved_pid)
|
||||
except OSError:
|
||||
pass
|
||||
if pid_is_alive(saved_pid):
|
||||
discovered.add(saved_pid)
|
||||
|
||||
# 2. Check port listeners with platform-native tools
|
||||
# 2a. lsof (macOS & Linux)
|
||||
|
||||
@@ -45,16 +45,12 @@ def _pid_alive(pid: int | None) -> bool:
|
||||
if not pid:
|
||||
return False
|
||||
try:
|
||||
import psutil
|
||||
pid_int = int(pid)
|
||||
except (TypeError, ValueError):
|
||||
return False
|
||||
from artemis.runtime.process_probe import pid_is_alive
|
||||
|
||||
proc = psutil.Process(int(pid))
|
||||
return proc.is_running() and proc.status() != psutil.STATUS_ZOMBIE
|
||||
except Exception:
|
||||
try:
|
||||
os.kill(int(pid), 0)
|
||||
return True
|
||||
except OSError:
|
||||
return False
|
||||
return pid_is_alive(pid_int)
|
||||
|
||||
|
||||
def _session_tracked_by_lock(trace_id: str) -> bool:
|
||||
|
||||
+235
-33
@@ -12,12 +12,30 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Trace storage and lifecycle management for MCP background automation tasks."""
|
||||
"""Trace storage and lifecycle management for MCP background automation tasks.
|
||||
|
||||
status.json is written and updated concurrently by real independent processes
|
||||
(the admin console daemon, MCP server tools, and spawned worker subprocesses),
|
||||
so this module provides:
|
||||
|
||||
- Atomic writes (temp file + fsync + ``os.replace``) so readers never observe
|
||||
a torn/truncated file.
|
||||
- A PID-aware ``status.json.lock`` file (O_CREAT | O_EXCL with bounded retry
|
||||
and stale-lock detection) serializing read-modify-write updates.
|
||||
- Three-state reads: missing file (silent ``None``), corrupt JSON (WARNING +
|
||||
the file is quarantined as ``status.json.corrupt``), other IO errors
|
||||
(WARNING).
|
||||
"""
|
||||
|
||||
from contextlib import contextmanager
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from typing import Any
|
||||
import uuid
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
try:
|
||||
from artemis.config.paths import get_traces_dir
|
||||
@@ -27,6 +45,153 @@ except Exception:
|
||||
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
TRACES_DIR = os.path.join(PROJECT_ROOT, "traces")
|
||||
|
||||
_LOCK_TIMEOUT_SECONDS = 5.0
|
||||
_LOCK_POLL_SECONDS = 0.05
|
||||
# A lock younger than this is assumed to be held by a live writer and is not
|
||||
# probed for staleness. Probing opens the lock file, and on Windows an open
|
||||
# reader handle blocks the owner's unlink (no FILE_SHARE_DELETE), so eager
|
||||
# probing can leak the lock forever under contention.
|
||||
_LOCK_STALE_PROBE_MIN_AGE_SECONDS = 0.5
|
||||
_LOCK_STALE_PROBE_INTERVAL_SECONDS = 1.0
|
||||
_MALFORMED_LOCK_GRACE_SECONDS = 10.0
|
||||
_REPLACE_RETRIES = 5
|
||||
_REPLACE_RETRY_DELAY_SECONDS = 0.02
|
||||
_UNLINK_RETRIES = 10
|
||||
_UNLINK_RETRY_DELAY_SECONDS = 0.01
|
||||
|
||||
|
||||
def _own_create_time() -> float:
|
||||
try:
|
||||
import psutil
|
||||
|
||||
return float(psutil.Process(os.getpid()).create_time())
|
||||
except Exception:
|
||||
return 0.0
|
||||
|
||||
|
||||
def _replace_with_retry(src: str, dst: str) -> None:
|
||||
"""``os.replace`` with retries for transient Windows sharing violations."""
|
||||
last_error: OSError | None = None
|
||||
for attempt in range(_REPLACE_RETRIES):
|
||||
try:
|
||||
os.replace(src, dst)
|
||||
return
|
||||
except OSError as exc:
|
||||
last_error = exc
|
||||
if attempt < _REPLACE_RETRIES - 1:
|
||||
time.sleep(_REPLACE_RETRY_DELAY_SECONDS)
|
||||
raise last_error # type: ignore[misc]
|
||||
|
||||
|
||||
def _atomic_write_json(path: str, data: dict[str, Any]) -> None:
|
||||
"""Write JSON atomically: temp file in the same directory + fsync + replace."""
|
||||
os.makedirs(os.path.dirname(path), exist_ok=True)
|
||||
temp_path = f"{path}.{uuid.uuid4().hex}.tmp"
|
||||
try:
|
||||
with open(temp_path, "w", encoding="utf-8") as f:
|
||||
json.dump(data, f, indent=2, ensure_ascii=False)
|
||||
f.flush()
|
||||
os.fsync(f.fileno())
|
||||
_replace_with_retry(temp_path, path)
|
||||
finally:
|
||||
# A leftover temp file (e.g. after a crash between write and replace)
|
||||
# never affects readers; remove it best-effort on the success path.
|
||||
try:
|
||||
os.unlink(temp_path)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
def _unlink_with_retry(path: str) -> bool:
|
||||
"""Best-effort unlink retrying transient Windows sharing violations."""
|
||||
for attempt in range(_UNLINK_RETRIES):
|
||||
try:
|
||||
os.unlink(path)
|
||||
return True
|
||||
except FileNotFoundError:
|
||||
return True
|
||||
except OSError:
|
||||
if attempt < _UNLINK_RETRIES - 1:
|
||||
time.sleep(_UNLINK_RETRY_DELAY_SECONDS)
|
||||
return False
|
||||
|
||||
|
||||
def _lock_is_stale(lock_path: str) -> bool:
|
||||
"""A status lock is stale when its owner process is gone or it is malformed and old."""
|
||||
try:
|
||||
if time.time() - os.path.getmtime(lock_path) < _LOCK_STALE_PROBE_MIN_AGE_SECONDS:
|
||||
# Young lock: assume a live writer holds it. Do not open the file
|
||||
# (an open reader handle would block the owner's unlink on Windows).
|
||||
return False
|
||||
except OSError:
|
||||
return False
|
||||
try:
|
||||
with open(lock_path, encoding="utf-8") as f:
|
||||
payload = json.load(f)
|
||||
pid = int(payload["pid"])
|
||||
created_at = float(payload.get("process_created_at", 0.0) or 0.0)
|
||||
except FileNotFoundError:
|
||||
return False
|
||||
except Exception:
|
||||
try:
|
||||
return time.time() - os.path.getmtime(lock_path) >= _MALFORMED_LOCK_GRACE_SECONDS
|
||||
except OSError:
|
||||
return False
|
||||
from artemis.runtime.process_probe import pid_is_alive
|
||||
|
||||
return not pid_is_alive(pid, created_at if created_at > 0 else None)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _status_lock(path: str):
|
||||
"""Bounded cross-process mutex for read-modify-write updates of ``path``.
|
||||
|
||||
Yields True when the lock was acquired. On timeout the caller proceeds
|
||||
without mutual exclusion (with a WARNING) rather than deadlocking: a
|
||||
last-writer-wins update is still strictly better than blocking MCP tools
|
||||
forever on a leaked lock.
|
||||
"""
|
||||
lock_path = f"{path}.lock"
|
||||
os.makedirs(os.path.dirname(lock_path), exist_ok=True)
|
||||
payload = json.dumps({"pid": os.getpid(), "process_created_at": _own_create_time()})
|
||||
deadline = time.monotonic() + _LOCK_TIMEOUT_SECONDS
|
||||
acquired = False
|
||||
last_stale_probe = 0.0
|
||||
while not acquired:
|
||||
try:
|
||||
fd = os.open(lock_path, os.O_CREAT | os.O_EXCL | os.O_WRONLY)
|
||||
except FileExistsError:
|
||||
now = time.monotonic()
|
||||
if now - last_stale_probe >= _LOCK_STALE_PROBE_INTERVAL_SECONDS:
|
||||
last_stale_probe = now
|
||||
if _lock_is_stale(lock_path):
|
||||
_unlink_with_retry(lock_path)
|
||||
continue
|
||||
if now >= deadline:
|
||||
logger.warning(
|
||||
f"Timed out acquiring trace status lock {lock_path}; "
|
||||
"proceeding without mutual exclusion"
|
||||
)
|
||||
break
|
||||
time.sleep(_LOCK_POLL_SECONDS)
|
||||
except OSError as exc:
|
||||
logger.warning(f"Could not create trace status lock {lock_path}: {exc}")
|
||||
break
|
||||
else:
|
||||
try:
|
||||
os.write(fd, payload.encode("utf-8"))
|
||||
finally:
|
||||
os.close(fd)
|
||||
acquired = True
|
||||
try:
|
||||
yield acquired
|
||||
finally:
|
||||
if acquired and not _unlink_with_retry(lock_path):
|
||||
logger.warning(
|
||||
f"Could not remove trace status lock {lock_path}; it will be "
|
||||
"reaped as stale once its owner exits"
|
||||
)
|
||||
|
||||
|
||||
def get_trace_dir(trace_id: str) -> str:
|
||||
"""Returns the absolute path to the directory for a given trace_id."""
|
||||
@@ -82,23 +247,41 @@ def init_trace(
|
||||
|
||||
|
||||
def read_status(trace_id: str) -> dict[str, Any] | None:
|
||||
"""Reads and returns the status.json for a given trace_id, or None if it doesn't exist."""
|
||||
"""Read the status.json for a trace_id, distinguishing three failure states.
|
||||
|
||||
- Missing file: normal (a trace that was never started) -> silent ``None``.
|
||||
- Corrupt JSON: WARNING with the path; the damaged file is renamed to
|
||||
``status.json.corrupt`` as evidence, then ``None`` is returned.
|
||||
- Other IO errors: WARNING, ``None``.
|
||||
"""
|
||||
path = get_status_path(trace_id)
|
||||
if not os.path.exists(path):
|
||||
return None
|
||||
try:
|
||||
with open(path, encoding="utf-8") as f:
|
||||
return json.load(f)
|
||||
except Exception:
|
||||
except FileNotFoundError:
|
||||
return None
|
||||
except (json.JSONDecodeError, UnicodeDecodeError, ValueError) as exc:
|
||||
logger.warning(
|
||||
f"Corrupt status.json for trace {trace_id} at {path}: {exc}; "
|
||||
f"quarantining as {path}.corrupt"
|
||||
)
|
||||
try:
|
||||
_replace_with_retry(path, f"{path}.corrupt")
|
||||
except OSError as move_exc:
|
||||
logger.warning(
|
||||
f"Could not quarantine corrupt status file {path}: {move_exc}"
|
||||
)
|
||||
return None
|
||||
except OSError as exc:
|
||||
logger.warning(f"Could not read status.json for trace {trace_id} at {path}: {exc}")
|
||||
return None
|
||||
|
||||
|
||||
def write_status(trace_id: str, data: dict[str, Any]) -> None:
|
||||
"""Writes the given status dictionary to the trace's status.json."""
|
||||
path = get_status_path(trace_id)
|
||||
os.makedirs(os.path.dirname(path), exist_ok=True)
|
||||
with open(path, "w", encoding="utf-8") as f:
|
||||
json.dump(data, f, indent=2, ensure_ascii=False)
|
||||
"""Atomically writes the given status dictionary to the trace's status.json."""
|
||||
_atomic_write_json(get_status_path(trace_id), data)
|
||||
|
||||
|
||||
def update_trace_status(
|
||||
@@ -108,37 +291,56 @@ def update_trace_status(
|
||||
result: Any | None = None,
|
||||
device_serial: str | None = None,
|
||||
) -> dict[str, Any] | None:
|
||||
"""Updates specific fields of the status.json for a given trace_id."""
|
||||
data = read_status(trace_id)
|
||||
if not data:
|
||||
return None
|
||||
"""Updates specific fields of the status.json for a given trace_id.
|
||||
|
||||
# "success" is a legacy alias for the canonical "completed" terminal
|
||||
# status and must never be persisted into status.json.
|
||||
if status == "success":
|
||||
status = "completed"
|
||||
The read-modify-write cycle holds a cross-process lock so concurrent
|
||||
updates from the daemon, MCP tools, and worker processes do not lose each
|
||||
other's fields.
|
||||
"""
|
||||
path = get_status_path(trace_id)
|
||||
with _status_lock(path):
|
||||
data = read_status(trace_id)
|
||||
if not data:
|
||||
if os.path.exists(path) or os.path.exists(f"{path}.corrupt"):
|
||||
logger.warning(
|
||||
f"Dropping status update ({status!r}) for trace {trace_id}: "
|
||||
"existing status.json is unreadable"
|
||||
)
|
||||
return None
|
||||
|
||||
data["status"] = status
|
||||
if status in ("completed", "failed", "cancelled"):
|
||||
data["end_time"] = time.time()
|
||||
# "success" is a legacy alias for the canonical "completed" terminal
|
||||
# status and must never be persisted into status.json.
|
||||
if status == "success":
|
||||
status = "completed"
|
||||
|
||||
if error is not None:
|
||||
data["error"] = error
|
||||
if result is not None:
|
||||
data["result"] = result
|
||||
if device_serial is not None:
|
||||
data["device_serial"] = device_serial
|
||||
data["status"] = status
|
||||
if status in ("completed", "failed", "cancelled"):
|
||||
data["end_time"] = time.time()
|
||||
|
||||
write_status(trace_id, data)
|
||||
return data
|
||||
if error is not None:
|
||||
data["error"] = error
|
||||
if result is not None:
|
||||
data["result"] = result
|
||||
if device_serial is not None:
|
||||
data["device_serial"] = device_serial
|
||||
|
||||
write_status(trace_id, data)
|
||||
return data
|
||||
|
||||
|
||||
def update_trace_device_serial(trace_id: str, device_serial: str) -> dict[str, Any] | None:
|
||||
"""Updates the device_serial field of the status.json for a given trace_id."""
|
||||
data = read_status(trace_id)
|
||||
if not data:
|
||||
return None
|
||||
path = get_status_path(trace_id)
|
||||
with _status_lock(path):
|
||||
data = read_status(trace_id)
|
||||
if not data:
|
||||
if os.path.exists(path) or os.path.exists(f"{path}.corrupt"):
|
||||
logger.warning(
|
||||
f"Dropping device_serial update for trace {trace_id}: "
|
||||
"existing status.json is unreadable"
|
||||
)
|
||||
return None
|
||||
|
||||
data["device_serial"] = device_serial
|
||||
write_status(trace_id, data)
|
||||
return data
|
||||
data["device_serial"] = device_serial
|
||||
write_status(trace_id, data)
|
||||
return data
|
||||
@@ -52,7 +52,7 @@ def clean_state(tmp_path, monkeypatch):
|
||||
state.current_profile = None
|
||||
state.active_session_id = None
|
||||
state.active_connections.clear()
|
||||
state.was_stopped_manually = False
|
||||
state.manually_stopped_run_ids.clear()
|
||||
state.cancelled_session_ids.clear()
|
||||
if state.worker_task and not state.worker_task.done():
|
||||
state.worker_task.cancel()
|
||||
@@ -66,7 +66,7 @@ def clean_state(tmp_path, monkeypatch):
|
||||
state.current_process = None
|
||||
state.current_goal = None
|
||||
state.current_profile = None
|
||||
state.was_stopped_manually = False
|
||||
state.manually_stopped_run_ids.clear()
|
||||
if hasattr(state, "recent_submissions"):
|
||||
state.recent_submissions.clear()
|
||||
if state.worker_task and not state.worker_task.done():
|
||||
@@ -448,7 +448,7 @@ def test_stop_tasks_terminates_external_global_owner_and_preserves_local_waiter(
|
||||
local_waiter.kill.assert_not_called()
|
||||
assert state.current_process is local_waiter
|
||||
assert state.queue_items[0]["session_id"] == "frontend-waiter"
|
||||
assert state.was_stopped_manually is False
|
||||
assert not state.manually_stopped_run_ids
|
||||
assert "frontend-waiter" not in state.cancelled_session_ids
|
||||
assert "mcp-session" not in state.active_connections
|
||||
update_status.assert_called_once()
|
||||
@@ -921,6 +921,111 @@ async def test_enqueue_tasks_deduplicates_by_session_id():
|
||||
assert res2["tasks"][0]["session_id"] == "sid-dedup-1"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_manual_stop_of_one_run_does_not_pollute_concurrent_run():
|
||||
"""Stopping run A must not flip run B's terminal status or its payload.
|
||||
|
||||
Regression test for the process-global ``was_stopped_manually`` flag that
|
||||
used to mark *every* in-flight run as manually stopped. Two concurrent runs
|
||||
on different devices: device A's task is stopped manually, device B's task
|
||||
then finishes normally and must still be reported as completed.
|
||||
"""
|
||||
ended_payloads: dict[str, dict] = {}
|
||||
|
||||
def capture(event_type, data):
|
||||
if event_type == "session_ended":
|
||||
ended_payloads[str(data.get("session_id"))] = dict(data)
|
||||
|
||||
class FakeProc:
|
||||
def __init__(self, pid):
|
||||
self.pid = pid
|
||||
self.returncode = None
|
||||
self.stdout = None
|
||||
self._done = asyncio.Event()
|
||||
|
||||
async def wait(self):
|
||||
await self._done.wait()
|
||||
return self.returncode
|
||||
|
||||
def finish(self, returncode):
|
||||
self.returncode = returncode
|
||||
self._done.set()
|
||||
|
||||
def kill(self):
|
||||
self.finish(-9)
|
||||
|
||||
procs: dict[str, FakeProc] = {}
|
||||
|
||||
async def fake_subprocess_exec(*args, **kwargs):
|
||||
goal = args[3]
|
||||
# Use our own pid so the reaped-process watchdog keeps waiting.
|
||||
proc = FakeProc(os.getpid())
|
||||
procs[goal] = proc
|
||||
return proc
|
||||
|
||||
item_a = {
|
||||
"session_id": "run-a",
|
||||
"goal": "Goal A",
|
||||
"status": "pending",
|
||||
"device_serial": "dev-a",
|
||||
}
|
||||
item_b = {
|
||||
"session_id": "run-b",
|
||||
"goal": "Goal B",
|
||||
"status": "pending",
|
||||
"device_serial": "dev-b",
|
||||
}
|
||||
state.queue_items = [item_a, item_b]
|
||||
state.ipc_subscribers.append(capture)
|
||||
|
||||
try:
|
||||
with (
|
||||
patch("asyncio.create_subprocess_exec", side_effect=fake_subprocess_exec),
|
||||
patch("apps.admin_console.services.task_queue_service.session_repo") as mock_repo,
|
||||
patch("apps.admin_console.services.task_queue_service.media_service"),
|
||||
patch(
|
||||
"apps.admin_console.services.task_queue_service.process_supervisor.terminate_tree",
|
||||
return_value=True,
|
||||
),
|
||||
):
|
||||
mock_repo.get_session_status.return_value = "running"
|
||||
mock_repo.get_video_recording_for_session.return_value = {"status": "ready"}
|
||||
|
||||
task_a = asyncio.create_task(TaskQueueService._execute_task_item(item_a))
|
||||
task_b = asyncio.create_task(TaskQueueService._execute_task_item(item_b))
|
||||
for _ in range(40):
|
||||
if "run-a" in state.active_runs and "run-b" in state.active_runs:
|
||||
break
|
||||
await asyncio.sleep(0.02)
|
||||
assert {"run-a", "run-b"} <= set(state.active_runs)
|
||||
# B launched last, so the legacy single-task mirror points at B.
|
||||
assert state.active_session_id == "run-b"
|
||||
|
||||
# Manually stop device A's task without naming its session id.
|
||||
assert task_queue_service.stop_tasks(clear_all=False, device_id="dev-a") is True
|
||||
await asyncio.wait_for(task_a, timeout=5.0)
|
||||
|
||||
# B finishes normally afterwards.
|
||||
procs["Goal B"].finish(0)
|
||||
await asyncio.wait_for(task_b, timeout=5.0)
|
||||
|
||||
assert ended_payloads["run-a"]["status"] == "cancelled"
|
||||
assert ended_payloads["run-a"]["was_stopped_manually"] is True
|
||||
assert ended_payloads["run-b"]["status"] == "completed"
|
||||
assert ended_payloads["run-b"]["was_stopped_manually"] is False
|
||||
|
||||
persisted = {
|
||||
call.args[0]: call.args[1]
|
||||
for call in mock_repo.update_session_status.call_args_list
|
||||
}
|
||||
assert persisted.get("run-a") == "cancelled"
|
||||
assert persisted.get("run-b") == "completed"
|
||||
assert "run-b" not in state.cancelled_session_ids
|
||||
assert "run-b" not in state.manually_stopped_run_ids
|
||||
finally:
|
||||
state.ipc_subscribers.remove(capture)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_enqueue_tasks_debounces_rapid_identical_submissions():
|
||||
state.queue_items = []
|
||||
|
||||
@@ -14,9 +14,12 @@
|
||||
|
||||
"""Unit tests for MCP trace store."""
|
||||
|
||||
import glob
|
||||
import logging
|
||||
import os
|
||||
import shutil
|
||||
import tempfile
|
||||
import threading
|
||||
import uuid
|
||||
import pytest
|
||||
|
||||
@@ -91,6 +94,125 @@ def test_read_nonexistent_status(temp_trace_env):
|
||||
assert trace_store.read_status(non_existent) is None
|
||||
|
||||
|
||||
def test_leftover_atomic_temp_file_does_not_affect_reads(temp_trace_env):
|
||||
"""A crash between temp-write and replace leaves a .tmp file that readers ignore."""
|
||||
trace_id = str(uuid.uuid4())
|
||||
trace_store.init_trace(trace_id, "Crash sim", "Flash")
|
||||
|
||||
status_path = trace_store.get_status_path(trace_id)
|
||||
# Simulate a writer that crashed after writing its temp file.
|
||||
with open(f"{status_path}.{uuid.uuid4().hex}.tmp", "w", encoding="utf-8") as f:
|
||||
f.write('{"status": "half-writ')
|
||||
|
||||
data = trace_store.read_status(trace_id)
|
||||
assert data is not None
|
||||
assert data["status"] == "running"
|
||||
assert data["task_desc"] == "Crash sim"
|
||||
|
||||
|
||||
def test_failed_atomic_replace_preserves_original(temp_trace_env):
|
||||
"""If the atomic replace fails, the previous status.json stays intact."""
|
||||
from unittest.mock import patch
|
||||
|
||||
trace_id = str(uuid.uuid4())
|
||||
trace_store.init_trace(trace_id, "Original", "Flash")
|
||||
|
||||
def broken_replace(src, dst):
|
||||
raise OSError("simulated replace failure")
|
||||
|
||||
with patch.object(trace_store, "_replace_with_retry", broken_replace):
|
||||
with pytest.raises(OSError):
|
||||
trace_store.write_status(trace_id, {"status": "clobbered"})
|
||||
|
||||
data = trace_store.read_status(trace_id)
|
||||
assert data is not None
|
||||
assert data["task_desc"] == "Original"
|
||||
# The failed writer's temp file must not linger.
|
||||
assert glob.glob(os.path.join(trace_store.get_trace_dir(trace_id), "*.tmp")) == []
|
||||
|
||||
|
||||
def test_corrupt_status_is_quarantined_and_read_returns_none(temp_trace_env, caplog):
|
||||
trace_id = str(uuid.uuid4())
|
||||
trace_store.init_trace(trace_id, "Corrupt me", "Flash")
|
||||
|
||||
status_path = trace_store.get_status_path(trace_id)
|
||||
with open(status_path, "w", encoding="utf-8") as f:
|
||||
f.write('{"status": "running", "trace_id"') # torn write
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="mcp_server.utils.trace_store"):
|
||||
assert trace_store.read_status(trace_id) is None
|
||||
|
||||
assert any("Corrupt status.json" in rec.message for rec in caplog.records)
|
||||
assert os.path.exists(f"{status_path}.corrupt")
|
||||
assert not os.path.exists(status_path)
|
||||
|
||||
|
||||
def test_missing_status_stays_silent(temp_trace_env, caplog):
|
||||
with caplog.at_level(logging.WARNING, logger="mcp_server.utils.trace_store"):
|
||||
assert trace_store.read_status(str(uuid.uuid4())) is None
|
||||
assert caplog.records == []
|
||||
|
||||
|
||||
def test_update_on_corrupt_status_logs_and_drops(temp_trace_env, caplog):
|
||||
trace_id = str(uuid.uuid4())
|
||||
trace_store.init_trace(trace_id, "Corrupt update", "Flash")
|
||||
|
||||
status_path = trace_store.get_status_path(trace_id)
|
||||
with open(status_path, "w", encoding="utf-8") as f:
|
||||
f.write("not json at all")
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="mcp_server.utils.trace_store"):
|
||||
assert trace_store.update_trace_status(trace_id, "completed") is None
|
||||
|
||||
assert any("Dropping status update" in rec.message for rec in caplog.records)
|
||||
# Updating a trace that never existed stays a silent no-op.
|
||||
caplog.clear()
|
||||
with caplog.at_level(logging.WARNING, logger="mcp_server.utils.trace_store"):
|
||||
assert trace_store.update_trace_status(str(uuid.uuid4()), "completed") is None
|
||||
assert not any("Dropping status update" in rec.message for rec in caplog.records)
|
||||
|
||||
|
||||
def test_concurrent_read_modify_write_does_not_lose_updates(temp_trace_env):
|
||||
"""Interleaved RMW updates from two writers must not lose each other's fields."""
|
||||
trace_id = str(uuid.uuid4())
|
||||
trace_store.init_trace(trace_id, "Concurrency", "Flash")
|
||||
iterations = 25
|
||||
failures: list[BaseException] = []
|
||||
|
||||
def set_errors():
|
||||
try:
|
||||
for i in range(iterations):
|
||||
assert trace_store.update_trace_status(
|
||||
trace_id, "running", error=f"err-{i}"
|
||||
) is not None
|
||||
except BaseException as exc: # noqa: BLE001 - surfaced in the main thread
|
||||
failures.append(exc)
|
||||
|
||||
def set_serials():
|
||||
try:
|
||||
for i in range(iterations):
|
||||
assert trace_store.update_trace_device_serial(
|
||||
trace_id, f"serial-{i}"
|
||||
) is not None
|
||||
except BaseException as exc: # noqa: BLE001 - surfaced in the main thread
|
||||
failures.append(exc)
|
||||
|
||||
threads = [threading.Thread(target=set_errors), threading.Thread(target=set_serials)]
|
||||
for t in threads:
|
||||
t.start()
|
||||
for t in threads:
|
||||
t.join(timeout=30.0)
|
||||
|
||||
assert failures == []
|
||||
final = trace_store.read_status(trace_id)
|
||||
assert final is not None
|
||||
# Every serialized RMW preserves the other writer's latest field, so both
|
||||
# final values must survive regardless of which thread finished last.
|
||||
assert final["error"] == f"err-{iterations - 1}"
|
||||
assert final["device_serial"] == f"serial-{iterations - 1}"
|
||||
assert final["status"] == "running"
|
||||
|
||||
|
||||
def test_trace_paths(temp_trace_env):
|
||||
trace_id = "test-paths-id"
|
||||
trace_dir = trace_store.get_trace_dir(trace_id)
|
||||
|
||||
@@ -185,6 +185,78 @@ def test_submission_reservations_preserve_order_before_workers_start():
|
||||
thread.join(timeout=2.0)
|
||||
|
||||
|
||||
def test_pending_reservation_at_queue_head_does_not_block_other_devices():
|
||||
"""An unclaimed pending submission must not gate every device's queue head.
|
||||
|
||||
Regression: a reservation whose device is still "pending" used to be
|
||||
counted into every device's FIFO queue, so one unclaimed head-of-queue
|
||||
ticket blocked all devices at once. It must now park on at most one idle
|
||||
candidate device (the first in sorted order) and leave the others
|
||||
schedulable.
|
||||
"""
|
||||
pending_ticket = DeviceExecutionLock.reserve("unclaimed submission")
|
||||
# device-alpha is a known idle device (a queued submission targets it), so
|
||||
# the pending ticket parks there; device-beta must stay schedulable.
|
||||
alpha_ticket = DeviceExecutionLock.reserve(
|
||||
"task for device alpha", device_id="device-alpha"
|
||||
)
|
||||
try:
|
||||
beta = DeviceExecutionLock("device-beta", "task for device beta")
|
||||
# Must acquire immediately even though the older pending ticket is at
|
||||
# the head of the global queue directory.
|
||||
beta.acquire(blocking=False)
|
||||
beta.release()
|
||||
finally:
|
||||
DeviceExecutionLock.cancel_reservation(pending_ticket)
|
||||
DeviceExecutionLock.cancel_reservation(alpha_ticket)
|
||||
|
||||
|
||||
def test_pending_reservation_is_served_with_original_fifo_position_after_claim():
|
||||
"""A pending reservation is not starved: once its worker claims a device it
|
||||
keeps its original submission-order position in that device's queue."""
|
||||
holder = DeviceExecutionLock("emulator-5554", "current holder")
|
||||
holder.acquire()
|
||||
|
||||
# Submitted while the device is busy; the reservation stays "pending".
|
||||
first_ticket = DeviceExecutionLock.reserve("first submitted task")
|
||||
|
||||
claimed = DeviceExecutionLock("emulator-5554", "claimed pending task", first_ticket)
|
||||
youngest = DeviceExecutionLock("emulator-5554", "youngest task")
|
||||
order = []
|
||||
claimed_acquired = threading.Event()
|
||||
youngest_acquired = threading.Event()
|
||||
|
||||
def run_claimed():
|
||||
claimed.acquire()
|
||||
order.append("claimed")
|
||||
claimed_acquired.set()
|
||||
|
||||
def run_youngest():
|
||||
youngest.acquire()
|
||||
order.append("youngest")
|
||||
youngest_acquired.set()
|
||||
|
||||
t1 = threading.Thread(target=run_claimed)
|
||||
t2 = threading.Thread(target=run_youngest)
|
||||
t1.start()
|
||||
t2.start()
|
||||
time.sleep(0.2)
|
||||
assert not claimed_acquired.is_set()
|
||||
assert not youngest_acquired.is_set()
|
||||
|
||||
holder.release()
|
||||
# The claimed pending ticket kept its original (oldest) FIFO timestamp and
|
||||
# must win the device before the younger concrete ticket.
|
||||
assert claimed_acquired.wait(timeout=2.0)
|
||||
assert not youngest_acquired.is_set()
|
||||
claimed.release()
|
||||
assert youngest_acquired.wait(timeout=2.0)
|
||||
youngest.release()
|
||||
t1.join(timeout=2.0)
|
||||
t2.join(timeout=2.0)
|
||||
assert order == ["claimed", "youngest"]
|
||||
|
||||
|
||||
def test_cancelled_reservation_cannot_execute():
|
||||
ticket = DeviceExecutionLock.reserve("cancelled task")
|
||||
assert DeviceExecutionLock.cancel_reservation(ticket)
|
||||
|
||||
@@ -0,0 +1,65 @@
|
||||
# Copyright 2026 Google LLC
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Unit tests for the unified PID liveness probe."""
|
||||
|
||||
import os
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import psutil
|
||||
|
||||
from artemis.runtime.process_probe import pid_is_alive
|
||||
|
||||
|
||||
def test_own_pid_is_alive():
|
||||
assert pid_is_alive(os.getpid()) is True
|
||||
|
||||
|
||||
def test_own_pid_with_matching_create_time_is_alive():
|
||||
created_at = psutil.Process(os.getpid()).create_time()
|
||||
assert pid_is_alive(os.getpid(), created_at) is True
|
||||
|
||||
|
||||
def test_recycled_pid_with_mismatched_create_time_is_dead():
|
||||
assert pid_is_alive(os.getpid(), 1.0) is False
|
||||
|
||||
|
||||
def test_invalid_pids_are_dead():
|
||||
assert pid_is_alive(0) is False
|
||||
assert pid_is_alive(-5) is False
|
||||
assert pid_is_alive(None) is False # type: ignore[arg-type]
|
||||
assert pid_is_alive("not-a-pid") is False # type: ignore[arg-type]
|
||||
|
||||
|
||||
def test_no_such_process_is_dead():
|
||||
with patch("psutil.Process", side_effect=psutil.NoSuchProcess(424242)):
|
||||
assert pid_is_alive(424242) is False
|
||||
|
||||
|
||||
def test_zombie_process_is_dead():
|
||||
proc = MagicMock()
|
||||
proc.is_running.return_value = True
|
||||
proc.status.return_value = psutil.STATUS_ZOMBIE
|
||||
with patch("psutil.Process", return_value=proc):
|
||||
assert pid_is_alive(1234) is False
|
||||
|
||||
|
||||
def test_access_denied_defaults_to_alive():
|
||||
with patch("psutil.Process", side_effect=psutil.AccessDenied(1234)):
|
||||
assert pid_is_alive(1234) is True
|
||||
|
||||
|
||||
def test_unexpected_error_defaults_to_alive():
|
||||
with patch("psutil.Process", side_effect=RuntimeError("boom")):
|
||||
assert pid_is_alive(1234) is True
|
||||
Reference in new issue
Block a user