feat: complete production workflow migration
This commit is contained in:
1 parent
7f215e79c4
commit
8df5266abb
448 files changed
+56937
-14619
No files matched your search
@@ -1,6 +1,11 @@
|
||||
"""Workflow definitions, executable steps, and orchestration engine."""
|
||||
|
||||
from .engine import WorkflowEngine, WorkflowRunResult
|
||||
from .graph import (
|
||||
WorkflowGraphState,
|
||||
compile_workflow_graph,
|
||||
initial_workflow_graph_state,
|
||||
)
|
||||
from .model import StepDefinition, WorkflowDefinition, WorkflowValidationError
|
||||
from .steps import CommandStep, StepExecution
|
||||
|
||||
@@ -10,6 +15,9 @@ __all__ = [
|
||||
"StepExecution",
|
||||
"WorkflowDefinition",
|
||||
"WorkflowEngine",
|
||||
"WorkflowGraphState",
|
||||
"WorkflowRunResult",
|
||||
"WorkflowValidationError",
|
||||
"compile_workflow_graph",
|
||||
"initial_workflow_graph_state",
|
||||
]
|
||||
@@ -10,6 +10,7 @@ from typing import Callable, Literal
|
||||
from gyxx_flow.core.context import RunContext
|
||||
from gyxx_flow.core.locks import LockManager
|
||||
from gyxx_flow.core.records import RunJournal
|
||||
from gyxx_flow.core.text import bounded_head_tail
|
||||
from gyxx_flow.ops.effects import (
|
||||
EffectAlreadyApplied,
|
||||
EffectClaim,
|
||||
@@ -17,12 +18,21 @@ from gyxx_flow.ops.effects import (
|
||||
EffectStateAmbiguous,
|
||||
)
|
||||
|
||||
from .graph import (
|
||||
GraphStepResult,
|
||||
WorkflowGraphState,
|
||||
compile_workflow_graph,
|
||||
initial_workflow_graph_state,
|
||||
)
|
||||
from .model import StepDefinition, WorkflowDefinition
|
||||
from .steps import StepExecution
|
||||
|
||||
StepRunStatus = Literal["success", "failed", "skipped"]
|
||||
WorkflowRunStatus = Literal["success", "failed"]
|
||||
|
||||
_STEP_FAILURE_DETAIL_MAX_CHARS = 1_000
|
||||
_WORKFLOW_ERROR_MAX_CHARS = 4_096
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class AttemptResult:
|
||||
@@ -80,60 +90,41 @@ class WorkflowEngine:
|
||||
f"{context.workflow_id!r} != {workflow.workflow_id!r}"
|
||||
)
|
||||
|
||||
results: dict[str, StepRunResult] = {}
|
||||
def run_step(
|
||||
step: StepDefinition,
|
||||
state: WorkflowGraphState,
|
||||
) -> GraphStepResult:
|
||||
prior_results = {
|
||||
step_id: _step_run_result(result)
|
||||
for step_id, result in state["step_results"].items()
|
||||
}
|
||||
result = self._execute_graph_step(
|
||||
step,
|
||||
results=prior_results,
|
||||
context=context,
|
||||
journal=journal,
|
||||
dry_run=dry_run,
|
||||
)
|
||||
return _graph_step_result(result)
|
||||
|
||||
ordered_steps = workflow.ordered_steps()
|
||||
graph = compile_workflow_graph(workflow, run_step=run_step)
|
||||
with self._lock_manager.acquire(
|
||||
f"workflow:{workflow.workflow_id}", owner=context.run_id
|
||||
):
|
||||
for step in workflow.ordered_steps():
|
||||
failed_dependency = any(
|
||||
_blocks_dependents(results[dependency])
|
||||
for dependency in step.depends_on
|
||||
)
|
||||
if failed_dependency:
|
||||
results[step.step_id] = self._record_policy_skip(
|
||||
step, journal, reason="dependency-failed"
|
||||
)
|
||||
continue
|
||||
if dry_run:
|
||||
results[step.step_id] = self._record_policy_skip(
|
||||
step, journal, reason="dry-run"
|
||||
)
|
||||
continue
|
||||
if context.shadow and (
|
||||
step.production_sink or step.official_notification
|
||||
):
|
||||
results[step.step_id] = self._record_policy_skip(
|
||||
step, journal, reason="shadow-policy"
|
||||
)
|
||||
continue
|
||||
claim: EffectClaim | None = None
|
||||
if step.production_sink and self._effect_ledger is not None:
|
||||
try:
|
||||
claim = self._effect_ledger.begin(
|
||||
workflow_id=context.workflow_id,
|
||||
business_date=context.business_date.isoformat(),
|
||||
step_id=step.step_id,
|
||||
run_id=context.run_id,
|
||||
)
|
||||
except EffectAlreadyApplied:
|
||||
results[step.step_id] = self._record_policy_skip(
|
||||
step, journal, reason="idempotency-replay"
|
||||
)
|
||||
continue
|
||||
except EffectStateAmbiguous:
|
||||
results[step.step_id] = self._record_policy_failure(
|
||||
step, journal, reason="effect-state-ambiguous"
|
||||
)
|
||||
continue
|
||||
result = self._execute_step(step, context=context, journal=journal)
|
||||
if claim is not None:
|
||||
if result.status == "success":
|
||||
self._effect_ledger.mark_applied(claim)
|
||||
else:
|
||||
self._effect_ledger.mark_ambiguous(
|
||||
claim, reason="step-did-not-succeed"
|
||||
)
|
||||
results[step.step_id] = result
|
||||
final_state = graph.invoke(
|
||||
initial_workflow_graph_state(context, dry_run=dry_run),
|
||||
config={"recursion_limit": max(25, len(workflow.steps) + 2)},
|
||||
)
|
||||
unordered_results = {
|
||||
step_id: _step_run_result(result)
|
||||
for step_id, result in final_state["step_results"].items()
|
||||
}
|
||||
results = {
|
||||
step.step_id: unordered_results[step.step_id]
|
||||
for step in ordered_steps
|
||||
}
|
||||
journal.normalize_step_order(tuple(results))
|
||||
|
||||
critical_failures = tuple(
|
||||
step_id
|
||||
@@ -148,7 +139,7 @@ class WorkflowEngine:
|
||||
status: WorkflowRunStatus = "failed" if critical_failures else "success"
|
||||
error = None
|
||||
if critical_failures:
|
||||
error = f"critical steps failed: {', '.join(critical_failures)}"
|
||||
error = _critical_failure_error(critical_failures, results)
|
||||
journal.finalize(status, error=error)
|
||||
return WorkflowRunResult(
|
||||
workflow_id=workflow.workflow_id,
|
||||
@@ -158,6 +149,66 @@ class WorkflowEngine:
|
||||
warnings=warnings,
|
||||
)
|
||||
|
||||
def _execute_graph_step(
|
||||
self,
|
||||
step: StepDefinition,
|
||||
*,
|
||||
results: dict[str, StepRunResult],
|
||||
context: RunContext,
|
||||
journal: RunJournal,
|
||||
dry_run: bool,
|
||||
) -> StepRunResult:
|
||||
failed_dependency = any(
|
||||
_blocks_dependents(results[dependency])
|
||||
for dependency in step.depends_on
|
||||
)
|
||||
if failed_dependency and not step.run_after_failure:
|
||||
return self._record_policy_skip(
|
||||
step, journal, reason="dependency-failed"
|
||||
)
|
||||
if dry_run:
|
||||
return self._record_policy_skip(step, journal, reason="dry-run")
|
||||
if context.shadow and (
|
||||
step.production_sink or step.official_notification
|
||||
):
|
||||
return self._record_policy_skip(
|
||||
step, journal, reason="shadow-policy"
|
||||
)
|
||||
|
||||
claim: EffectClaim | None = None
|
||||
if (
|
||||
step.production_sink
|
||||
and step.replay_policy == "guarded"
|
||||
and self._effect_ledger is not None
|
||||
):
|
||||
try:
|
||||
claim = self._effect_ledger.begin(
|
||||
workflow_id=context.workflow_id,
|
||||
business_date=context.business_date.isoformat(),
|
||||
step_id=step.step_id,
|
||||
run_id=context.run_id,
|
||||
)
|
||||
except EffectAlreadyApplied:
|
||||
return self._record_policy_skip(
|
||||
step, journal, reason="idempotency-replay"
|
||||
)
|
||||
except EffectStateAmbiguous:
|
||||
return self._record_policy_failure(
|
||||
step, journal, reason="effect-state-ambiguous"
|
||||
)
|
||||
|
||||
result = self._execute_step(step, context=context, journal=journal)
|
||||
if claim is not None:
|
||||
if result.status == "success":
|
||||
self._effect_ledger.mark_applied(claim)
|
||||
elif result.status == "skipped":
|
||||
self._effect_ledger.cancel(claim)
|
||||
else:
|
||||
self._effect_ledger.mark_ambiguous(
|
||||
claim, reason="step-did-not-succeed"
|
||||
)
|
||||
return result
|
||||
|
||||
def _execute_step(
|
||||
self,
|
||||
step: StepDefinition,
|
||||
@@ -320,3 +371,80 @@ def _blocks_dependents(result: StepRunResult) -> bool:
|
||||
|
||||
def _is_failure(result: StepRunResult) -> bool:
|
||||
return result.status == "failed" or result.reason == "dependency-failed"
|
||||
|
||||
|
||||
def _critical_failure_error(
|
||||
step_ids: tuple[str, ...],
|
||||
results: dict[str, StepRunResult],
|
||||
) -> str:
|
||||
details = "; ".join(
|
||||
f"{step_id}: {_step_failure_detail(results[step_id])}"
|
||||
for step_id in step_ids
|
||||
)
|
||||
return _bounded_text(
|
||||
f"critical steps failed: {details}",
|
||||
_WORKFLOW_ERROR_MAX_CHARS,
|
||||
)
|
||||
|
||||
|
||||
def _step_failure_detail(result: StepRunResult) -> str:
|
||||
attempt = result.attempts[-1] if result.attempts else None
|
||||
detail = (
|
||||
(attempt.error if attempt is not None else None)
|
||||
or (attempt.reason if attempt is not None else None)
|
||||
or result.reason
|
||||
)
|
||||
if not detail:
|
||||
detail = (
|
||||
f"exit code {result.final_exit_code}"
|
||||
if result.final_exit_code is not None
|
||||
else "unknown failure"
|
||||
)
|
||||
return _bounded_text(
|
||||
" ".join(str(detail).split()),
|
||||
_STEP_FAILURE_DETAIL_MAX_CHARS,
|
||||
)
|
||||
|
||||
|
||||
def _bounded_text(value: str, limit: int) -> str:
|
||||
return bounded_head_tail(value, limit)
|
||||
|
||||
|
||||
def _graph_step_result(result: StepRunResult) -> GraphStepResult:
|
||||
return {
|
||||
"step_id": result.step_id,
|
||||
"status": result.status,
|
||||
"attempts": [
|
||||
{
|
||||
"attempt": attempt.attempt,
|
||||
"status": attempt.status,
|
||||
"exit_code": attempt.exit_code,
|
||||
"error": attempt.error,
|
||||
"reason": attempt.reason,
|
||||
}
|
||||
for attempt in result.attempts
|
||||
],
|
||||
"final_exit_code": result.final_exit_code,
|
||||
"critical": result.critical,
|
||||
"reason": result.reason,
|
||||
}
|
||||
|
||||
|
||||
def _step_run_result(result: GraphStepResult) -> StepRunResult:
|
||||
return StepRunResult(
|
||||
step_id=result["step_id"],
|
||||
status=result["status"],
|
||||
attempts=tuple(
|
||||
AttemptResult(
|
||||
attempt=attempt["attempt"],
|
||||
status=attempt["status"],
|
||||
exit_code=attempt["exit_code"],
|
||||
error=attempt["error"],
|
||||
reason=attempt["reason"],
|
||||
)
|
||||
for attempt in result["attempts"]
|
||||
),
|
||||
final_exit_code=result["final_exit_code"],
|
||||
critical=result["critical"],
|
||||
reason=result["reason"],
|
||||
)
|
||||
@@ -0,0 +1,87 @@
|
||||
"""Shared catalog-to-LangGraph workflow composition."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
from gyxx_flow.adapters.native import DeferredModuleCommandStep, ModuleCommandFactory
|
||||
from gyxx_flow.catalog import WorkflowEntry
|
||||
|
||||
from .model import StepDefinition, WorkflowDefinition
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _ConfiguredStep:
|
||||
step_id: str
|
||||
entry: str
|
||||
args: tuple[str, ...]
|
||||
depends_on: tuple[str, ...]
|
||||
run_after_failure: bool = False
|
||||
timeout_seconds: float | None = None
|
||||
replay_policy: str | None = None
|
||||
|
||||
|
||||
def build_catalog_workflow(
|
||||
entry: WorkflowEntry,
|
||||
*,
|
||||
default_step_id: str,
|
||||
timeout_seconds: float,
|
||||
resource: str,
|
||||
command_factory: ModuleCommandFactory | None = None,
|
||||
official_notification: bool = False,
|
||||
independent_step_resources: bool = False,
|
||||
default_replay_policy: str = "guarded",
|
||||
) -> WorkflowDefinition:
|
||||
"""Build one executable graph without adding one class per business task."""
|
||||
|
||||
if entry.steps:
|
||||
configured_steps = entry.steps
|
||||
else:
|
||||
configured_steps = (
|
||||
_ConfiguredStep(
|
||||
step_id=default_step_id,
|
||||
entry=entry.entry,
|
||||
args=entry.args,
|
||||
depends_on=(),
|
||||
),
|
||||
)
|
||||
|
||||
definitions: list[StepDefinition] = []
|
||||
for configured in configured_steps:
|
||||
action_kwargs = {
|
||||
"command_entry": configured.entry,
|
||||
"command_args": configured.args,
|
||||
}
|
||||
if command_factory is not None:
|
||||
action_kwargs["command_factory"] = command_factory
|
||||
action = DeferredModuleCommandStep(entry, **action_kwargs)
|
||||
step_resources = (
|
||||
(f"{resource}:{configured.step_id}",)
|
||||
if independent_step_resources
|
||||
else (resource,)
|
||||
)
|
||||
definitions.append(
|
||||
StepDefinition(
|
||||
configured.step_id,
|
||||
action,
|
||||
depends_on=configured.depends_on,
|
||||
timeout_seconds=(
|
||||
configured.timeout_seconds
|
||||
if configured.timeout_seconds is not None
|
||||
else timeout_seconds
|
||||
),
|
||||
max_attempts=1,
|
||||
resources=step_resources,
|
||||
production_sink=True,
|
||||
official_notification=official_notification,
|
||||
replay_policy=(
|
||||
getattr(configured, "replay_policy", None)
|
||||
or default_replay_policy
|
||||
),
|
||||
run_after_failure=getattr(configured, "run_after_failure", False),
|
||||
)
|
||||
)
|
||||
return WorkflowDefinition(entry.workflow_id, tuple(definitions))
|
||||
|
||||
|
||||
__all__ = ["build_catalog_workflow"]
|
||||
@@ -0,0 +1,119 @@
|
||||
"""LangGraph compilation for deterministic workflow definitions."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Annotated, Callable, Literal, TypedDict
|
||||
|
||||
from langgraph.graph import END, START, StateGraph
|
||||
from langgraph.graph.state import CompiledStateGraph
|
||||
|
||||
from gyxx_flow.core.context import RunContext
|
||||
|
||||
from .model import StepDefinition, WorkflowDefinition
|
||||
|
||||
GraphStepStatus = Literal["success", "failed", "skipped"]
|
||||
|
||||
|
||||
class GraphAttemptResult(TypedDict):
|
||||
"""JSON-serializable attempt result stored in graph state."""
|
||||
|
||||
attempt: int
|
||||
status: GraphStepStatus
|
||||
exit_code: int | None
|
||||
error: str | None
|
||||
reason: str | None
|
||||
|
||||
|
||||
class GraphStepResult(TypedDict):
|
||||
"""JSON-serializable step result stored in graph state."""
|
||||
|
||||
step_id: str
|
||||
status: GraphStepStatus
|
||||
attempts: list[GraphAttemptResult]
|
||||
final_exit_code: int | None
|
||||
critical: bool
|
||||
reason: str | None
|
||||
|
||||
|
||||
class WorkflowGraphState(TypedDict):
|
||||
"""Serializable state passed between one workflow's LangGraph nodes."""
|
||||
|
||||
workflow_id: str
|
||||
run_id: str
|
||||
business_date: str
|
||||
started_at: str
|
||||
shadow: bool
|
||||
dry_run: bool
|
||||
step_results: Annotated[
|
||||
dict[str, GraphStepResult],
|
||||
_merge_step_results,
|
||||
]
|
||||
|
||||
|
||||
GraphStepRunner = Callable[
|
||||
[StepDefinition, WorkflowGraphState],
|
||||
GraphStepResult,
|
||||
]
|
||||
|
||||
|
||||
def initial_workflow_graph_state(
|
||||
context: RunContext,
|
||||
*,
|
||||
dry_run: bool,
|
||||
) -> WorkflowGraphState:
|
||||
"""Create a JSON-serializable initial state from the public run context."""
|
||||
|
||||
return {
|
||||
"workflow_id": context.workflow_id,
|
||||
"run_id": context.run_id,
|
||||
"business_date": context.business_date.isoformat(),
|
||||
"started_at": context.started_at.isoformat(),
|
||||
"shadow": context.shadow,
|
||||
"dry_run": dry_run,
|
||||
"step_results": {},
|
||||
}
|
||||
|
||||
|
||||
def compile_workflow_graph(
|
||||
workflow: WorkflowDefinition,
|
||||
*,
|
||||
run_step: GraphStepRunner,
|
||||
) -> CompiledStateGraph:
|
||||
"""Compile dependencies into a real LangGraph DAG."""
|
||||
|
||||
builder = StateGraph(WorkflowGraphState)
|
||||
ordered_steps = workflow.ordered_steps()
|
||||
dependent_ids = {
|
||||
dependency
|
||||
for step in ordered_steps
|
||||
for dependency in step.depends_on
|
||||
}
|
||||
for step in ordered_steps:
|
||||
builder.add_node(step.step_id, _step_node(step, run_step))
|
||||
if step.depends_on:
|
||||
builder.add_edge(list(step.depends_on), step.step_id)
|
||||
else:
|
||||
builder.add_edge(START, step.step_id)
|
||||
if step.step_id not in dependent_ids:
|
||||
builder.add_edge(step.step_id, END)
|
||||
return builder.compile()
|
||||
|
||||
|
||||
def _step_node(
|
||||
step: StepDefinition,
|
||||
run_step: GraphStepRunner,
|
||||
) -> Callable[[WorkflowGraphState], dict[str, object]]:
|
||||
def execute(state: WorkflowGraphState) -> dict[str, object]:
|
||||
result = run_step(step, state)
|
||||
return {"step_results": {step.step_id: result}}
|
||||
|
||||
return execute
|
||||
|
||||
|
||||
def _merge_step_results(
|
||||
current: dict[str, GraphStepResult],
|
||||
update: dict[str, GraphStepResult],
|
||||
) -> dict[str, GraphStepResult]:
|
||||
"""Merge branch-local results emitted in the same LangGraph superstep."""
|
||||
|
||||
return {**current, **update}
|
||||
@@ -3,9 +3,10 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import heapq
|
||||
import math
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Protocol
|
||||
from typing import TYPE_CHECKING, Literal, Protocol
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from gyxx_flow.core.context import RunContext
|
||||
@@ -15,6 +16,7 @@ if TYPE_CHECKING:
|
||||
|
||||
_ID = re.compile(r"^[a-z][a-z0-9]*(?:[._-][a-z0-9]+)*$")
|
||||
_RESOURCE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9:._-]*$")
|
||||
ReplayPolicy = Literal["guarded", "idempotent"]
|
||||
|
||||
|
||||
class WorkflowValidationError(ValueError):
|
||||
@@ -47,6 +49,8 @@ class StepDefinition:
|
||||
resources: tuple[str, ...] = ()
|
||||
production_sink: bool = False
|
||||
official_notification: bool = False
|
||||
replay_policy: ReplayPolicy = "guarded"
|
||||
run_after_failure: bool = False
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
object.__setattr__(self, "depends_on", tuple(self.depends_on))
|
||||
@@ -65,7 +69,10 @@ class StepDefinition:
|
||||
if dependency == self.step_id:
|
||||
raise WorkflowValidationError(f"step {self.step_id!r} depends on itself")
|
||||
if self.timeout_seconds is not None and (
|
||||
isinstance(self.timeout_seconds, bool) or self.timeout_seconds <= 0
|
||||
isinstance(self.timeout_seconds, bool)
|
||||
or not isinstance(self.timeout_seconds, (int, float))
|
||||
or not math.isfinite(self.timeout_seconds)
|
||||
or self.timeout_seconds <= 0
|
||||
):
|
||||
raise WorkflowValidationError(
|
||||
f"timeout_seconds must be positive for step {self.step_id!r}"
|
||||
@@ -89,6 +96,20 @@ class StepDefinition:
|
||||
raise WorkflowValidationError(
|
||||
f"invalid resource for step {self.step_id!r}: {resource!r}"
|
||||
)
|
||||
if not isinstance(self.run_after_failure, bool):
|
||||
raise WorkflowValidationError(
|
||||
f"run_after_failure must be boolean for step {self.step_id!r}"
|
||||
)
|
||||
if self.replay_policy not in {"guarded", "idempotent"}:
|
||||
raise WorkflowValidationError(
|
||||
f"invalid replay_policy for step {self.step_id!r}: "
|
||||
f"{self.replay_policy!r}"
|
||||
)
|
||||
if self.replay_policy == "idempotent" and not self.production_sink:
|
||||
raise WorkflowValidationError(
|
||||
f"idempotent replay_policy requires production_sink for step "
|
||||
f"{self.step_id!r}"
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
||||
+361
-30
@@ -2,12 +2,129 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import signal
|
||||
import subprocess
|
||||
import tempfile
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Mapping
|
||||
from typing import IO, Mapping
|
||||
|
||||
import psutil
|
||||
|
||||
from gyxx_flow.core.context import RunContext
|
||||
from gyxx_flow.core.exit_codes import COOKIE_SKIP_EXIT_CODE
|
||||
from gyxx_flow.core.text import bounded_head_tail
|
||||
|
||||
_CREATE_SUSPENDED = 0x00000004
|
||||
_TERMINATION_GRACE_SECONDS = 0.5
|
||||
_TERMINATION_CONFIRM_SECONDS = 5.0
|
||||
_FAILURE_OUTPUT_MAX_CHARS = 16_000
|
||||
|
||||
|
||||
class _WindowsJob:
|
||||
"""A private Job Object that contains exactly one command process tree."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
if os.name != "nt":
|
||||
raise OSError("Windows Job Objects are unavailable on this platform")
|
||||
|
||||
import ctypes
|
||||
from ctypes import wintypes
|
||||
|
||||
class BasicAccountingInformation(ctypes.Structure):
|
||||
_fields_ = [
|
||||
("TotalUserTime", ctypes.c_longlong),
|
||||
("TotalKernelTime", ctypes.c_longlong),
|
||||
("ThisPeriodTotalUserTime", ctypes.c_longlong),
|
||||
("ThisPeriodTotalKernelTime", ctypes.c_longlong),
|
||||
("TotalPageFaultCount", wintypes.DWORD),
|
||||
("TotalProcesses", wintypes.DWORD),
|
||||
("ActiveProcesses", wintypes.DWORD),
|
||||
("TotalTerminatedProcesses", wintypes.DWORD),
|
||||
]
|
||||
|
||||
kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)
|
||||
kernel32.CreateJobObjectW.argtypes = (wintypes.LPVOID, wintypes.LPCWSTR)
|
||||
kernel32.CreateJobObjectW.restype = wintypes.HANDLE
|
||||
kernel32.OpenProcess.argtypes = (
|
||||
wintypes.DWORD,
|
||||
wintypes.BOOL,
|
||||
wintypes.DWORD,
|
||||
)
|
||||
kernel32.OpenProcess.restype = wintypes.HANDLE
|
||||
kernel32.AssignProcessToJobObject.argtypes = (
|
||||
wintypes.HANDLE,
|
||||
wintypes.HANDLE,
|
||||
)
|
||||
kernel32.AssignProcessToJobObject.restype = wintypes.BOOL
|
||||
kernel32.TerminateJobObject.argtypes = (wintypes.HANDLE, wintypes.UINT)
|
||||
kernel32.TerminateJobObject.restype = wintypes.BOOL
|
||||
kernel32.QueryInformationJobObject.argtypes = (
|
||||
wintypes.HANDLE,
|
||||
ctypes.c_int,
|
||||
wintypes.LPVOID,
|
||||
wintypes.DWORD,
|
||||
wintypes.LPVOID,
|
||||
)
|
||||
kernel32.QueryInformationJobObject.restype = wintypes.BOOL
|
||||
kernel32.CloseHandle.argtypes = (wintypes.HANDLE,)
|
||||
kernel32.CloseHandle.restype = wintypes.BOOL
|
||||
|
||||
handle = kernel32.CreateJobObjectW(None, None)
|
||||
if not handle:
|
||||
raise ctypes.WinError(ctypes.get_last_error())
|
||||
self._ctypes = ctypes
|
||||
self._kernel32 = kernel32
|
||||
self._accounting_type = BasicAccountingInformation
|
||||
self._handle = handle
|
||||
|
||||
def assign(self, pid: int) -> None:
|
||||
process_set_quota = 0x0100
|
||||
process_terminate = 0x0001
|
||||
process_handle = self._kernel32.OpenProcess(
|
||||
process_set_quota | process_terminate,
|
||||
False,
|
||||
pid,
|
||||
)
|
||||
if not process_handle:
|
||||
raise self._ctypes.WinError(self._ctypes.get_last_error())
|
||||
try:
|
||||
if not self._kernel32.AssignProcessToJobObject(
|
||||
self._handle,
|
||||
process_handle,
|
||||
):
|
||||
raise self._ctypes.WinError(self._ctypes.get_last_error())
|
||||
finally:
|
||||
self._kernel32.CloseHandle(process_handle)
|
||||
|
||||
def terminate(self, exit_code: int) -> None:
|
||||
if not self._kernel32.TerminateJobObject(self._handle, exit_code):
|
||||
raise self._ctypes.WinError(self._ctypes.get_last_error())
|
||||
|
||||
def wait_empty(self, timeout_seconds: float) -> None:
|
||||
deadline = time.monotonic() + timeout_seconds
|
||||
while True:
|
||||
information = self._accounting_type()
|
||||
if not self._kernel32.QueryInformationJobObject(
|
||||
self._handle,
|
||||
1,
|
||||
self._ctypes.byref(information),
|
||||
self._ctypes.sizeof(information),
|
||||
None,
|
||||
):
|
||||
raise self._ctypes.WinError(self._ctypes.get_last_error())
|
||||
if information.ActiveProcesses == 0:
|
||||
return
|
||||
if time.monotonic() >= deadline:
|
||||
raise RuntimeError("command process tree did not terminate")
|
||||
time.sleep(0.01)
|
||||
|
||||
def close(self) -> None:
|
||||
if self._handle:
|
||||
self._kernel32.CloseHandle(self._handle)
|
||||
self._handle = None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
@@ -59,43 +176,257 @@ class CommandStep:
|
||||
del context
|
||||
if dry_run:
|
||||
return StepExecution(exit_code=0, skipped=True, reason="dry-run")
|
||||
try:
|
||||
completed = subprocess.run(
|
||||
list(self.argv),
|
||||
cwd=self.cwd,
|
||||
env=dict(self.env),
|
||||
shell=False,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=timeout_seconds,
|
||||
check=False,
|
||||
)
|
||||
except subprocess.TimeoutExpired as exc:
|
||||
stdout = _to_text(exc.stdout)
|
||||
stderr = _to_text(exc.stderr)
|
||||
with (
|
||||
tempfile.TemporaryFile(
|
||||
mode="w+", encoding="utf-8", errors="replace", newline=""
|
||||
) as stdout_file,
|
||||
tempfile.TemporaryFile(
|
||||
mode="w+", encoding="utf-8", errors="replace", newline=""
|
||||
) as stderr_file,
|
||||
):
|
||||
try:
|
||||
process, process_scope = _start_process(
|
||||
list(self.argv),
|
||||
cwd=self.cwd,
|
||||
env=dict(self.env),
|
||||
stdout=stdout_file,
|
||||
stderr=stderr_file,
|
||||
)
|
||||
except OSError as exc:
|
||||
return StepExecution(exit_code=127, error=f"cannot start command: {exc}")
|
||||
|
||||
try:
|
||||
return_code = process.wait(timeout=timeout_seconds)
|
||||
except subprocess.TimeoutExpired as exc:
|
||||
_terminate_process_tree(process, process_scope)
|
||||
stdout = _read_output(stdout_file)
|
||||
stderr = _read_output(stderr_file)
|
||||
return StepExecution(
|
||||
exit_code=124,
|
||||
stdout=stdout,
|
||||
stderr=stderr,
|
||||
error=f"command timed out after {exc.timeout} seconds",
|
||||
)
|
||||
except BaseException:
|
||||
_terminate_process_tree(process, process_scope)
|
||||
raise
|
||||
finally:
|
||||
if process_scope is not None:
|
||||
process_scope.close()
|
||||
|
||||
stdout = _read_output(stdout_file)
|
||||
stderr = _read_output(stderr_file)
|
||||
|
||||
if return_code == COOKIE_SKIP_EXIT_CODE:
|
||||
return StepExecution(
|
||||
exit_code=124,
|
||||
exit_code=return_code,
|
||||
stdout=stdout,
|
||||
stderr=stderr,
|
||||
error=f"command timed out after {exc.timeout} seconds",
|
||||
skipped=True,
|
||||
reason="cookie-preflight",
|
||||
)
|
||||
except OSError as exc:
|
||||
return StepExecution(exit_code=127, error=f"cannot start command: {exc}")
|
||||
|
||||
error = None
|
||||
if completed.returncode != 0:
|
||||
error = completed.stderr.strip() or f"command exited with code {completed.returncode}"
|
||||
if return_code != 0:
|
||||
error = _command_failure_detail(
|
||||
stdout,
|
||||
stderr,
|
||||
return_code=return_code,
|
||||
)
|
||||
return StepExecution(
|
||||
exit_code=completed.returncode,
|
||||
stdout=completed.stdout,
|
||||
stderr=completed.stderr,
|
||||
exit_code=return_code,
|
||||
stdout=stdout,
|
||||
stderr=stderr,
|
||||
error=error,
|
||||
)
|
||||
|
||||
|
||||
def _to_text(value: str | bytes | None) -> str:
|
||||
if value is None:
|
||||
return ""
|
||||
if isinstance(value, bytes):
|
||||
return value.decode(errors="replace")
|
||||
return value
|
||||
def _start_process(
|
||||
argv: list[str],
|
||||
*,
|
||||
cwd: Path,
|
||||
env: dict[str, str],
|
||||
stdout: IO[str],
|
||||
stderr: IO[str],
|
||||
) -> tuple[subprocess.Popen[str], _WindowsJob | None]:
|
||||
common = {
|
||||
"cwd": cwd,
|
||||
"env": env,
|
||||
"shell": False,
|
||||
"stdout": stdout,
|
||||
"stderr": stderr,
|
||||
"text": True,
|
||||
"encoding": "utf-8",
|
||||
"errors": "replace",
|
||||
}
|
||||
if os.name != "nt":
|
||||
return subprocess.Popen(argv, start_new_session=True, **common), None
|
||||
|
||||
job = _WindowsJob()
|
||||
process: subprocess.Popen[str] | None = None
|
||||
assigned = False
|
||||
try:
|
||||
process = subprocess.Popen(
|
||||
argv,
|
||||
creationflags=(
|
||||
_CREATE_SUSPENDED | subprocess.CREATE_NEW_PROCESS_GROUP
|
||||
),
|
||||
**common,
|
||||
)
|
||||
job.assign(process.pid)
|
||||
assigned = True
|
||||
psutil.Process(process.pid).resume()
|
||||
return process, job
|
||||
except BaseException:
|
||||
if process is not None:
|
||||
if assigned:
|
||||
try:
|
||||
job.terminate(127)
|
||||
job.wait_empty(_TERMINATION_CONFIRM_SECONDS)
|
||||
except (OSError, RuntimeError):
|
||||
pass
|
||||
else:
|
||||
try:
|
||||
process.kill()
|
||||
except OSError:
|
||||
pass
|
||||
try:
|
||||
process.wait(timeout=_TERMINATION_CONFIRM_SECONDS)
|
||||
except (OSError, subprocess.TimeoutExpired):
|
||||
pass
|
||||
job.close()
|
||||
raise
|
||||
|
||||
|
||||
def start_managed_process(
|
||||
argv: list[str],
|
||||
*,
|
||||
cwd: Path,
|
||||
env: dict[str, str],
|
||||
stdout: IO[str],
|
||||
stderr: IO[str],
|
||||
) -> tuple[subprocess.Popen[str], _WindowsJob | None]:
|
||||
"""Start a process in a scope that can terminate its complete child tree."""
|
||||
|
||||
return _start_process(
|
||||
argv,
|
||||
cwd=cwd,
|
||||
env=env,
|
||||
stdout=stdout,
|
||||
stderr=stderr,
|
||||
)
|
||||
|
||||
|
||||
def _terminate_process_tree(
|
||||
process: subprocess.Popen[str],
|
||||
process_scope: _WindowsJob | None,
|
||||
) -> None:
|
||||
tracked = _process_tree_snapshot(process.pid)
|
||||
if os.name == "nt":
|
||||
if process_scope is None:
|
||||
raise RuntimeError("Windows command process is missing its Job Object")
|
||||
process_scope.terminate(124)
|
||||
process_scope.wait_empty(_TERMINATION_CONFIRM_SECONDS)
|
||||
else:
|
||||
_terminate_posix_process_group(process.pid)
|
||||
|
||||
try:
|
||||
process.wait(timeout=_TERMINATION_CONFIRM_SECONDS)
|
||||
except subprocess.TimeoutExpired as exc:
|
||||
raise RuntimeError("command root process did not terminate") from exc
|
||||
_confirm_processes_stopped(tracked)
|
||||
|
||||
|
||||
def terminate_managed_process_tree(
|
||||
process: subprocess.Popen[str],
|
||||
process_scope: _WindowsJob | None,
|
||||
) -> None:
|
||||
"""Terminate and confirm the complete tree of a managed process."""
|
||||
|
||||
_terminate_process_tree(process, process_scope)
|
||||
|
||||
|
||||
def _terminate_posix_process_group(process_group_id: int) -> None:
|
||||
try:
|
||||
os.killpg(process_group_id, signal.SIGTERM)
|
||||
except ProcessLookupError:
|
||||
return
|
||||
except PermissionError as exc:
|
||||
raise RuntimeError("cannot terminate command process group") from exc
|
||||
|
||||
deadline = time.monotonic() + _TERMINATION_GRACE_SECONDS
|
||||
while _posix_group_has_live_processes(process_group_id):
|
||||
if time.monotonic() >= deadline:
|
||||
break
|
||||
time.sleep(0.01)
|
||||
if _posix_group_has_live_processes(process_group_id):
|
||||
try:
|
||||
os.killpg(process_group_id, signal.SIGKILL)
|
||||
except ProcessLookupError:
|
||||
return
|
||||
except PermissionError as exc:
|
||||
raise RuntimeError("cannot kill command process group") from exc
|
||||
|
||||
deadline = time.monotonic() + _TERMINATION_CONFIRM_SECONDS
|
||||
while _posix_group_has_live_processes(process_group_id):
|
||||
if time.monotonic() >= deadline:
|
||||
raise RuntimeError("command process group did not terminate")
|
||||
time.sleep(0.01)
|
||||
|
||||
|
||||
def _posix_group_has_live_processes(process_group_id: int) -> bool:
|
||||
for candidate in psutil.process_iter(("pid", "status")):
|
||||
try:
|
||||
if (
|
||||
os.getpgid(candidate.pid) == process_group_id
|
||||
and candidate.info["status"] != psutil.STATUS_ZOMBIE
|
||||
):
|
||||
return True
|
||||
except (OSError, psutil.Error):
|
||||
continue
|
||||
return False
|
||||
|
||||
|
||||
def _process_tree_snapshot(pid: int) -> tuple[psutil.Process, ...]:
|
||||
try:
|
||||
root = psutil.Process(pid)
|
||||
return (*root.children(recursive=True), root)
|
||||
except psutil.Error:
|
||||
return ()
|
||||
|
||||
|
||||
def _confirm_processes_stopped(processes: tuple[psutil.Process, ...]) -> None:
|
||||
_, alive = psutil.wait_procs(
|
||||
list(processes),
|
||||
timeout=_TERMINATION_CONFIRM_SECONDS,
|
||||
)
|
||||
live = []
|
||||
for process in alive:
|
||||
try:
|
||||
if process.is_running() and process.status() != psutil.STATUS_ZOMBIE:
|
||||
live.append(process.pid)
|
||||
except psutil.Error:
|
||||
continue
|
||||
if live:
|
||||
raise RuntimeError("command descendants did not terminate")
|
||||
|
||||
|
||||
def _read_output(handle) -> str: # type: ignore[no-untyped-def]
|
||||
handle.flush()
|
||||
handle.seek(0)
|
||||
return handle.read()
|
||||
|
||||
|
||||
def _command_failure_detail(
|
||||
stdout: str,
|
||||
stderr: str,
|
||||
*,
|
||||
return_code: int,
|
||||
) -> str:
|
||||
chunks = [chunk for chunk in (stdout.strip(), stderr.strip()) if chunk]
|
||||
if not chunks:
|
||||
return f"command exited with code {return_code}"
|
||||
return bounded_head_tail(
|
||||
"\n".join(chunks),
|
||||
_FAILURE_OUTPUT_MAX_CHARS,
|
||||
)
|
||||
Reference in new issue
Block a user