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_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 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 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()}" if status == self.db_status: return logger.debug(f"Updating status for {self} to {status}") self.db_status = status self.db.update_unit(self.db_id, status=status)
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 """ 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 and self.assignments_generator_done is not False: self.db.update_task_run(self.db_id, is_completed=True) self.__is_completed = True
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()}" if status == self.db_status: return logger.debug(f"Updating status for {self} to {status}") ACTIVE_UNIT_STATUSES.labels( status=self.db_status, unit_type=INDEX_TO_TYPE_MAP[self.unit_index]).dec() ACTIVE_UNIT_STATUSES.labels( status=status, unit_type=INDEX_TO_TYPE_MAP[self.unit_index]).inc() self.db_status = status self.db.update_unit(self.db_id, status=status)
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 if self.db_status in AssignmentState.final_unit(): if self.agent_id is None: return None return Agent.get(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 unit_copy = Unit.get(self.db, self.db_id) self.agent_id = unit_copy.agent_id if self.agent_id is not None: return Agent.get(self.db, self.agent_id) return 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
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}")
INDEX_TO_TYPE_MAP: DefaultDict[int, str] = defaultdict( lambda: "standard", { 0: "standard", SCREENING_UNIT_INDEX: "screening_unit", GOLD_UNIT_INDEX: "gold_unit", COMPENSATION_UNIT_INDEX: "compensation_unit", }, ) ACTIVE_UNIT_STATUSES = Gauge( "active_unit_statuses", "Tracking of all units current statuses", ["status", "unit_type"], ) for status in AssignmentState.valid_unit(): for unit_type in INDEX_TO_TYPE_MAP.values(): ACTIVE_UNIT_STATUSES.labels(status=status, unit_type=unit_type) class Unit(MephistoDataModelComponentMixin, metaclass=MephistoDBBackedABCMeta): """ 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,
def get_valid_units_for_worker(self, worker: "Worker") -> List["Unit"]: """ Get any units that the given worker could work on in this task run """ config = self.get_task_args() if config.allowed_concurrent != 0 or config.maximum_units_per_worker: current_units = self.db.find_units( task_run_id=self.db_id, worker_id=worker.db_id, status=AssignmentState.ASSIGNED, ) currently_active = len(current_units) if config.allowed_concurrent != 0: if currently_active >= config.allowed_concurrent: logger.debug( f"{worker} at maximum concurrent units {currently_active}" ) return [ ] # currently at the maximum number of concurrent units if config.maximum_units_per_worker != 0: completed_types = AssignmentState.completed() related_units = self.db.find_units( task_id=self.task_id, worker_id=worker.db_id, ) currently_completed = len([ u for u in related_units if u.db_status in completed_types ]) if (currently_active + currently_completed >= config.maximum_units_per_worker): logger.debug( f"{worker} at maximum units {currently_active}, {currently_completed}" ) return [ ] # Currently at the maximum number of units for this task task_units: List["Unit"] = self.get_units() unit_assigns: Dict[str, List["Unit"]] = {} for unit in task_units: assignment_id = unit.assignment_id if assignment_id not in unit_assigns: unit_assigns[assignment_id] = [] unit_assigns[assignment_id].append(unit) # Cannot pair with self units: List["Unit"] = [] for unit_set in unit_assigns.values(): is_self_set = map(lambda u: u.worker_id == worker.db_id, unit_set) if not any(is_self_set): units += unit_set # Valid units must be launched and must not be special units (negative indices) valid_units = [ u for u in units if u.get_status() == AssignmentState.LAUNCHED and u.unit_index >= 0 ] logger.debug(f"Found {len(valid_units)} available units") # Should load cached blueprint for SharedTaskState blueprint = self.get_blueprint() ret_units = [ u for u in valid_units if blueprint.shared_state.worker_can_do_unit(worker, u) ] logger.debug(f"This worker is qualified for {len(ret_units)} unit.") logger.debug(f"Found {ret_units[:3]} for {worker}.") return ret_units