def get_test_assignment(self) -> Assignment:
"""Create a test assignment for self.task_run using mock agents"""
task_run = self.task_run
assignment_id = self.db.new_assignment(
task_run.task_id,
task_run.db_id,
task_run.requester_id,
task_run.task_type,
task_run.provider_type,
)
assign = Assignment(self.db, assignment_id)
unit_id = self.db.new_unit(
task_run.task_id,
task_run.db_id,
task_run.requester_id,
assignment_id,
0,
0,
task_run.provider_type,
task_run.task_type,
)
unit = MockUnit(self.db, unit_id)
worker_id = self.db.new_worker("MOCK_TEST_WORKER", MOCK_PROVIDER_TYPE)
worker = MockWorker(self.db, worker_id)
agent_id = self.db.new_agent(
worker.db_id,
unit_id,
task_run.task_id,
task_run.db_id,
assignment_id,
task_run.task_type,
task_run.provider_type,
)
Agent = MockAgent(self.db, agent_id)
return assign
def _create_single_assignment(self, assignment_data) -> None:
""" Create a single assignment in the database using its read assignment_data """
task_run = self.task_run
task_config = task_run.get_task_config()
assignment_id = self.db.new_assignment(
task_run.task_id,
task_run.db_id,
task_run.requester_id,
task_run.task_type,
task_run.provider_type,
task_run.sandbox,
)
assignment = Assignment(self.db, assignment_id)
assignment.write_assignment_data(assignment_data)
self.assignments.append(assignment)
unit_count = len(assignment_data.unit_data)
for unit_idx in range(unit_count):
unit_id = self.db.new_unit(
task_run.task_id,
task_run.db_id,
task_run.requester_id,
assignment_id,
unit_idx,
task_config.task_reward,
task_run.provider_type,
task_run.task_type,
task_run.sandbox,
)
self.units.append(Unit(self.db, unit_id))
with self.unlaunched_units_access_condition:
self.unlaunched_units[unit_id] = Unit(self.db, unit_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)})
def get_assignment(self) -> "Assignment":
"""
Return the assignment that this Unit is part of.
"""
if self.__assignment is None:
from mephisto.data_model.assignment import Assignment
self.__assignment = Assignment(self.db, self.assignment_id)
return self.__assignment
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 test_unit_fails(self) -> None:
"""Ensure units fail to be created or loaded under failure conditions"""
assert self.db is not None, "No db initialized"
db: MephistoDB = self.db
# Cant get non-existent entry
with self.assertRaises(EntryDoesNotExistException):
unit = Unit(db, self.get_fake_id("Unit"))
assignment_id = get_test_assignment(db)
assignment = Assignment(db, assignment_id)
unit_index = 0
pay_amount = 15.0
provider_type = PROVIDER_TYPE
# Can't use invalid assignment_id name
with self.assertRaises(EntryDoesNotExistException):
unit_id = db.new_unit(
assignment.task_id,
assignment.task_run_id,
assignment.requester_id,
self.get_fake_id("Assignment"),
unit_index,
pay_amount,
provider_type,
assignment.sandbox,
)
unit_id = db.new_unit(
assignment.task_id,
assignment.task_run_id,
assignment.requester_id,
assignment.db_id,
unit_index,
pay_amount,
provider_type,
assignment.sandbox,
)
# Can't create same unit again
with self.assertRaises(EntryAlreadyExistsException):
unit_id = db.new_unit(
assignment.task_id,
assignment.task_run_id,
assignment.requester_id,
assignment.db_id,
unit_index,
pay_amount,
provider_type,
assignment.sandbox,
)
# Ensure no units were created
units = db.find_units()
self.assertEqual(len(units), 1)
def test_unit(self) -> None:
"""Test creation and querying of units"""
assert self.db is not None, "No db initialized"
db: MephistoDB = self.db
# Check creation and retrieval of a unit
assignment_id = get_test_assignment(db)
assignment = Assignment(db, assignment_id)
unit_index = 0
pay_amount = 15.0
provider_type = PROVIDER_TYPE
unit_id = db.new_unit(
assignment.task_id,
assignment.task_run_id,
assignment.requester_id,
assignment.db_id,
unit_index,
pay_amount,
provider_type,
assignment.sandbox,
)
self.assertIsNotNone(unit_id)
self.assertTrue(isinstance(unit_id, str))
unit_row = db.get_unit(unit_id)
self.assertEqual(unit_row["assignment_id"], assignment_id)
self.assertEqual(unit_row["pay_amount"], pay_amount)
self.assertEqual(unit_row["status"], AssignmentState.CREATED)
unit = Unit(db, unit_id)
self.assertEqual(unit.assignment_id, assignment_id)
# Check finding for units
units = db.find_units()
self.assertEqual(len(units), 1)
self.assertTrue(isinstance(units[0], Unit))
self.assertEqual(units[0].db_id, unit_id)
self.assertEqual(units[0].assignment_id, assignment_id)
self.assertEqual(units[0].pay_amount, pay_amount)
# Check finding for specific units
units = db.find_units(assignment_id=assignment_id)
self.assertEqual(len(units), 1)
self.assertTrue(isinstance(units[0], Unit))
self.assertEqual(units[0].db_id, unit_id)
self.assertEqual(units[0].assignment_id, assignment_id)
self.assertEqual(units[0].pay_amount, pay_amount)
units = db.find_units(assignment_id=self.get_fake_id("Assignment"))
self.assertEqual(len(units), 0)
def get_test_unit(db: MephistoDB, unit_index=0) -> str:
# Check creation and retrieval of a unit
assignment_id = get_test_assignment(db)
pay_amount = 15.0
assignment = Assignment(db, assignment_id)
return db.new_unit(
assignment.task_id,
assignment.task_run_id,
assignment.requester_id,
assignment.db_id,
0,
pay_amount,
assignment.provider_type,
assignment.task_type,
)
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 find_assignments(
self,
task_run_id: Optional[str] = None,
task_id: Optional[str] = None,
requester_id: Optional[str] = None,
task_type: Optional[str] = None,
provider_type: Optional[str] = None,
sandbox: Optional[bool] = None,
) -> List[Assignment]:
"""
Try to find any task that matches the above. When called with no arguments,
return all tasks.
"""
with self.table_access_condition:
conn = self._get_connection()
c = conn.cursor()
c.execute(
"""
SELECT * from assignments
WHERE (?1 IS NULL OR task_run_id = ?1)
AND (?2 IS NULL OR task_id = ?2)
AND (?3 IS NULL OR requester_id = ?3)
AND (?4 IS NULL OR task_type = ?4)
AND (?5 IS NULL OR provider_type = ?5)
AND (?6 IS NULL OR sandbox = ?6)
""",
(
nonesafe_int(task_run_id),
nonesafe_int(task_id),
nonesafe_int(requester_id),
task_type,
provider_type,
sandbox,
),
)
rows = c.fetchall()
return [Assignment(self, str(r["assignment_id"]), row=r) for r in rows]
"inputs": {"something": True, "something else": False},
"outputs": {"some": "annotations"},
}
]
# Write a new task, and then complete it
for annotation in test_annotations:
assignment_id = db.new_assignment(
task_run.task_id,
task_run.db_id,
task_run.requester_id,
task_run.task_type,
task_run.provider_type,
task_run.sandbox,
)
assignment = Assignment(db, assignment_id)
assignment.write_assignment_data(
InitializationData(unit_data={}, shared=annotation["inputs"])
)
unit_id = db.new_unit(
task_run.task_id,
task_run.db_id,
task_run.requester_id,
assignment_id,
0, # Unit_index
0, # reward
task_run.provider_type,
task_run.task_type,
task_run.sandbox,
)
