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_total_spend(self) -> float: """ Return the total amount spent on this run, based on any assignments that are still in a payable state. """ assigns = self.get_assignments() total_amount = 0.0 for assign in assigns: total_amount += assign.get_cost_of_statuses(AssignmentState.payable()) return total_amount
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_assignment_statuses(self) -> Dict[str, int]: """ Get the statistics for all of the assignments for this run. """ assigns = self.get_assignments() assigns_with_status = [(x, x.get_status()) for x in assigns] return { status: len( [x for x, had_status in assigns_with_status if had_status == status] ) for status in AssignmentState.valid() }
def get_assignments(self, status: Optional[str] = None) -> List["Assignment"]: """ Get assignments for this run, optionally filtering by their current status """ assert ( status is None or status in AssignmentState.valid() ), "Invalid assignment status" assignments = self.db.find_assignments(task_run_id=self.db_id) if status is not None: assignments = [a for a in assignments if a.get_status() == status] return assignments
def sync_completion_status(self) -> None: """ Update the is_complete status for this task run based on completion of subassignments. If this task run has no subassignments yet, it is not complete """ # TODO(#99) revisit when/if it's possible to add tasks to a completed run if not self.__is_completed and self.get_has_assignments(): statuses = self.get_assignment_statuses() has_incomplete = False for status in AssignmentState.incomplete(): if statuses[status] > 0: has_incomplete = True if not has_incomplete: self.db.update_task_run(self.db_id, is_completed=True) self.__is_completed = True
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
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}")