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_unit(self) -> "Unit":
"""
Return the Unit that this agent is working on.
"""
if self._unit is None:
from mephisto.data_model.unit import Unit
self._unit = Unit(self.db, self.unit_id)
return self._unit
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 find_units(
self,
task_id: Optional[str] = None,
task_run_id: Optional[str] = None,
requester_id: Optional[str] = None,
assignment_id: Optional[str] = None,
unit_index: Optional[int] = None,
provider_type: Optional[str] = None,
task_type: Optional[str] = None,
agent_id: Optional[str] = None,
worker_id: Optional[str] = None,
sandbox: Optional[bool] = None,
status: Optional[str] = None,
) -> List[Unit]:
"""
Try to find any unit that matches the above. When called with no arguments,
return all units.
"""
with self.table_access_condition:
conn = self._get_connection()
c = conn.cursor()
c.execute(
"""
SELECT * from units
WHERE (?1 IS NULL OR task_id = ?1)
AND (?2 IS NULL OR task_run_id = ?2)
AND (?3 IS NULL OR requester_id = ?3)
AND (?4 IS NULL OR assignment_id = ?4)
AND (?5 IS NULL OR unit_index = ?5)
AND (?6 IS NULL OR provider_type = ?6)
AND (?7 IS NULL OR task_type = ?7)
AND (?8 IS NULL OR agent_id = ?8)
AND (?9 IS NULL OR worker_id = ?9)
AND (?10 IS NULL OR sandbox = ?10)
AND (?11 IS NULL OR status = ?11)
""",
(
nonesafe_int(task_id),
nonesafe_int(task_run_id),
nonesafe_int(requester_id),
nonesafe_int(assignment_id),
unit_index,
provider_type,
task_type,
nonesafe_int(agent_id),
nonesafe_int(worker_id),
sandbox,
status,
),
)
rows = c.fetchall()
return [Unit(self, str(r["unit_id"]), row=r) for r in rows]
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 test_agent(self) -> None:
"""Test creation and querying of agents"""
assert self.db is not None, "No db initialized"
db: MephistoDB = self.db
# Check creation and retrieval of a agent
worker_name, worker_id = get_test_worker(db)
unit_id = get_test_unit(db)
unit = Unit(db, unit_id)
agent_id = db.new_agent(
worker_id,
unit_id,
unit.task_id,
unit.task_run_id,
unit.assignment_id,
unit.task_type,
unit.provider_type,
)
self.assertIsNotNone(agent_id)
self.assertTrue(isinstance(agent_id, str))
agent_row = db.get_agent(agent_id)
self.assertEqual(agent_row["worker_id"], worker_id)
self.assertEqual(agent_row["unit_id"], unit_id)
self.assertEqual(agent_row["status"], AgentState.STATUS_NONE)
# ensure the unit is assigned now
units = db.find_units(status=AssignmentState.ASSIGNED)
self.assertEqual(len(units), 1)
agent = Agent(db, agent_id)
self.assertEqual(agent.worker_id, worker_id)
# Check finding for agents
agents = db.find_agents()
self.assertEqual(len(agents), 1)
self.assertTrue(isinstance(agents[0], Agent))
self.assertEqual(agents[0].db_id, agent_id)
self.assertEqual(agents[0].worker_id, worker_id)
# Check finding for specific agents
agents = db.find_agents(worker_id=worker_id)
self.assertEqual(len(agents), 1)
self.assertTrue(isinstance(agents[0], Agent))
self.assertEqual(agents[0].db_id, agent_id)
self.assertEqual(agents[0].worker_id, worker_id)
agents = db.find_agents(worker_id=self.get_fake_id("Worker"))
self.assertEqual(len(agents), 0)
def get_mturk_ids_from_unit_id(db, unit_id: str) -> Dict[str, Optional[str]]:
"""
Find the relevant mturk ids from the given mephisto unit id
"""
mturk_unit = Unit(db, unit_id)
assignment_id = mturk_unit.get_mturk_assignment_id()
hit_id = mturk_unit.get_mturk_hit_id()
agent = mturk_unit.get_assigned_agent()
worker_id = None
if agent is not None:
worker_id = agent.get_worker().get_mturk_worker_id()
return {
"assignment_id": assignment_id,
"hit_id": hit_id,
"worker_id": worker_id
}
def test_agent_fails(self) -> None:
"""Ensure agents 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):
agent = Agent(db, self.get_fake_id("Agent"))
unit_id = get_test_unit(db)
worker_name, worker_id = get_test_worker(db)
unit = Unit(db, unit_id)
# Can't use invalid worker id
with self.assertRaises(EntryDoesNotExistException):
agent_id = db.new_agent(
self.get_fake_id("Worker"),
unit_id,
unit.task_id,
unit.task_run_id,
unit.assignment_id,
unit.task_type,
unit.provider_type,
)
# Can't use invalid unit id
with self.assertRaises(EntryDoesNotExistException):
agent_id = db.new_agent(
worker_id,
self.get_fake_id("Unit"),
unit.task_id,
unit.task_run_id,
unit.assignment_id,
unit.task_type,
unit.provider_type,
)
# Ensure no agents were created
agents = db.find_agents()
self.assertEqual(len(agents), 0)
def format_for_printing_data(data):
# Custom tasks can define methods for how to display their data in a relevant way
worker_name = Worker(db, data["worker_id"]).worker_name
contents = data["data"]
duration = contents["times"]["task_end"] - contents["times"]["task_start"]
metadata_string = (
f"Worker: {worker_name}\nUnit: {data['unit_id']}\n"
f"Duration: {int(duration)}\nStatus: {data['status']}\n")
inputs = contents["inputs"]
inputs_string = f"Character: {inputs['character_name']}\nDescription: {inputs['character_description']}\n"
outputs = contents["outputs"]
output_string = f" Rating: {outputs['rating']}\n"
found_files = outputs.get("files")
if found_files is not None:
file_dir = Unit(db,
data["unit_id"]).get_assigned_agent().get_data_dir()
output_string += f" Files: {found_files}\n"
output_string += f" File directory {file_dir}\n"
else:
output_string += f" Files: No files attached\n"
return f"-------------------\n{metadata_string}{inputs_string}{output_string}"
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,
)
unit = Unit(db, unit_id)
agent = MockAgent.new(db, worker, unit)
agent.state.state["inputs"] = annotation["inputs"]
agent.state.state["outputs"] = annotation["outputs"]
agent.state.save_data()
agent.mark_done()
agent.update_status(AgentState.STATUS_COMPLETED)
# Show tasks appear in MephistoDB:
mephisto_data_browser = MephistoDataBrowser(db=db)
units = mephisto_data_browser.get_units_for_task_name(input("Input task name: "))
for unit in units:
print(mephisto_data_browser.get_data_from_unit(unit))
