48 lines
1.1 KiB
Python
48 lines
1.1 KiB
Python
"""Host-local execution metadata for Cloud planner accounting."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from contextlib import contextmanager
|
|
from contextvars import ContextVar
|
|
from dataclasses import dataclass
|
|
from typing import TYPE_CHECKING
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Iterator
|
|
|
|
from cloud.internal_api.models import AssignmentModel
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class PlannerExecutionContext:
|
|
task_id: str
|
|
attempt: int
|
|
lease_id: str
|
|
|
|
|
|
_context: ContextVar[PlannerExecutionContext | None] = ContextVar(
|
|
"host_agent_planner_execution_context",
|
|
default=None,
|
|
)
|
|
|
|
|
|
def current_planner_execution_context() -> PlannerExecutionContext | None:
|
|
return _context.get()
|
|
|
|
|
|
@contextmanager
|
|
def bind_planner_execution_context(
|
|
assignment: "AssignmentModel",
|
|
) -> "Iterator[None]":
|
|
token = _context.set(
|
|
PlannerExecutionContext(
|
|
task_id=assignment.task_id,
|
|
attempt=assignment.attempt,
|
|
lease_id=assignment.lease_id,
|
|
)
|
|
)
|
|
try:
|
|
yield
|
|
finally:
|
|
_context.reset(token)
|