def _handle_updated_agent_status(self, status_map: Dict[str, str]):
"""
Handle updating the local statuses for agents based on
the previously reported agent statuses.
Takes as input a mapping from agent_id to server-side status
"""
for agent_id, status in status_map.items():
if status not in AgentState.valid():
logger.warning(
f"Invalid status for agent {agent_id}: {status}")
continue
if agent_id not in self.agents:
# no longer tracking agent
continue
agent = self.agents[agent_id].agent
db_status = agent.get_status()
if agent.has_updated_status.is_set():
continue # Incoming info may be stale if we have new info to send
if status == AgentState.STATUS_NONE:
# Stale or reconnect, send a status update
self._send_status_update(self.agents[agent_id])
continue
if status != db_status:
if db_status in AgentState.complete():
logger.info(
f"Got updated status {status} when already final: {agent.db_status}"
)
continue
elif status == AgentState.STATUS_COMPLETED:
continue # COMPLETED can only be marked locally
agent.update_status(status)
pass
def get_agents(self, status: Optional[str] = None) -> List["Agent"]:
"""
Get the list of agents that this worker was responsible for, by the given status
if needed
"""
assert status is None or status in AgentState.valid(), "Invalid agent status"
return self.db.find_agents(worker_id=self.db_id, status=status)
def _handle_updated_agent_status(self, status_map: Dict[str, str]):
"""
Handle updating the local statuses for agents based on
the previously reported agent statuses.
Takes as input a mapping from agent_id to server-side status
"""
for agent_id, status in status_map.items():
if status not in AgentState.valid():
logger.warning(
f"Invalid status for agent {agent_id}: {status}")
continue
if agent_id not in self.agents:
# no longer tracking agent
continue
agent = self.agents[agent_id].agent
db_status = agent.get_status()
if agent.has_updated_status.is_set():
continue # Incoming info may be stale if we have new info to send
if status != db_status:
if status == AgentState.STATUS_COMPLETED:
# Frontend agent completed but hasn't confirmed yet
continue
if status != AgentState.STATUS_DISCONNECT:
# Stale or reconnect, send a status update
self._send_status_update(self.agents[agent_id])
continue # Only DISCONNECT can be marked remotely, rest are mismatch (except STATUS_COMPLETED)
agent.update_status(status)
pass
def handle_updated_agent_status(self, status_map: Dict[str, str]):
"""
Handle updating the local statuses for agents based on
the previously reported agent statuses.
Takes as input a mapping from agent_id to server-side status
"""
live_run = self.get_live_run()
for agent_id, status in status_map.items():
if status not in AgentState.valid():
logger.warning(
f"Invalid status for agent {agent_id}: {status}")
continue
agent = self.get_agent_for_id(agent_id)
if agent is None:
# no longer tracking agent
continue
db_status = agent.get_status()
if status != db_status:
if status == AgentState.STATUS_COMPLETED:
# Frontend agent completed but hasn't confirmed yet
continue
if status != AgentState.STATUS_DISCONNECT:
# Stale or reconnect, send a status update
live_run.loop_wrap.execute_coro(
self.push_status_update(self.agents[agent_id]))
continue # Only DISCONNECT can be marked remotely, rest are mismatch (except STATUS_COMPLETED)
agent.update_status(status)
pass
from mephisto.abstractions.database import MephistoDB
from mephisto.data_model.packet import Packet
from mephisto.data_model.task import Task
from mephisto.data_model.task_run import TaskRun
from mephisto.operations.datatypes import LiveTaskRun
from mephisto.utils.logger_core import get_logger, warn_once
logger = get_logger(name=__name__)
ACTIVE_AGENT_STATUSES = Gauge(
"active_agent_statuses",
"Tracking of all units current statuses",
["status", "agent_type"],
)
for status in AgentState.valid():
ACTIVE_AGENT_STATUSES.labels(status=status, agent_type="main")
ACTIVE_AGENT_STATUSES.labels(status=status, agent_type="onboarding")
ACTIVE_WORKERS = Gauge(
"active_workers",
"Tracking of active workers and how many agents they have",
["worker_id", "agent_type"],
)
# TODO(CLEAN) can probably refactor out some kind of AgentBase
class Agent(MephistoDataModelComponentMixin,
metaclass=MephistoDBBackedABCMeta):
"""
This class encompasses a worker as they are working on an individual assignment.
It maintains details for the current task at hand such as start and end time,
