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,