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}")