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
