def setUp(self): self.data_dir = tempfile.mkdtemp() database_path = os.path.join(self.data_dir, "mephisto.db") self.db = LocalMephistoDB(database_path) self.task_id = self.db.new_task("test_mock", MockBlueprint.BLUEPRINT_TYPE) self.task_run_id = get_test_task_run(self.db) self.task_run = TaskRun(self.db, self.task_run_id) architect_config = OmegaConf.structured( MephistoConfig( architect=MockArchitectArgs(should_run_server=True), )) self.architect = MockArchitect(self.db, architect_config, EMPTY_STATE, self.task_run, self.data_dir) self.architect.prepare() self.architect.deploy() self.urls = self.architect._get_socket_urls() # FIXME self.url = self.urls[0] self.provider = MockProvider(self.db) self.provider.setup_resources_for_task_run(self.task_run, self.task_run.args, self.url) self.launcher = TaskLauncher(self.db, self.task_run, self.get_mock_assignment_data_array()) self.launcher.create_assignments() self.launcher.launch_units(self.url) self.sup = None
def get_task_run(self) -> "TaskRun": """Return the TaskRun this agent is working within""" if self._task_run is None: from mephisto.data_model.task import TaskRun self._task_run = TaskRun(self.db, self.task_run_id) return self._task_run
def get_task_run(self) -> TaskRun: """ Return the task run that this assignment is part of """ if self.__task_run is None: self.__task_run = TaskRun(self.db, self.task_run_id) return self.__task_run
def get_task_run(self) -> "TaskRun": """Return the TaskRun this agent is working within""" if self._task_run is None: if self._unit is not None: self._task_run = self._unit.get_task_run() elif self._assignment is not None: self._task_run = self._assignment.get_task_run() else: from mephisto.data_model.task import TaskRun self._task_run = TaskRun(self.db, self.task_run_id) return self._task_run
def new( db: "MephistoDB", task_run: TaskRun, assignment_data: Optional[Dict[str, Any]] ) -> "Assignment": """ Create an assignment for the given task. Initialize the folders for storing the results for this assignment. Can take assignment_data to save and load for this particular assignment. """ # TODO(101) consider offloading this state management to the MephistoDB # as it is data handling and can theoretically be done differently # in different implementations db_id = db.new_assignment( task_run.db_id, task_run.requester_id, task_run.task_type, task_run.provider_type, task_run.sandbox, ) run_dir = task_run.get_run_dir() assign_dir = os.path.join(run_dir, db_id) os.makedirs(assign_dir) if assignment_data is not None: with open( os.path.join(assign_dir, ASSIGNMENT_DATA_FILE), "w+" ) as json_file: json.dump(assignment_data, json_file) return Assignment(db, db_id)
def test_task_run(self) -> None: """Test creation and querying of task_runs""" assert self.db is not None, "No db initialized" db: MephistoDB = self.db task_name, task_id = get_test_task(db) requester_name, requester_id = get_test_requester(db) # Check creation and retrieval of a task_run init_params = json.dumps( OmegaConf.to_yaml(TaskConfig.get_mock_params())) task_run_id = db.new_task_run(task_id, requester_id, init_params, "mock", "mock") self.assertIsNotNone(task_run_id) self.assertTrue(isinstance(task_run_id, str)) task_run_row = db.get_task_run(task_run_id) self.assertEqual(task_run_row["init_params"], init_params) task_run = TaskRun(db, task_run_id) self.assertEqual(task_run.task_id, task_id) # Check finding for task_runs task_runs = db.find_task_runs() self.assertEqual(len(task_runs), 1) self.assertTrue(isinstance(task_runs[0], TaskRun)) self.assertEqual(task_runs[0].db_id, task_run_id) self.assertEqual(task_runs[0].task_id, task_id) self.assertEqual(task_runs[0].requester_id, requester_id) # Check finding for specific task_runs task_runs = db.find_task_runs(task_id=task_id) self.assertEqual(len(task_runs), 1) self.assertTrue(isinstance(task_runs[0], TaskRun)) self.assertEqual(task_runs[0].db_id, task_run_id) self.assertEqual(task_runs[0].task_id, task_id) self.assertEqual(task_runs[0].requester_id, requester_id) task_runs = db.find_task_runs(requester_id=requester_id) self.assertEqual(len(task_runs), 1) self.assertTrue(isinstance(task_runs[0], TaskRun)) self.assertEqual(task_runs[0].db_id, task_run_id) self.assertEqual(task_runs[0].task_id, task_id) self.assertEqual(task_runs[0].requester_id, requester_id) task_runs = db.find_task_runs(task_id=self.get_fake_id("TaskRun")) self.assertEqual(len(task_runs), 0) task_runs = db.find_task_runs(is_completed=True) self.assertEqual(len(task_runs), 0) # Test updating the completion status, requery db.update_task_run(task_run_id, True) task_runs = db.find_task_runs(is_completed=True) self.assertEqual(len(task_runs), 1) self.assertTrue(isinstance(task_runs[0], TaskRun)) self.assertEqual(task_runs[0].db_id, task_run_id)
def get_test_assignment(db: MephistoDB) -> str: """Helper to create an assignment for tests""" task_run_id = get_test_task_run(db) task_run = TaskRun(db, task_run_id) return db.new_assignment( task_run.task_id, task_run_id, task_run.requester_id, task_run.task_type, task_run.provider_type, )
def setUp(self) -> None: """ Setup should put together any requirements for starting the database for a test. """ self.data_dir = tempfile.mkdtemp() self.build_dir = tempfile.mkdtemp() database_path = os.path.join(self.data_dir, "mephisto.db") self.db = LocalMephistoDB(database_path) # TODO(#97) we need to actually pull the task type from the Blueprint self.task_run = TaskRun(self.db, get_test_task_run(self.db)) # TODO(#97) create a mock agent with the given task type? self.TaskRunnerClass = self.BlueprintClass.TaskRunnerClass self.AgentStateClass = self.BlueprintClass.AgentStateClass self.TaskBuilderClass = self.BlueprintClass.TaskBuilderClass
def test_onboarding_agents(self) -> None: """Ensure that the db can create and manipulate onboarding agents""" assert self.db is not None, "No db initialized" db: MephistoDB = self.db task_run_id = get_test_task_run(db) task_run = TaskRun(db, task_run_id) task = task_run.get_task() worker_name, worker_id = get_test_worker(db) onboarding_agent_id = db.new_onboarding_agent(worker_id, task.db_id, task_run_id, "mock") self.assertIsNotNone(onboarding_agent_id) onboarding_agent = OnboardingAgent(db, onboarding_agent_id) self.assertIsInstance(onboarding_agent, OnboardingAgent) found_agents = db.find_onboarding_agents(worker_id=worker_id) self.assertEqual(len(found_agents), 1) self.assertIsInstance(found_agents[0], OnboardingAgent) found_agent = found_agents[0] self.assertEqual(found_agent.db_id, onboarding_agent_id) self.assertEqual(found_agent.get_status(), AgentState.STATUS_NONE)
def get_reviewable_task_runs(): """ Find reviewable task runs by querying for all reviewable tasks and getting their runs """ db = app.extensions["db"] units = db.find_units(status=AssignmentState.COMPLETED) reviewable_count = len(units) task_run_ids = set( [u.get_assignment().get_task_run().db_id for u in units]) task_runs = [TaskRun(db, db_id) for db_id in task_run_ids] dict_tasks = [t.to_dict() for t in task_runs] # TODO(OWN) maybe include warning for auto approve date once that's tracked return jsonify({ "task_runs": dict_tasks, "total_reviewable": reviewable_count })
def test_assignment(self) -> None: """Test creation and querying of assignments""" assert self.db is not None, "No db initialized" db: MephistoDB = self.db task_run_id = get_test_task_run(db) task_run = TaskRun(db, task_run_id) # Check creation and retrieval of an assignment assignment_id = db.new_assignment( task_run.task_id, task_run_id, task_run.requester_id, task_run.task_type, task_run.provider_type, task_run.sandbox, ) self.assertIsNotNone(assignment_id) self.assertTrue(isinstance(assignment_id, str)) assignment_row = db.get_assignment(assignment_id) self.assertEqual(assignment_row["task_run_id"], task_run_id) assignment = Assignment(db, assignment_id) self.assertEqual(assignment.task_run_id, task_run_id) # Check finding for assignments assignments = db.find_assignments() self.assertEqual(len(assignments), 1) self.assertTrue(isinstance(assignments[0], Assignment)) self.assertEqual(assignments[0].db_id, assignment_id) self.assertEqual(assignments[0].task_run_id, task_run_id) # Check finding for specific assignments assignments = db.find_assignments(task_run_id=task_run_id) self.assertEqual(len(assignments), 1) self.assertTrue(isinstance(assignments[0], Assignment)) self.assertEqual(assignments[0].db_id, assignment_id) self.assertEqual(assignments[0].task_run_id, task_run_id) assignments = db.find_assignments( task_run_id=self.get_fake_id("Assignment")) self.assertEqual(len(assignments), 0)
def test_assignment_fails(self) -> None: """Ensure assignments fail to be created or loaded under failure conditions""" assert self.db is not None, "No db initialized" db: MephistoDB = self.db task_run_id = get_test_task_run(db) task_run = TaskRun(db, task_run_id) # Can't create task run with invalid ids with self.assertRaises(EntryDoesNotExistException): assignment_id = db.new_assignment( task_run.task_id, self.get_fake_id("TaskRun"), task_run.requester_id, task_run.task_type, task_run.provider_type, task_run.sandbox, ) # Ensure no assignments were created assignments = db.find_assignments() self.assertEqual(len(assignments), 0)
def setUp(self) -> None: """ Setup should put together any requirements for starting the database for a test. """ try: _ = self.ArchitectClass except: raise unittest.SkipTest("Skipping test as no ArchitectClass set") if not self.warned_about_setup: print( "Architect tests may require using an account with the server provider " "in order to function properly. Make sure these are configured before testing." ) self.warned_about_setup = True self.data_dir = tempfile.mkdtemp() database_path = os.path.join(self.data_dir, "mephisto.db") self.db = LocalMephistoDB(database_path) self.build_dir = tempfile.mkdtemp() self.task_run = TaskRun(self.db, get_test_task_run(self.db)) builder = MockTaskBuilder(self.task_run, {}) builder.build_in_dir(self.build_dir)
class OnboardingAgent(ABC): """ Onboarding agents are a special extension of agents used in tasks that have a separate onboarding step. These agents are designed to work without being linked to an explicit unit, and instead are tied to the task run and task name. Blueprints that require OnboardingAgents should implement an OnboardingAgentState (to process the special task), and their TaskRunners should have a run_onboarding and cleanup_onboarding method. """ DISPLAY_PREFIX = "onboarding_" def __init__(self, db: "MephistoDB", db_id: str, row: Optional[Mapping[str, Any]] = None): self.db: "MephistoDB" = db if row is None: row = db.get_onboarding_agent(db_id) assert row is not None, f"Given db_id {db_id} did not exist in given db" self.db_id: str = row["onboarding_agent_id"] self.db_status = row["status"] self.worker_id = row["worker_id"] self.task_type = row["task_type"] self.pending_observations: List["Packet"] = [] self.pending_actions: List["Packet"] = [] self.has_action = threading.Event() self.has_action.clear() self.wants_action = threading.Event() self.wants_action.clear() self.has_updated_status = threading.Event() self.task_run_id = row["task_run_id"] self.task_id = row["task_id"] self.did_submit = threading.Event() # Deferred loading of related entities self._worker: Optional["Worker"] = None self._task_run: Optional["TaskRun"] = None self._task: Optional["Task"] = None # Follow-up initialization self.state = AgentState(self) # type: ignore def get_agent_id(self) -> str: """Return an id to use for onboarding agent requests""" return f"{self.DISPLAY_PREFIX}{self.db_id}" @classmethod def is_onboarding_id(cls, agent_id: str) -> bool: """return if the given id is for an onboarding agent""" return agent_id.startswith(cls.DISPLAY_PREFIX) @classmethod def get_db_id_from_agent_id(cls, agent_id: str) -> str: """Extract the db_id for an onboarding_agent""" assert agent_id.startswith( cls.DISPLAY_PREFIX ), f"Provided id {agent_id} is not an onboarding_id" return agent_id[len(cls.DISPLAY_PREFIX):] def get_worker(self) -> Worker: """ Return the worker that is using this agent for a task """ if self._worker is None: self._worker = Worker(self.db, self.worker_id) return self._worker def get_task_run(self) -> "TaskRun": """Return the TaskRun this agent is working within""" if self._task_run is None: from mephisto.data_model.task import TaskRun self._task_run = TaskRun(self.db, self.task_run_id) return self._task_run def get_task(self) -> "Task": """Return the Task this agent is working within""" if self._task is None: if self._task_run is not None: self._task = self._task_run.get_task() else: from mephisto.data_model.task import Task self._task = Task(self.db, self.task_id) return self._task def get_data_dir(self) -> str: """ Return the directory to be storing any agent state for this agent into """ task_run_dir = self.get_task_run().get_run_dir() return os.path.join(task_run_dir, "onboarding", self.get_agent_id()) def update_status(self, new_status: str) -> None: """Update the database status of this agent, and possibly send a message to the frontend agent informing them of this update""" if self.db_status == new_status: return # Noop, this is already the case if self.db_status in AgentState.complete(): print(f"Updating a final status, was {self.db_status} " f"and want to set to {new_status}") self.db.update_onboarding_agent(self.db_id, status=new_status) self.db_status = new_status self.has_updated_status.set() if new_status in [ AgentState.STATUS_RETURNED, AgentState.STATUS_DISCONNECT ]: # Disconnect statuses should free any pending acts self.has_action.set() self.did_submit.set() def observe(self, packet: "Packet") -> None: """ Pass the observed information to the AgentState, then queue the information to be pushed to the user """ sending_packet = packet.copy() sending_packet.receiver_id = self.get_agent_id() self.state.update_data(sending_packet) self.pending_observations.append(sending_packet) def act(self, timeout: Optional[int] = None) -> Optional["Packet"]: """ Request information from the Agent's frontend. If non-blocking, (timeout is None) should return None if no actions are ready to be returned. """ if len(self.pending_actions) == 0: self.wants_action.set() if timeout is None or timeout == 0: return None self.has_action.wait(timeout) if len(self.pending_actions) == 0: # various disconnect cases status = self.get_status() if status == AgentState.STATUS_DISCONNECT: raise AgentDisconnectedError(self.db_id) elif status == AgentState.STATUS_RETURNED: raise AgentReturnedError(self.db_id) self.update_status(AgentState.STATUS_TIMEOUT) raise AgentTimeoutError(timeout, self.db_id) assert len( self.pending_actions) > 0, "has_action released without an action!" act = self.pending_actions.pop(0) if "MEPHISTO_is_submit" in act.data and act.data["MEPHISTO_is_submit"]: self.did_submit.set() if len(self.pending_actions) == 0: self.has_action.clear() self.state.update_data(act) return act def get_status(self) -> str: """Get the status of this agent in their work on their unit""" if self.db_status not in AgentState.complete(): row = self.db.get_onboarding_agent(self.db_id) if row["status"] != self.db_status: if row["status"] in [ AgentState.STATUS_RETURNED, AgentState.STATUS_DISCONNECT, ]: # Disconnect statuses should free any pending acts self.has_action.set() self.has_updated_status.set() self.db_status = row["status"] return self.db_status def mark_done(self) -> None: """Mark this agent as done by setting the status to a terminal onboarding state""" # TODO the logic for when onboarding gets marked as waiting or approved/rejected # should likely be cleaned up to remove these conditionals. if self.get_status not in [ AgentState.STATUS_APPROVED, AgentState.STATUS_REJECTED, ]: self.update_status(AgentState.STATUS_WAITING) @staticmethod def new(db: "MephistoDB", worker: Worker, task_run: "TaskRun") -> "OnboardingAgent": """ Create an OnboardingAgent for a worker to use as part of a task run """ db_id = db.new_onboarding_agent(worker.db_id, task_run.task_id, task_run.db_id, task_run.task_type) return OnboardingAgent(db, db_id)
class Assignment: """ This class tracks an individual run of a specific task, and handles state management for the set of units within via abstracted database helpers """ def __init__( self, db: "MephistoDB", db_id: str, row: Optional[Mapping[str, Any]] = None ): self.db: "MephistoDB" = db if row is None: row = db.get_assignment(db_id) assert row is not None, f"Given db_id {db_id} did not exist in given db" self.db_id: str = row["assignment_id"] self.task_run_id = row["task_run_id"] self.sandbox = row["sandbox"] self.task_id = row["task_id"] self.requester_id = row["requester_id"] self.task_type = row["task_type"] self.provider_type = row["provider_type"] # Deferred loading of related entities self.__task_run: Optional["TaskRun"] = None self.__task: Optional["Task"] = None self.__requester: Optional["Requester"] = None def get_data_dir(self) -> str: """Return the directory we expect to find assignment data in""" task_run = self.get_task_run() run_dir = task_run.get_run_dir() return os.path.join(run_dir, self.db_id) def get_assignment_data(self) -> InitializationData: """Return the specific assignment data for this assignment""" assign_data_filename = os.path.join(self.get_data_dir(), ASSIGNMENT_DATA_FILE) assert os.path.exists(assign_data_filename), "No data exists for assignment" with open(assign_data_filename, "r") as json_file: return InitializationData.loadFromJSON(json_file) def write_assignment_data(self, data: InitializationData) -> None: """Set the assignment data for this assignment""" assign_data_filename = os.path.join(self.get_data_dir(), ASSIGNMENT_DATA_FILE) os.makedirs(self.get_data_dir(), exist_ok=True) with open(assign_data_filename, "w+") as json_file: data.dumpJSON(json_file) def get_agents(self) -> List[Optional["Agent"]]: """ Return all of the agents for this assignment """ units = self.get_units() return [u.get_assigned_agent() for u in units] def get_status(self) -> str: """ Get the status of this assignment, as determined by the status of the units """ units = self.get_units() statuses = set(unit.get_status() for unit in units) if len(statuses) == 1: return statuses.pop() if len(statuses) == 0: return AssignmentState.CREATED if AssignmentState.CREATED in statuses: # TODO(#99) handle the case where new units are created after # everything else is launched return AssignmentState.CREATED if any([s == AssignmentState.LAUNCHED for s in statuses]): # If any are only launched, consider the whole thing launched return AssignmentState.LAUNCHED if any([s == AssignmentState.ASSIGNED for s in statuses]): # If any are still assigned, consider the whole thing assigned return AssignmentState.ASSIGNED if all( [ s in [AssignmentState.ACCEPTED, AssignmentState.REJECTED] for s in statuses ] ): return AssignmentState.MIXED if all([s in AssignmentState.final_agent() for s in statuses]): return AssignmentState.COMPLETED raise NotImplementedError(f"Unexpected set of unit statuses {statuses}") def get_task_run(self) -> TaskRun: """ Return the task run that this assignment is part of """ if self.__task_run is None: self.__task_run = TaskRun(self.db, self.task_run_id) return self.__task_run def get_task(self) -> Task: """ Return the task run that this assignment is part of """ if self.__task is None: if self.__task_run is not None: self.__task = self.__task_run.get_task() else: self.__task = Task(self.db, self.task_id) return self.__task def get_requester(self) -> Requester: """ Return the requester who offered this Assignment """ if self.__requester is None: if self.__task_run is not None: self.__requester = self.__task_run.get_requester() else: self.__requester = Requester(self.db, self.requester_id) return self.__requester def get_units(self, status: Optional[str] = None) -> List["Unit"]: """ Get units for this assignment, optionally constrained by the specific status. """ assert ( status is None or status in AssignmentState.valid_unit() ), "Invalid assignment status" units = self.db.find_units(assignment_id=self.db_id) if status is not None: units = [u for u in units if u.get_status() == status] return units def get_workers(self) -> List["Worker"]: """ Get the list of workers that have worked on this specific assignment """ units = self.get_units() pos_agents = [s.get_assigned_agent() for s in units] agents = [a for a in pos_agents if a is not None] workers = [a.get_worker() for a in agents] return workers def get_cost_of_statuses(self, statuses: List[str]) -> float: """ Return the sum of all pay_amounts for every unit of this assignment with any of the given statuses """ units = [u for u in self.get_units() if u.get_status() in statuses] sum_cost = 0.0 for unit in units: sum_cost += unit.get_pay_amount() return sum_cost # TODO(100) add helpers to manage retrieving results as well @staticmethod def new( db: "MephistoDB", task_run: TaskRun, assignment_data: Optional[Dict[str, Any]] ) -> "Assignment": """ Create an assignment for the given task. Initialize the folders for storing the results for this assignment. Can take assignment_data to save and load for this particular assignment. """ # TODO(101) consider offloading this state management to the MephistoDB # as it is data handling and can theoretically be done differently # in different implementations db_id = db.new_assignment( task_run.db_id, task_run.requester_id, task_run.task_type, task_run.provider_type, task_run.sandbox, ) run_dir = task_run.get_run_dir() assign_dir = os.path.join(run_dir, db_id) os.makedirs(assign_dir) if assignment_data is not None: with open( os.path.join(assign_dir, ASSIGNMENT_DATA_FILE), "w+" ) as json_file: json.dump(assignment_data, json_file) return Assignment(db, db_id)
class Unit(ABC): """ This class tracks the status of an individual worker's contribution to a higher level assignment. It is the smallest 'unit' of work to complete the assignment, and this class is only responsible for checking the status of that work itself being done. It should be extended for usage with a specific crowd provider """ def __init__( self, db: "MephistoDB", db_id: str, row: Optional[Mapping[str, Any]] = None ): self.db: "MephistoDB" = db if row is None: row = db.get_unit(db_id) assert row is not None, f"Given db_id {db_id} did not exist in given db" self.db_id: str = row["unit_id"] self.assignment_id = row["assignment_id"] self.unit_index = row["unit_index"] self.pay_amount = row["pay_amount"] self.agent_id = row["agent_id"] self.provider_type = row["provider_type"] self.db_status = row["status"] self.task_type = row["task_type"] self.task_id = row["task_id"] self.task_run_id = row["task_run_id"] self.sandbox = row["sandbox"] self.requester_id = row["requester_id"] self.worker_id = row["worker_id"] # Deferred loading of related entities self.__task: Optional["Task"] = None self.__task_run: Optional["TaskRun"] = None self.__assignment: Optional["Assignment"] = None self.__requester: Optional["Requester"] = None self.__agent: Optional["Agent"] = None self.__worker: Optional["Worker"] = None def __new__( cls, db: "MephistoDB", db_id: str, row: Optional[Mapping[str, Any]] = None ) -> "Unit": """ The new method is overridden to be able to automatically generate the expected Unit class without needing to specifically find it for a given db_id. As such it is impossible to create a Unit as you will instead be returned the correct Unit class according to the crowdprovider associated with this Unit. """ if cls == Unit: # We are trying to construct a Unit, find what type to use and # create that instead from mephisto.core.registry import get_crowd_provider_from_type if row is None: row = db.get_unit(db_id) assert row is not None, f"Given db_id {db_id} did not exist in given db" correct_class = get_crowd_provider_from_type(row["provider_type"]).UnitClass return super().__new__(correct_class) else: # We are constructing another instance directly return super().__new__(cls) def get_crowd_provider_class(self) -> Type["CrowdProvider"]: """Get the CrowdProvider class that manages this Unit""" from mephisto.core.registry import get_crowd_provider_from_type return get_crowd_provider_from_type(self.provider_type) def get_assignment_data(self) -> Optional[Dict[str, Any]]: """Return the specific assignment data for this assignment""" return self.get_assignment().get_assignment_data() def sync_status(self) -> None: """ Ensure that the queried status from this unit and the db status are up to date """ # TODO(102) this will need to be run periodically/on crashes # to sync any lost state self.set_db_status(self.get_status()) def get_db_status(self) -> str: """ Return the status as currently stored in the database """ if self.db_status in AssignmentState.final_unit(): return self.db_status row = self.db.get_unit(self.db_id) assert row is not None, f"Unit {self.db_id} stopped existing in the db..." return row["status"] def set_db_status(self, status: str) -> None: """ Set the status reflected in the database for this Unit """ assert ( status in AssignmentState.valid_unit() ), f"{status} not valid Assignment Status, not in {AssignmentState.valid_unit()}" self.db_status = status self.db.update_unit(self.db_id, status=status) def get_assignment(self) -> Assignment: """ Return the assignment that this Unit is part of. """ if self.__assignment is None: self.__assignment = Assignment(self.db, self.assignment_id) return self.__assignment def get_task_run(self) -> TaskRun: """ Return the task run that this assignment is part of """ if self.__task_run is None: if self.__assignment is not None: self.__task_run = self.__assignment.get_task_run() else: self.__task_run = TaskRun(self.db, self.task_run_id) return self.__task_run def get_task(self) -> Task: """ Return the task that this assignment is part of """ if self.__task is None: if self.__assignment is not None: self.__task = self.__assignment.get_task() elif self.__task_run is not None: self.__task = self.__task_run.get_task() else: self.__task = Task(self.db, self.task_id) return self.__task def get_requester(self) -> "Requester": """ Return the requester who offered this Unit """ if self.__requester is None: if self.__assignment is not None: self.__requester = self.__assignment.get_requester() elif self.__task_run is not None: self.__requester = self.__task_run.get_requester() else: self.__requester = Requester(self.db, self.requester_id) return self.__requester def clear_assigned_agent(self) -> None: """Clear the agent that is assigned to this unit""" self.db.clear_unit_agent_assignment(self.db_id) self.agent_id = None self.__agent = None def get_assigned_agent(self) -> Optional[Agent]: """ Get the agent assigned to this Unit if there is one, else return None """ # In these statuses, we know the agent isn't changing anymore, and thus will # not need to be re-queried # TODO(#97) add test to ensure this behavior/assumption holds always if self.db_status in AssignmentState.final_unit(): if self.agent_id is None: return None return Agent(self.db, self.agent_id) # Query the database to get the most up-to-date assignment, as this can # change after instantiation if the Unit status isn't final # TODO(#101) this may not be particularly efficient row = self.db.get_unit(self.db_id) assert row is not None, f"Unit {self.db_id} stopped existing in the db..." agent_id = row["agent_id"] if agent_id is not None: return Agent(self.db, agent_id) return None @staticmethod def _register_unit( db: "MephistoDB", assignment: Assignment, index: int, pay_amount: float, provider_type: str, ) -> "Unit": """ Create an entry for this unit in the database """ db_id = db.new_unit( assignment.task_id, assignment.task_run_id, assignment.requester_id, assignment.db_id, index, pay_amount, provider_type, assignment.task_type, ) return Unit(db, db_id) def get_pay_amount(self) -> float: """ Return the amount that this Unit is costing against the budget, calculating additional fees as relevant """ return self.pay_amount # Children classes may need to override the following def get_status(self) -> str: """ Get the status of this unit, as determined by whether there's a worker working on it at the moment, and any other possible states. Should return one of UNIT_STATUSES Accurate status is crowd-provider dependent, and thus this method should be defined in the child class to ensure that the local record matches the ground truth in the provider """ from mephisto.data_model.blueprint import AgentState db_status = self.db_status computed_status = AssignmentState.LAUNCHED agent = self.get_assigned_agent() if agent is None: row = self.db.get_unit(self.db_id) computed_status = row["status"] else: agent_status = agent.get_status() if agent_status == AgentState.STATUS_NONE: computed_status = AssignmentState.LAUNCHED elif agent_status in [ AgentState.STATUS_ACCEPTED, AgentState.STATUS_ONBOARDING, AgentState.STATUS_PARTNER_DISCONNECT, AgentState.STATUS_WAITING, AgentState.STATUS_IN_TASK, ]: computed_status = AssignmentState.ASSIGNED elif agent_status in [AgentState.STATUS_COMPLETED]: computed_status = AssignmentState.COMPLETED elif agent_status in [AgentState.STATUS_SOFT_REJECTED]: computed_status = AssignmentState.SOFT_REJECTED elif agent_status in [AgentState.STATUS_EXPIRED]: computed_status = AssignmentState.EXPIRED elif agent_status in [ AgentState.STATUS_DISCONNECT, AgentState.STATUS_RETURNED, ]: computed_status = AssignmentState.ASSIGNED elif agent_status == AgentState.STATUS_APPROVED: computed_status = AssignmentState.ACCEPTED elif agent_status == AgentState.STATUS_REJECTED: computed_status = AssignmentState.REJECTED if computed_status != db_status: self.set_db_status(computed_status) return computed_status # Children classes should implement the below methods def launch(self, task_url: str) -> None: """ Make this Unit available on the crowdsourcing vendor. Depending on the task type, this could mean a number of different setup steps. Some crowd providers require setting up a configuration for the very first launch, and this method should call a helper to manage that step if necessary. """ raise NotImplementedError() def expire(self) -> float: """ Expire this unit, removing it from being workable on the vendor. Return the maximum time needed to wait before we know it's taken down. """ raise NotImplementedError() def is_expired(self) -> bool: """Determine if this unit is expired as according to the vendor.""" raise NotImplementedError() @staticmethod def new( db: "MephistoDB", assignment: Assignment, index: int, pay_amount: float ) -> "Unit": """ Create a Unit for the given assignment Implementation should return the result of _register_unit when sure the unit can be successfully created to have it put into the db. """ raise NotImplementedError()
class TestSupervisor(unittest.TestCase): """ Unit testing for the Mephisto Supervisor, uses WebsocketChannel and MockArchitect """ def setUp(self): self.data_dir = tempfile.mkdtemp() database_path = os.path.join(self.data_dir, "mephisto.db") self.db = LocalMephistoDB(database_path) self.task_id = self.db.new_task("test_mock", MockBlueprint.BLUEPRINT_TYPE) self.task_run_id = get_test_task_run(self.db) self.task_run = TaskRun(self.db, self.task_run_id) architect_config = OmegaConf.structured( MephistoConfig( architect=MockArchitectArgs(should_run_server=True), )) self.architect = MockArchitect(self.db, architect_config, EMPTY_STATE, self.task_run, self.data_dir) self.architect.prepare() self.architect.deploy() self.urls = self.architect._get_socket_urls() # FIXME self.url = self.urls[0] self.provider = MockProvider(self.db) self.provider.setup_resources_for_task_run(self.task_run, self.task_run.args, self.url) self.launcher = TaskLauncher(self.db, self.task_run, self.get_mock_assignment_data_array()) self.launcher.create_assignments() self.launcher.launch_units(self.url) self.sup = None def tearDown(self): if self.sup is not None: self.sup.shutdown() self.launcher.expire_units() self.architect.cleanup() self.architect.shutdown() self.db.shutdown() shutil.rmtree(self.data_dir) def get_mock_assignment_data_array(self) -> List[InitializationData]: mock_data = MockTaskRunner.get_mock_assignment_data() return [mock_data, mock_data] def test_initialize_supervisor(self): """Ensure that the supervisor object can even be created""" sup = Supervisor(self.db) self.assertIsNotNone(sup) self.assertDictEqual(sup.agents, {}) self.assertDictEqual(sup.channels, {}) sup.shutdown() def test_channel_operations(self): """ Initialize a channel, and ensure the basic startup and shutdown functions are working """ sup = Supervisor(self.db) self.sup = sup TaskRunnerClass = MockBlueprint.TaskRunnerClass args = MockBlueprint.ArgsClass() config = OmegaConf.structured(MephistoConfig(blueprint=args)) task_runner = TaskRunnerClass(self.task_run, config, EMPTY_STATE) test_job = Job( architect=self.architect, task_runner=task_runner, provider=self.provider, qualifications=[], registered_channel_ids=[], ) channels = self.architect.get_channels(sup._on_channel_open, sup._on_catastrophic_disconnect, sup._on_message) channel = channels[0] channel.open() channel_id = channel.channel_id self.assertIsNotNone(channel_id) channel.close() self.assertTrue(channel.is_closed()) def test_register_concurrent_job(self): """Test registering and running a job that requires multiple workers""" # Handle baseline setup sup = Supervisor(self.db) self.sup = sup TaskRunnerClass = MockBlueprint.TaskRunnerClass args = MockBlueprint.ArgsClass() args.timeout_time = 5 args.is_concurrent = False config = OmegaConf.structured(MephistoConfig(blueprint=args)) task_runner = TaskRunnerClass(self.task_run, config, EMPTY_STATE) sup.register_job(self.architect, task_runner, self.provider) self.assertEqual(len(sup.channels), 1) channel_info = list(sup.channels.values())[0] self.assertIsNotNone(channel_info) self.assertTrue(channel_info.channel.is_alive) channel_id = channel_info.channel_id task_runner = channel_info.job.task_runner self.assertIsNotNone(channel_id) self.assertEqual( len(self.architect.server.subs), 1, "MockServer doesn't see registered channel", ) self.assertIsNotNone( self.architect.server.last_alive_packet, "No alive packet received by server", ) sup.launch_sending_thread() self.assertIsNotNone(sup.sending_thread) # Register a worker mock_worker_name = "MOCK_WORKER" self.architect.server.register_mock_worker(mock_worker_name) workers = self.db.find_workers(worker_name=mock_worker_name) self.assertEqual(len(workers), 1, "Worker not successfully registered") worker = workers[0] self.architect.server.register_mock_worker(mock_worker_name) workers = self.db.find_workers(worker_name=mock_worker_name) self.assertEqual(len(workers), 1, "Worker potentially re-registered") worker_id = workers[0].db_id self.assertEqual(len(task_runner.running_assignments), 0) # Register an agent mock_agent_details = "FAKE_ASSIGNMENT" self.architect.server.register_mock_agent(worker_id, mock_agent_details) agents = self.db.find_agents() self.assertEqual(len(agents), 1, "Agent was not created properly") self.architect.server.register_mock_agent(worker_id, mock_agent_details) agents = self.db.find_agents() self.assertEqual(len(agents), 1, "Agent may have been duplicated") agent = agents[0] self.assertIsNotNone(agent) self.assertEqual(len(sup.agents), 1, "Agent not registered with supervisor") self.assertEqual(len(task_runner.running_units), 1, "Ready task was not launched") # Register another worker mock_worker_name = "MOCK_WORKER_2" self.architect.server.register_mock_worker(mock_worker_name) workers = self.db.find_workers(worker_name=mock_worker_name) worker_id = workers[0].db_id # Register an agent mock_agent_details = "FAKE_ASSIGNMENT_2" self.architect.server.register_mock_agent(worker_id, mock_agent_details) self.assertEqual(len(task_runner.running_units), 2, "Tasks were not launched") agents = [a.agent for a in sup.agents.values()] # Make both agents act agent_id_1, agent_id_2 = agents[0].db_id, agents[1].db_id agent_1_data = agents[0].datastore.agent_data[agent_id_1] agent_2_data = agents[1].datastore.agent_data[agent_id_2] self.architect.server.send_agent_act(agent_id_1, {"text": "message1"}) self.architect.server.send_agent_act(agent_id_2, {"text": "message2"}) # Give up to 1 seconds for the actual operations to occur start_time = time.time() TIMEOUT_TIME = 1 while time.time() - start_time < TIMEOUT_TIME: if len(agent_1_data["acts"]) > 0: break time.sleep(0.1) self.assertLess(time.time() - start_time, TIMEOUT_TIME, "Did not process messages in time") # Give up to 1 seconds for the task to complete afterwards start_time = time.time() TIMEOUT_TIME = 1 while time.time() - start_time < TIMEOUT_TIME: if len(task_runner.running_units) == 0: break time.sleep(0.1) self.assertLess(time.time() - start_time, TIMEOUT_TIME, "Did not complete task in time") # Give up to 1 seconds for all messages to propogate start_time = time.time() TIMEOUT_TIME = 1 while time.time() - start_time < TIMEOUT_TIME: if self.architect.server.actions_observed == 2: break time.sleep(0.1) self.assertLess(time.time() - start_time, TIMEOUT_TIME, "Not all actions observed in time") sup.shutdown() self.assertTrue(channel_info.channel.is_closed) def test_register_job(self): """Test registering and running a job run asynchronously""" # Handle baseline setup sup = Supervisor(self.db) self.sup = sup TaskRunnerClass = MockBlueprint.TaskRunnerClass args = MockBlueprint.ArgsClass() args.timeout_time = 5 config = OmegaConf.structured(MephistoConfig(blueprint=args)) task_runner = TaskRunnerClass(self.task_run, config, EMPTY_STATE) sup.register_job(self.architect, task_runner, self.provider) self.assertEqual(len(sup.channels), 1) channel_info = list(sup.channels.values())[0] self.assertIsNotNone(channel_info) self.assertTrue(channel_info.channel.is_alive()) channel_id = channel_info.channel_id task_runner = channel_info.job.task_runner self.assertIsNotNone(channel_id) self.assertEqual( len(self.architect.server.subs), 1, "MockServer doesn't see registered channel", ) self.assertIsNotNone( self.architect.server.last_alive_packet, "No alive packet received by server", ) sup.launch_sending_thread() self.assertIsNotNone(sup.sending_thread) # Register a worker mock_worker_name = "MOCK_WORKER" self.architect.server.register_mock_worker(mock_worker_name) workers = self.db.find_workers(worker_name=mock_worker_name) self.assertEqual(len(workers), 1, "Worker not successfully registered") worker = workers[0] self.architect.server.register_mock_worker(mock_worker_name) workers = self.db.find_workers(worker_name=mock_worker_name) self.assertEqual(len(workers), 1, "Worker potentially re-registered") worker_id = workers[0].db_id self.assertEqual(len(task_runner.running_assignments), 0) # Register an agent mock_agent_details = "FAKE_ASSIGNMENT" self.architect.server.register_mock_agent(worker_id, mock_agent_details) agents = self.db.find_agents() self.assertEqual(len(agents), 1, "Agent was not created properly") self.architect.server.register_mock_agent(worker_id, mock_agent_details) agents = self.db.find_agents() self.assertEqual(len(agents), 1, "Agent may have been duplicated") agent = agents[0] self.assertIsNotNone(agent) self.assertEqual(len(sup.agents), 1, "Agent not registered with supervisor") self.assertEqual(len(task_runner.running_assignments), 0, "Task was not yet ready") # Register another worker mock_worker_name = "MOCK_WORKER_2" self.architect.server.register_mock_worker(mock_worker_name) workers = self.db.find_workers(worker_name=mock_worker_name) worker_id = workers[0].db_id # Register an agent mock_agent_details = "FAKE_ASSIGNMENT_2" self.architect.server.register_mock_agent(worker_id, mock_agent_details) self.assertEqual(len(task_runner.running_assignments), 1, "Task was not launched") agents = [a.agent for a in sup.agents.values()] # Make both agents act agent_id_1, agent_id_2 = agents[0].db_id, agents[1].db_id agent_1_data = agents[0].datastore.agent_data[agent_id_1] agent_2_data = agents[1].datastore.agent_data[agent_id_2] self.architect.server.send_agent_act(agent_id_1, {"text": "message1"}) self.architect.server.send_agent_act(agent_id_2, {"text": "message2"}) # Give up to 1 seconds for the actual operation to occur start_time = time.time() TIMEOUT_TIME = 1 while time.time() - start_time < TIMEOUT_TIME: if len(agent_1_data["acts"]) > 0: break time.sleep(0.1) self.assertLess(time.time() - start_time, TIMEOUT_TIME, "Did not process messages in time") # Give up to 1 seconds for the task to complete afterwards start_time = time.time() TIMEOUT_TIME = 1 while time.time() - start_time < TIMEOUT_TIME: if len(task_runner.running_assignments) == 0: break time.sleep(0.1) self.assertLess(time.time() - start_time, TIMEOUT_TIME, "Did not complete task in time") # Give up to 1 seconds for all messages to propogate start_time = time.time() TIMEOUT_TIME = 1 while time.time() - start_time < TIMEOUT_TIME: if self.architect.server.actions_observed == 2: break time.sleep(0.1) self.assertLess(time.time() - start_time, TIMEOUT_TIME, "Not all actions observed in time") sup.shutdown() self.assertTrue(channel_info.channel.is_closed()) def test_register_concurrent_job_with_onboarding(self): """Test registering and running a job with onboarding""" # Handle baseline setup sup = Supervisor(self.db) self.sup = sup TEST_QUALIFICATION_NAME = "test_onboarding_qualification" task_run_args = self.task_run.args task_run_args.blueprint.use_onboarding = True task_run_args.blueprint.onboarding_qualification = TEST_QUALIFICATION_NAME task_run_args.blueprint.timeout_time = 5 task_run_args.blueprint.is_concurrent = True self.task_run.get_task_config() # Supervisor expects that blueprint setup has already occurred blueprint = self.task_run.get_blueprint() TaskRunnerClass = MockBlueprint.TaskRunnerClass task_runner = TaskRunnerClass(self.task_run, task_run_args, EMPTY_STATE) sup.register_job(self.architect, task_runner, self.provider) self.assertEqual(len(sup.channels), 1) channel_info = list(sup.channels.values())[0] self.assertIsNotNone(channel_info) self.assertTrue(channel_info.channel.is_alive()) channel_id = channel_info.channel_id task_runner = channel_info.job.task_runner self.assertIsNotNone(channel_id) self.assertEqual( len(self.architect.server.subs), 1, "MockServer doesn't see registered channel", ) self.assertIsNotNone( self.architect.server.last_alive_packet, "No alive packet received by server", ) sup.launch_sending_thread() self.assertIsNotNone(sup.sending_thread) self.assertEqual(len(task_runner.running_units), 0) # Fail to register an agent who fails onboarding mock_worker_name = "BAD_WORKER" self.architect.server.register_mock_worker(mock_worker_name) workers = self.db.find_workers(worker_name=mock_worker_name) self.assertEqual(len(workers), 1, "Worker not successfully registered") worker_0 = workers[0] self.architect.server.register_mock_worker(mock_worker_name) workers = self.db.find_workers(worker_name=mock_worker_name) self.assertEqual(len(workers), 1, "Worker potentially re-registered") worker_id = workers[0].db_id mock_agent_details = "FAKE_ASSIGNMENT" self.architect.server.register_mock_agent(worker_id, mock_agent_details) agents = self.db.find_agents() self.assertEqual(len(agents), 0, "Agent should not be created yet - need onboarding") onboard_agents = self.db.find_onboarding_agents() self.assertEqual(len(onboard_agents), 1, "Onboarding agent should have been created") time.sleep(0.1) last_packet = self.architect.server.last_packet self.assertIsNotNone(last_packet) self.assertIn("onboard_data", last_packet["data"], "Onboarding not triggered") self.architect.server.last_packet = None # Submit onboarding from the agent onboard_data = {"should_pass": False} self.architect.server.register_mock_agent_after_onboarding( worker_id, onboard_agents[0].get_agent_id(), onboard_data) agents = self.db.find_agents() self.assertEqual(len(agents), 0, "Failed agent created after onboarding") # Re-register as if refreshing self.architect.server.register_mock_agent(worker_id, mock_agent_details) agents = self.db.find_agents() self.assertEqual(len(agents), 0, "Failed agent created after onboarding") self.assertEqual(len(sup.agents), 0, "Failed agent registered with supervisor") self.assertEqual( len(task_runner.running_units), 0, "Task should not launch with failed worker", ) # Register a worker mock_worker_name = "MOCK_WORKER" self.architect.server.register_mock_worker(mock_worker_name) workers = self.db.find_workers(worker_name=mock_worker_name) self.assertEqual(len(workers), 1, "Worker not successfully registered") worker_1 = workers[0] self.architect.server.register_mock_worker(mock_worker_name) workers = self.db.find_workers(worker_name=mock_worker_name) self.assertEqual(len(workers), 1, "Worker potentially re-registered") worker_id = workers[0].db_id self.assertEqual(len(task_runner.running_assignments), 0) # Fail to register a blocked agent mock_agent_details = "FAKE_ASSIGNMENT" qualification_id = blueprint.onboarding_qualification_id self.db.grant_qualification(qualification_id, worker_1.db_id, 0) self.architect.server.register_mock_agent(worker_id, mock_agent_details) agents = self.db.find_agents() self.assertEqual(len(agents), 0, "Agent should not be created yet, failed onboarding") time.sleep(0.1) last_packet = self.architect.server.last_packet self.assertIsNotNone(last_packet) self.assertNotIn( "onboard_data", last_packet["data"], "Onboarding triggered for disqualified worker", ) self.assertIsNone(last_packet["data"]["agent_id"], "worker assigned real agent id") self.architect.server.last_packet = None self.db.revoke_qualification(qualification_id, worker_id) # Register an onboarding agent successfully mock_agent_details = "FAKE_ASSIGNMENT" self.architect.server.register_mock_agent(worker_id, mock_agent_details) agents = self.db.find_agents() self.assertEqual(len(agents), 0, "Agent should not be created yet - need onboarding") onboard_agents = self.db.find_onboarding_agents() self.assertEqual(len(onboard_agents), 2, "Onboarding agent should have been created") time.sleep(0.1) last_packet = self.architect.server.last_packet self.assertIsNotNone(last_packet) self.assertIn("onboard_data", last_packet["data"], "Onboarding not triggered") self.architect.server.last_packet = None # Submit onboarding from the agent onboard_data = {"should_pass": True} self.architect.server.register_mock_agent_after_onboarding( worker_id, onboard_agents[1].get_agent_id(), onboard_data) agents = self.db.find_agents() self.assertEqual(len(agents), 1, "Agent not created after onboarding") # Re-register as if refreshing self.architect.server.register_mock_agent(worker_id, mock_agent_details) agents = self.db.find_agents() self.assertEqual(len(agents), 1, "Agent may have been duplicated") agent = agents[0] self.assertIsNotNone(agent) self.assertEqual(len(sup.agents), 1, "Agent not registered with supervisor") self.assertEqual( len(task_runner.running_assignments), 0, "Task was not yet ready, should not launch", ) # Register another worker mock_worker_name = "MOCK_WORKER_2" self.architect.server.register_mock_worker(mock_worker_name) workers = self.db.find_workers(worker_name=mock_worker_name) worker_2 = workers[0] worker_id = worker_2.db_id # Register an agent that is already qualified mock_agent_details = "FAKE_ASSIGNMENT_2" self.db.grant_qualification(qualification_id, worker_2.db_id, 1) self.architect.server.register_mock_agent(worker_id, mock_agent_details) time.sleep(0.1) last_packet = self.architect.server.last_packet self.assertIsNotNone(last_packet) self.assertNotIn( "onboard_data", last_packet["data"], "Onboarding triggered for qualified agent", ) agents = self.db.find_agents() self.assertEqual(len(agents), 2, "Second agent not created without onboarding") self.assertEqual(len(task_runner.running_assignments), 1, "Task was not launched") self.assertFalse(worker_0.is_qualified(TEST_QUALIFICATION_NAME)) self.assertTrue(worker_0.is_disqualified(TEST_QUALIFICATION_NAME)) self.assertTrue(worker_1.is_qualified(TEST_QUALIFICATION_NAME)) self.assertFalse(worker_1.is_disqualified(TEST_QUALIFICATION_NAME)) self.assertTrue(worker_2.is_qualified(TEST_QUALIFICATION_NAME)) self.assertFalse(worker_2.is_disqualified(TEST_QUALIFICATION_NAME)) agents = [a.agent for a in sup.agents.values()] # Make both agents act agent_id_1, agent_id_2 = agents[0].db_id, agents[1].db_id agent_1_data = agents[0].datastore.agent_data[agent_id_1] agent_2_data = agents[1].datastore.agent_data[agent_id_2] self.architect.server.send_agent_act(agent_id_1, {"text": "message1"}) self.architect.server.send_agent_act(agent_id_2, {"text": "message2"}) # Give up to 1 seconds for the actual operation to occur start_time = time.time() TIMEOUT_TIME = 1 while time.time() - start_time < TIMEOUT_TIME: if len(agent_1_data["acts"]) > 0: break time.sleep(0.1) self.assertLess(time.time() - start_time, TIMEOUT_TIME, "Did not process messages in time") # Give up to 1 seconds for the task to complete afterwards start_time = time.time() TIMEOUT_TIME = 1 while time.time() - start_time < TIMEOUT_TIME: if len(task_runner.running_assignments) == 0: break time.sleep(0.1) self.assertLess(time.time() - start_time, TIMEOUT_TIME, "Did not complete task in time") # Give up to 1 seconds for all messages to propogate start_time = time.time() TIMEOUT_TIME = 1 while time.time() - start_time < TIMEOUT_TIME: if self.architect.server.actions_observed == 2: break time.sleep(0.1) self.assertLess(time.time() - start_time, TIMEOUT_TIME, "Not all actions observed in time") sup.shutdown() self.assertTrue(channel_info.channel.is_closed()) def test_register_job_with_onboarding(self): """Test registering and running a job with onboarding""" # Handle baseline setup sup = Supervisor(self.db) self.sup = sup TEST_QUALIFICATION_NAME = "test_onboarding_qualification" # Register onboarding arguments for blueprint task_run_args = self.task_run.args task_run_args.blueprint.use_onboarding = True task_run_args.blueprint.onboarding_qualification = TEST_QUALIFICATION_NAME task_run_args.blueprint.timeout_time = 5 task_run_args.blueprint.is_concurrent = False self.task_run.get_task_config() # Supervisor expects that blueprint setup has already occurred blueprint = self.task_run.get_blueprint() TaskRunnerClass = MockBlueprint.TaskRunnerClass task_runner = TaskRunnerClass(self.task_run, task_run_args, EMPTY_STATE) sup.register_job(self.architect, task_runner, self.provider) self.assertEqual(len(sup.channels), 1) channel_info = list(sup.channels.values())[0] self.assertIsNotNone(channel_info) self.assertTrue(channel_info.channel.is_alive()) channel_id = channel_info.channel_id task_runner = channel_info.job.task_runner self.assertIsNotNone(channel_id) self.assertEqual( len(self.architect.server.subs), 1, "MockServer doesn't see registered channel", ) self.assertIsNotNone( self.architect.server.last_alive_packet, "No alive packet received by server", ) sup.launch_sending_thread() self.assertIsNotNone(sup.sending_thread) # Register a worker mock_worker_name = "MOCK_WORKER" self.architect.server.register_mock_worker(mock_worker_name) workers = self.db.find_workers(worker_name=mock_worker_name) self.assertEqual(len(workers), 1, "Worker not successfully registered") worker_1 = workers[0] self.architect.server.register_mock_worker(mock_worker_name) workers = self.db.find_workers(worker_name=mock_worker_name) self.assertEqual(len(workers), 1, "Worker potentially re-registered") worker_id = workers[0].db_id self.assertEqual(len(task_runner.running_units), 0) # Fail to register a blocked agent mock_agent_details = "FAKE_ASSIGNMENT" qualification_id = blueprint.onboarding_qualification_id self.db.grant_qualification(qualification_id, worker_1.db_id, 0) self.architect.server.register_mock_agent(worker_id, mock_agent_details) agents = self.db.find_agents() self.assertEqual(len(agents), 0, "Agent should not be created yet, failed onboarding") time.sleep(0.1) last_packet = self.architect.server.last_packet self.assertIsNotNone(last_packet) self.assertNotIn( "onboard_data", last_packet["data"], "Onboarding triggered for disqualified worker", ) self.assertIsNone(last_packet["data"]["agent_id"], "worker assigned real agent id") self.architect.server.last_packet = None self.db.revoke_qualification(qualification_id, worker_id) # Register an agent successfully mock_agent_details = "FAKE_ASSIGNMENT" self.architect.server.register_mock_agent(worker_id, mock_agent_details) agents = self.db.find_agents() self.assertEqual(len(agents), 0, "Agent should not be created yet - need onboarding") onboard_agents = self.db.find_onboarding_agents() self.assertEqual(len(onboard_agents), 1, "Onboarding agent should have been created") time.sleep(0.1) last_packet = self.architect.server.last_packet self.assertIsNotNone(last_packet) self.assertIn("onboard_data", last_packet["data"], "Onboarding not triggered") self.architect.server.last_packet = None # Submit onboarding from the agent onboard_data = {"should_pass": False} self.architect.server.register_mock_agent_after_onboarding( worker_id, onboard_agents[0].get_agent_id(), onboard_data) agents = self.db.find_agents() self.assertEqual(len(agents), 0, "Failed agent created after onboarding") # Re-register as if refreshing self.architect.server.register_mock_agent(worker_id, mock_agent_details) agents = self.db.find_agents() self.assertEqual(len(agents), 0, "Failed agent created after onboarding") self.assertEqual(len(sup.agents), 0, "Failed agent registered with supervisor") self.assertEqual( len(task_runner.running_units), 0, "Task should not launch with failed worker", ) # Register another worker mock_worker_name = "MOCK_WORKER_2" self.architect.server.register_mock_worker(mock_worker_name) workers = self.db.find_workers(worker_name=mock_worker_name) worker_2 = workers[0] worker_id = worker_2.db_id # Register an agent that is already qualified mock_agent_details = "FAKE_ASSIGNMENT_2" self.db.grant_qualification(qualification_id, worker_2.db_id, 1) self.architect.server.register_mock_agent(worker_id, mock_agent_details) time.sleep(0.1) last_packet = self.architect.server.last_packet self.assertIsNotNone(last_packet) self.assertNotIn( "onboard_data", last_packet["data"], "Onboarding triggered for qualified agent", ) agents = self.db.find_agents() self.assertEqual(len(agents), 1, "Second agent not created without onboarding") self.assertEqual(len(task_runner.running_units), 1, "Tasks were not launched") self.assertFalse(worker_1.is_qualified(TEST_QUALIFICATION_NAME)) self.assertTrue(worker_1.is_disqualified(TEST_QUALIFICATION_NAME)) self.assertTrue(worker_2.is_qualified(TEST_QUALIFICATION_NAME)) self.assertFalse(worker_2.is_disqualified(TEST_QUALIFICATION_NAME)) # Register another worker mock_worker_name = "MOCK_WORKER_3" self.architect.server.register_mock_worker(mock_worker_name) workers = self.db.find_workers(worker_name=mock_worker_name) worker_3 = workers[0] worker_id = worker_3.db_id mock_agent_details = "FAKE_ASSIGNMENT_3" self.architect.server.register_mock_agent(worker_id, mock_agent_details) agents = self.db.find_agents() self.assertEqual(len(agents), 1, "Agent should not be created yet - need onboarding") onboard_agents = self.db.find_onboarding_agents() self.assertEqual(len(onboard_agents), 2, "Onboarding agent should have been created") time.sleep(0.1) last_packet = self.architect.server.last_packet self.assertIsNotNone(last_packet) self.assertIn("onboard_data", last_packet["data"], "Onboarding not triggered") self.architect.server.last_packet = None # Submit onboarding from the agent onboard_data = {"should_pass": True} self.architect.server.register_mock_agent_after_onboarding( worker_id, onboard_agents[1].get_agent_id(), onboard_data) agents = self.db.find_agents() self.assertEqual(len(agents), 2, "Agent not created after onboarding") # Re-register as if refreshing self.architect.server.register_mock_agent(worker_id, mock_agent_details) agents = self.db.find_agents() self.assertEqual(len(agents), 2, "Duplicate agent created after onboarding") agent = agents[1] self.assertIsNotNone(agent) self.assertEqual(len(sup.agents), 2, "Agent not registered supervisor after onboarding") self.assertEqual(len(task_runner.running_units), 2, "Task not launched after onboarding") agents = [a.agent for a in sup.agents.values()] # Make both agents act agent_id_1, agent_id_2 = agents[0].db_id, agents[1].db_id agent_1_data = agents[0].datastore.agent_data[agent_id_1] agent_2_data = agents[1].datastore.agent_data[agent_id_2] self.architect.server.send_agent_act(agent_id_1, {"text": "message1"}) self.architect.server.send_agent_act(agent_id_2, {"text": "message2"}) # Give up to 1 seconds for the actual operation to occur start_time = time.time() TIMEOUT_TIME = 1 while time.time() - start_time < TIMEOUT_TIME: if len(agent_1_data["acts"]) > 0: break time.sleep(0.1) self.assertLess(time.time() - start_time, TIMEOUT_TIME, "Did not process messages in time") # Give up to 1 seconds for the task to complete afterwards start_time = time.time() TIMEOUT_TIME = 1 while time.time() - start_time < TIMEOUT_TIME: if len(task_runner.running_units) == 0: break time.sleep(0.1) self.assertLess(time.time() - start_time, TIMEOUT_TIME, "Did not complete task in time") # Give up to 1 seconds for all messages to propogate start_time = time.time() TIMEOUT_TIME = 1 while time.time() - start_time < TIMEOUT_TIME: if self.architect.server.actions_observed == 2: break time.sleep(0.1) self.assertLess(time.time() - start_time, TIMEOUT_TIME, "Not all actions observed in time") sup.shutdown() self.assertTrue(channel_info.channel.is_closed())
def parse_and_launch_run( self, arg_list: Optional[List[str]] = None, extra_args: Optional[Dict[str, Any]] = None, ) -> str: """ Parse the given arguments and launch a job. """ if extra_args is None: extra_args = {} # Extract the abstractions being used parser = self._get_baseline_argparser() type_args, task_args_string = parser.parse_known_args(arg_list) requesters = self.db.find_requesters(requester_name=type_args.requester_name) if len(requesters) == 0: raise EntryDoesNotExistException( f"No requester found with name {type_args.requester_name}" ) requester = requesters[0] requester_id = requester.db_id provider_type = requester.provider_type # Parse the arguments for the abstractions to ensure # everything required is set BlueprintClass = get_blueprint_from_type(type_args.blueprint_type) ArchitectClass = get_architect_from_type(type_args.architect_type) CrowdProviderClass = get_crowd_provider_from_type(provider_type) task_args, _unknown = self._parse_args_from_classes( BlueprintClass, ArchitectClass, CrowdProviderClass, task_args_string ) task_args.update(extra_args) # Load the classes to force argument validation before anything # is actually created in the database # TODO(#94) perhaps parse the arguments for these things one at a time? BlueprintClass.assert_task_args(task_args) ArchitectClass.assert_task_args(task_args) CrowdProviderClass.assert_task_args(task_args) # Find an existing task or create a new one task_name = task_args.get("task_name") if task_name is None: task_name = type_args.blueprint_type logger.warning( f"Task is using the default blueprint name {task_name} as a name, as no task_name is provided" ) tasks = self.db.find_tasks(task_name=task_name) task_id = None if len(tasks) == 0: task_id = self.db.new_task(task_name, type_args.blueprint_type) else: task_id = tasks[0].db_id logger.info(f"Creating a task run under task name: {task_name}") # Create a new task run new_run_id = self.db.new_task_run( task_id, requester_id, " ".join([shlex.quote(x) for x in task_args_string]), provider_type, type_args.blueprint_type, requester.is_sandbox(), ) task_run = TaskRun(self.db, new_run_id) try: # If anything fails after here, we have to cleanup the architect build_dir = os.path.join(task_run.get_run_dir(), "build") os.makedirs(build_dir, exist_ok=True) architect = ArchitectClass(self.db, task_args, task_run, build_dir) # Register the blueprint with args to the task run, # ensure cached blueprint = BlueprintClass(task_run, task_args) task_run.get_blueprint(opts=task_args) # Setup and deploy the server built_dir = architect.prepare() task_url = architect.deploy() # TODO(#102) maybe the cleanup (destruction of the server configuration?) should only # happen after everything has already been reviewed, this way it's possible to # retrieve the exact build directory to review a task for real architect.cleanup() # Create the backend runner task_runner = BlueprintClass.TaskRunnerClass(task_run, task_args) # Small hack for auto appending block qualification existing_qualifications = task_args.get("qualifications", []) if task_args.get("block_qualification") is not None: existing_qualifications.append( make_qualification_dict( task_args["block_qualification"], QUAL_NOT_EXIST, None ) ) if task_args.get("onboarding_qualification") is not None: existing_qualifications.append( make_qualification_dict( OnboardingRequired.get_failed_qual( task_args["onboarding_qualification"] ), QUAL_NOT_EXIST, None, ) ) task_args["qualifications"] = existing_qualifications # Register the task with the provider provider = CrowdProviderClass(self.db) provider.setup_resources_for_task_run(task_run, task_args, task_url) initialization_data_array = blueprint.get_initialization_data() # Link the job together job = self.supervisor.register_job( architect, task_runner, provider, existing_qualifications ) if self.supervisor.sending_thread is None: self.supervisor.launch_sending_thread() except (KeyboardInterrupt, Exception) as e: logger.error( "Encountered error while launching run, shutting down", exc_info=True ) try: architect.shutdown() except (KeyboardInterrupt, Exception) as architect_exception: logger.exception( f"Could not shut down architect: {architect_exception}", exc_info=True, ) raise e launcher = TaskLauncher(self.db, task_run, initialization_data_array) launcher.create_assignments() launcher.launch_units(task_url) self._task_runs_tracked[task_run.db_id] = TrackedRun( task_run=task_run, task_launcher=launcher, task_runner=task_runner, architect=architect, job=job, ) return task_run.db_id
def setUp(self): self.data_dir = tempfile.mkdtemp() database_path = os.path.join(self.data_dir, "mephisto.db") self.db = LocalMephistoDB(database_path) self.task_run_id = get_test_task_run(self.db) self.task_run = TaskRun(self.db, self.task_run_id)
class Agent(ABC): """ 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, connection status, etc. """ def __init__(self, db: "MephistoDB", db_id: str, row: Optional[Mapping[str, Any]] = None): self.db: "MephistoDB" = db if row is None: row = db.get_agent(db_id) assert row is not None, f"Given db_id {db_id} did not exist in given db" self.db_id: str = row["agent_id"] self.db_status = row["status"] self.worker_id = row["worker_id"] self.unit_id = row["unit_id"] self.task_type = row["task_type"] self.provider_type = row["provider_type"] self.pending_observations: List["Packet"] = [] self.pending_actions: List["Packet"] = [] self.has_action = threading.Event() self.has_action.clear() self.wants_action = threading.Event() self.wants_action.clear() self.has_updated_status = threading.Event() self.assignment_id = row["assignment_id"] self.task_run_id = row["task_run_id"] self.task_id = row["task_id"] self.did_submit = threading.Event() # Deferred loading of related entities self._worker: Optional["Worker"] = None self._unit: Optional["Unit"] = None self._assignment: Optional["Assignment"] = None self._task_run: Optional["TaskRun"] = None self._task: Optional["Task"] = None # Follow-up initialization self.state = AgentState(self) # type: ignore def __new__(cls, db: "MephistoDB", db_id: str, row: Optional[Mapping[str, Any]] = None) -> "Agent": """ The new method is overridden to be able to automatically generate the expected Agent class without needing to specifically find it for a given db_id. As such it is impossible to create a base Agent as you will instead be returned the correct Agent class according to the crowdprovider associated with this Agent. """ from mephisto.core.registry import get_crowd_provider_from_type if cls == Agent: # We are trying to construct a Agent, find what type to use and # create that instead if row is None: row = db.get_agent(db_id) assert row is not None, f"Given db_id {db_id} did not exist in given db" correct_class = get_crowd_provider_from_type( row["provider_type"]).AgentClass return super().__new__(correct_class) else: # We are constructing another instance directly return super().__new__(cls) def get_agent_id(self) -> str: """Return this agent's id""" return self.db_id def get_worker(self) -> Worker: """ Return the worker that is using this agent for a task """ if self._worker is None: self._worker = Worker(self.db, self.worker_id) return self._worker def get_unit(self) -> "Unit": """ Return the Unit that this agent is working on. """ if self._unit is None: from mephisto.data_model.assignment import Unit self._unit = Unit(self.db, self.unit_id) return self._unit def get_assignment(self) -> "Assignment": """Return the assignment this agent is working on""" if self._assignment is None: if self._unit is not None: self._assignment = self._unit.get_assignment() else: from mephisto.data_model.assignment import Assignment self._assignment = Assignment(self.db, self.assignment_id) return self._assignment def get_task_run(self) -> "TaskRun": """Return the TaskRun this agent is working within""" if self._task_run is None: if self._unit is not None: self._task_run = self._unit.get_task_run() elif self._assignment is not None: self._task_run = self._assignment.get_task_run() else: from mephisto.data_model.task import TaskRun self._task_run = TaskRun(self.db, self.task_run_id) return self._task_run def get_task(self) -> "Task": """Return the Task this agent is working within""" if self._task is None: if self._unit is not None: self._task = self._unit.get_task() elif self._assignment is not None: self._task = self._assignment.get_task() elif self._task_run is not None: self._task = self._task_run.get_task() else: from mephisto.data_model.task import Task self._task = Task(self.db, self.task_id) return self._task def get_data_dir(self) -> str: """ Return the directory to be storing any agent state for this agent into """ assignment_dir = self.get_assignment().get_data_dir() return os.path.join(assignment_dir, self.db_id) def update_status(self, new_status: str) -> None: """Update the database status of this agent, and possibly send a message to the frontend agent informing them of this update""" if self.db_status == new_status: return # Noop, this is already the case if self.db_status in AgentState.complete(): print(f"Updating a final status, was {self.db_status} " f"and want to set to {new_status}") self.db.update_agent(self.db_id, status=new_status) self.db_status = new_status self.has_updated_status.set() if new_status in [ AgentState.STATUS_RETURNED, AgentState.STATUS_DISCONNECT ]: # Disconnect statuses should free any pending acts self.has_action.set() self.did_submit.set() @staticmethod def _register_agent(db: "MephistoDB", worker: Worker, unit: "Unit", provider_type: str) -> "Agent": """ Create this agent in the mephisto db with the correct setup """ db_id = db.new_agent( worker.db_id, unit.db_id, unit.task_id, unit.task_run_id, unit.assignment_id, unit.task_type, provider_type, ) a = Agent(db, db_id) a.update_status(AgentState.STATUS_ACCEPTED) return a # Specialized child cases may need to implement the following @classmethod def new_from_provider_data( cls, db: "MephistoDB", worker: Worker, unit: "Unit", provider_data: Dict[str, Any], ) -> "Agent": """ Wrapper around the new method that allows registering additional bookkeeping information from a crowd provider for this agent """ agent = cls.new(db, worker, unit) unit.worker_id = worker.db_id agent._unit = unit return agent def observe(self, packet: "Packet") -> None: """ Pass the observed information to the AgentState, then queue the information to be pushed to the user """ sending_packet = packet.copy() sending_packet.receiver_id = self.db_id self.state.update_data(sending_packet) self.pending_observations.append(sending_packet) def act(self, timeout: Optional[int] = None) -> Optional["Packet"]: """ Request information from the Agent's frontend. If non-blocking, (timeout is None) should return None if no actions are ready to be returned. """ if len(self.pending_actions) == 0: self.wants_action.set() if timeout is None or timeout == 0: return None self.has_action.wait(timeout) if len(self.pending_actions) == 0: # various disconnect cases status = self.get_status() if status == AgentState.STATUS_DISCONNECT: raise AgentDisconnectedError(self.db_id) elif status == AgentState.STATUS_RETURNED: raise AgentReturnedError(self.db_id) self.update_status(AgentState.STATUS_TIMEOUT) raise AgentTimeoutError(timeout, self.db_id) assert len( self.pending_actions) > 0, "has_action released without an action!" act = self.pending_actions.pop(0) if "MEPHISTO_is_submit" in act.data and act.data["MEPHISTO_is_submit"]: self.did_submit.set() if len(self.pending_actions) == 0: self.has_action.clear() self.state.update_data(act) return act def get_status(self) -> str: """Get the status of this agent in their work on their unit""" if self.db_status not in AgentState.complete(): row = self.db.get_agent(self.db_id) if row["status"] != self.db_status: if row["status"] in [ AgentState.STATUS_RETURNED, AgentState.STATUS_DISCONNECT, ]: # Disconnect statuses should free any pending acts self.has_action.set() self.has_updated_status.set() self.db_status = row["status"] return self.db_status # Children classes should implement the following methods def approve_work(self) -> None: """Approve the work done on this agent's specific Unit""" raise NotImplementedError() def soft_reject_work(self) -> None: """ Pay a worker for attempted work, but mark it as below the quality bar for this assignment """ # TODO(OWN) extend this method to assign a soft block # qualification automatically if a threshold of # soft rejects as a proportion of total accepts # is exceeded self.approve_work() self.update_status(AgentState.STATUS_SOFT_REJECTED) def reject_work(self, reason) -> None: """Reject the work done on this agent's specific Unit""" raise NotImplementedError() def mark_done(self) -> None: """ Take any required step with the crowd_provider to ensure that the worker can submit their work and be marked as complete via a call to get_status """ raise NotImplementedError() @staticmethod def new(db: "MephistoDB", worker: Worker, unit: "Unit") -> "Agent": """ Create an agent for this worker to be used for work on the given Unit. Implementation should return the result of _register_agent when sure the agent can be successfully created to have it put into the db. """ raise NotImplementedError()
def get_units_for_run_id(self, run_id: str) -> List[Unit]: task_run = TaskRun(self.db, run_id) return self._get_units_for_task_runs([task_run])
def validate_and_run_config_or_die( self, run_config: DictConfig, shared_state: Optional[SharedTaskState] = None, ) -> str: """ Parse the given arguments and launch a job. """ if shared_state is None: shared_state = SharedTaskState() # First try to find the requester: requester_name = run_config.provider.requester_name requesters = self.db.find_requesters(requester_name=requester_name) if len(requesters) == 0: if run_config.provider.requester_name == "MOCK_REQUESTER": requesters = [get_mock_requester(self.db)] else: raise EntryDoesNotExistException( f"No requester found with name {requester_name}") requester = requesters[0] requester_id = requester.db_id provider_type = requester.provider_type assert provider_type == run_config.provider._provider_type, ( f"Found requester for name {requester_name} is not " f"of the specified type {run_config.provider._provider_type}, " f"but is instead {provider_type}.") # Next get the abstraction classes, and run validation # before anything is actually created in the database blueprint_type = run_config.blueprint._blueprint_type architect_type = run_config.architect._architect_type BlueprintClass = get_blueprint_from_type(blueprint_type) ArchitectClass = get_architect_from_type(architect_type) CrowdProviderClass = get_crowd_provider_from_type(provider_type) BlueprintClass.assert_task_args(run_config, shared_state) ArchitectClass.assert_task_args(run_config, shared_state) CrowdProviderClass.assert_task_args(run_config, shared_state) # Find an existing task or create a new one task_name = run_config.task.get("task_name", None) if task_name is None: task_name = blueprint_type logger.warning( f"Task is using the default blueprint name {task_name} as a name, " "as no task_name is provided") tasks = self.db.find_tasks(task_name=task_name) task_id = None if len(tasks) == 0: task_id = self.db.new_task(task_name, blueprint_type) else: task_id = tasks[0].db_id logger.info(f"Creating a task run under task name: {task_name}") # Create a new task run new_run_id = self.db.new_task_run( task_id, requester_id, json.dumps(OmegaConf.to_container(run_config, resolve=True)), provider_type, blueprint_type, requester.is_sandbox(), ) task_run = TaskRun(self.db, new_run_id) try: # If anything fails after here, we have to cleanup the architect build_dir = os.path.join(task_run.get_run_dir(), "build") os.makedirs(build_dir, exist_ok=True) architect = ArchitectClass(self.db, run_config, shared_state, task_run, build_dir) # Register the blueprint with args to the task run, # ensure cached blueprint = BlueprintClass(task_run, run_config, shared_state) task_run.get_blueprint(args=run_config, shared_state=shared_state) # Setup and deploy the server built_dir = architect.prepare() task_url = architect.deploy() # TODO(#102) maybe the cleanup (destruction of the server configuration?) should only # happen after everything has already been reviewed, this way it's possible to # retrieve the exact build directory to review a task for real architect.cleanup() # Create the backend runner task_runner = BlueprintClass.TaskRunnerClass( task_run, run_config, shared_state) # Small hack for auto appending block qualification existing_qualifications = shared_state.qualifications if run_config.blueprint.get("block_qualification", None) is not None: existing_qualifications.append( make_qualification_dict( run_config.blueprint.block_qualification, QUAL_NOT_EXIST, None)) if run_config.blueprint.get("onboarding_qualification", None) is not None: existing_qualifications.append( make_qualification_dict( OnboardingRequired.get_failed_qual( run_config.blueprint.onboarding_qualification, ), QUAL_NOT_EXIST, None, )) shared_state.qualifications = existing_qualifications # Register the task with the provider provider = CrowdProviderClass(self.db) provider.setup_resources_for_task_run(task_run, run_config, task_url) initialization_data_array = blueprint.get_initialization_data() # Link the job together job = self.supervisor.register_job(architect, task_runner, provider, existing_qualifications) if self.supervisor.sending_thread is None: self.supervisor.launch_sending_thread() except (KeyboardInterrupt, Exception) as e: logger.error( "Encountered error while launching run, shutting down", exc_info=True) try: architect.shutdown() except (KeyboardInterrupt, Exception) as architect_exception: logger.exception( f"Could not shut down architect: {architect_exception}", exc_info=True, ) raise e launcher = TaskLauncher(self.db, task_run, initialization_data_array) launcher.create_assignments() launcher.launch_units(task_url) self._task_runs_tracked[task_run.db_id] = TrackedRun( task_run=task_run, task_launcher=launcher, task_runner=task_runner, architect=architect, job=job, ) task_run.update_completion_progress(status=False) return task_run.db_id
def get_submitted_data(): try: task_run_ids = request.args.getlist("task_run_id") task_names = request.args.getlist("task_name") assignment_ids = request.args.getlist("assignment_id") unit_ids = request.args.getlist("unit_ids") statuses = request.args.getlist("status") db = app.extensions["db"] units = [] assignments = [] assert len( task_names) == 0, "Searching via task names not yet supported" task_runs = [TaskRun(db, task_run_id) for task_run_id in task_run_ids] for task_run in task_runs: assignments += task_run.get_assignments() assignments += [ Assignment(db, assignment_id) for assignment_id in assignment_ids ] if len(statuses) == 0: statuses = [ AssignmentState.COMPLETED, AssignmentState.ACCEPTED, AssignmentState.REJECTED, ] filtered_assignments = [ a for a in assignments if a.get_status() in statuses ] for assignment in assignments: units += assignment.get_units() units += [Unit(db, unit_id) for unit_id in unit_ids] all_unit_data = [] for unit in units: unit_data = { "assignment_id": unit.assignment_id, "task_run_id": unit.task_run_id, "status": unit.db_status, "unit_id": unit.db_id, "worker_id": unit.worker_id, "data": None, } agent = unit.get_assigned_agent() if agent is not None: unit_data["data"] = agent.state.get_data() unit_data["worker_id"] = agent.worker_id all_unit_data.append(unit_data) print(all_unit_data) return jsonify({"success": True, "units": all_unit_data}) except Exception as e: import traceback traceback.print_exc() return jsonify({"success": False, "msg": str(e)})