def _get_init_data(self, packet, channel_info: ChannelInfo):
"""Get the initialization data for the assigned agent's task"""
task_runner = channel_info.job.task_runner
agent_id = packet.data["provider_data"]["agent_id"]
agent_info = self.agents[agent_id]
assert isinstance(
agent_info.agent, Agent
), f"Can only get init unit data for Agents, not OnboardingAgents, got {agent_info}"
unit_data = task_runner.get_init_data_for_agent(agent_info.agent)
agent_data_packet = Packet(
packet_type=PACKET_TYPE_INIT_DATA,
sender_id=SYSTEM_CHANNEL_ID,
receiver_id=channel_info.channel_id,
data={
"request_id": packet.data["request_id"],
"init_data": unit_data
},
)
self.message_queue.append(agent_data_packet)
if isinstance(unit_data,
dict) and unit_data.get("raw_messages") is not None:
# TODO bring these into constants somehow
for message in unit_data["raw_messages"]:
packet = Packet.from_dict(message)
packet.receiver_id = agent_id
agent_info.agent.pending_observations.append(packet)
def run_unit(self, unit: "Unit", agent: "Agent") -> None:
"""
ParlAI runners will initialize a task world, then run them to completion
if possible
"""
agents = [agent]
opt: Dict[str, Any] = self.shared_state.world_opt
parlai_agents = [MephistoAgentWrapper(a) for a in agents]
world = self.parlai_world_module.make_world(
opt, parlai_agents) # type: ignore
world_id = self.get_world_id("unit", unit.db_id)
self.id_to_worlds[world_id] = world
while not world.episode_done() and unit.db_id in self.running_units:
world.parley()
# TODO(WISH) it would be nice to have individual agents be able to submit their
# final things without needing to wait for their partner, such
# as if one needs to rate and the other doesn't
world.shutdown()
if hasattr(world, "prep_save_data"):
agent.observe(
Packet(
packet_type=PACKET_TYPE_AGENT_ACTION,
sender_id="mephisto",
receiver_id=agent.db_id,
data={
"id": "SUBMIT_WORLD_DATA",
"WORLD_DATA": world.prep_save_data(parlai_agents),
"text": "",
},
))
def _send_alive(self, channel_id: str) -> bool:
logger.info("Sending alive")
return self.channels[channel_id].enqueue_send(
Packet(
packet_type=PACKET_TYPE_ALIVE,
subject_id=SYSTEM_CHANNEL_ID,
))
def _send_status_update(self, agent_info: AgentInfo) -> None:
"""
Handle telling the frontend agent about a change in their
active status. (Pushing a change in AgentState)
"""
status = agent_info.agent.db_status
if isinstance(agent_info.agent, OnboardingAgent):
if status in [
AgentState.STATUS_APPROVED, AgentState.STATUS_REJECTED
]:
# We don't expose the updated status directly to the frontend here
# Can be simplified if we improve how bootstrap-chat handles
# the transition of onboarding states
status = AgentState.STATUS_WAITING
send_packet = Packet(
packet_type=PACKET_TYPE_UPDATE_AGENT_STATUS,
sender_id=SYSTEM_CHANNEL_ID,
receiver_id=agent_info.agent.get_agent_id(),
data={
"status": status,
"state": {
"done_text": STATUS_TO_TEXT_MAP.get(status),
"task_done":
status == AgentState.STATUS_PARTNER_DISCONNECT,
},
},
)
channel_info = self.channels[agent_info.used_channel_id]
channel_info.channel.send(send_packet)
def _mark_agent_done(self, agent_info: AgentInfo) -> None:
"""
Handle marking an agent as done, and telling the frontend agent
that they have successfully completed their task.
If the agent is in a final non-successful status, or already
told of partner disconnect, skip
"""
if agent_info.agent.db_status in AgentState.complete() + [
AgentState.STATUS_PARTNER_DISCONNECT
]:
return
send_packet = Packet(
packet_type=PACKET_TYPE_UPDATE_AGENT_STATUS,
sender_id=SYSTEM_CHANNEL_ID,
receiver_id=agent_info.agent.get_agent_id(),
data={
"status": "completed",
"state": {
"done_text":
"You have completed this task. Please submit.",
"task_done": True,
},
},
)
channel_info = self.channels[agent_info.used_channel_id]
channel_info.channel.send(send_packet)
def run_onboarding(self, agent: "OnboardingAgent") -> None:
"""
ParlAI Onboarding will initialize an onboarding
world, then run it to completion if possible
"""
opt: Dict[str, Any] = self.shared_state.onboarding_world_opt
parlai_agent = MephistoAgentWrapper(agent)
world = self.parlai_world_module.make_onboarding_world( # type: ignore
opt, parlai_agent)
world_id = self.get_world_id("onboard", agent.get_agent_id())
self.id_to_worlds[world_id] = world
while (not world.episode_done()
and agent.get_agent_id() in self.running_onboardings):
world.parley()
world.shutdown()
if hasattr(world, "prep_save_data"):
agent.observe(
Packet(
packet_type=PACKET_TYPE_AGENT_ACTION,
sender_id="mephisto",
receiver_id=agent.db_id,
data={
"id": "SUBMIT_WORLD_DATA",
"WORLD_DATA": world.prep_save_data([parlai_agent]),
"text": "",
},
))
def send_live_update(self, agent_id: str, data: Dict[str, Any]):
"""Send a live data packet to the given agent id"""
data_packet = Packet(
packet_type=PACKET_TYPE_CLIENT_BOUND_LIVE_UPDATE,
subject_id=agent_id,
data=data,
)
self._get_channel_for_agent(agent_id).enqueue_send(data_packet)
def _send_alive(self, channel_info: ChannelInfo) -> bool:
logger.info("Sending alive")
return channel_info.channel.send(
Packet(
packet_type=PACKET_TYPE_ALIVE,
sender_id=SYSTEM_CHANNEL_ID,
receiver_id=channel_info.channel_id,
))
def send_status_update(self, agent_id: str, status: str):
"""Update the status for the given agent"""
status_packet = Packet(
packet_type=PACKET_TYPE_UPDATE_STATUS,
subject_id=agent_id,
data={
"status": status,
},
)
self._get_channel_for_agent(agent_id).enqueue_send(status_packet)
def on_message(msg_json):
"""Incoming message handler defers to the internal handler"""
try:
packet_dict = json.loads(msg_json)
packet = Packet.from_dict(packet_dict)
self.on_message(self.channel_id, packet)
except Exception as e:
# TODO(CLEAN) properly handle only failed from_dict calls
logger.exception(repr(e), exc_info=True)
raise
def _request_status_update(self) -> None:
"""
Check last round of statuses, then request
an update from the server on all agent's current status
"""
for channel_id, channel in self.channels.items():
send_packet = Packet(
packet_type=PACKET_TYPE_REQUEST_STATUSES,
subject_id=SYSTEM_CHANNEL_ID,
data={},
)
channel.enqueue_send(send_packet)
def observe(self, act):
"""
ParlAI Agents observe a dict, we must convert these to packets?
"""
if act.get("message_id") is None:
act["message_id"] = str(uuid4())
packaged_act = Packet(
packet_type=PACKET_TYPE_AGENT_ACTION,
sender_id="mephisto",
receiver_id=self.__mephisto_agent_id,
data=act,
)
self.mephisto_agent.observe(packaged_act)
def _register_worker(self, packet: Packet, channel_info: ChannelInfo):
"""Process a worker registration packet to register a worker"""
crowd_data = packet.data["provider_data"]
crowd_provider = channel_info.job.provider
worker_name = crowd_data["worker_name"]
workers = self.db.find_workers(worker_name=worker_name)
if len(workers) == 0:
# TODO(WISH) get rid of sandbox designation
workers = self.db.find_workers(worker_name=worker_name +
"_sandbox")
if len(workers) == 0:
worker = crowd_provider.WorkerClass.new_from_provider_data(
self.db, crowd_data)
else:
worker = workers[0]
if not worker_is_qualified(worker, channel_info.job.qualifications):
self.message_queue.append(
Packet(
packet_type=PACKET_TYPE_PROVIDER_DETAILS,
sender_id=SYSTEM_CHANNEL_ID,
receiver_id=channel_info.channel_id,
data={
"request_id": packet.data["request_id"],
"worker_id": None
},
))
else:
self.message_queue.append(
Packet(
packet_type=PACKET_TYPE_PROVIDER_DETAILS,
sender_id=SYSTEM_CHANNEL_ID,
receiver_id=channel_info.channel_id,
data={
"request_id": packet.data["request_id"],
"worker_id": worker.db_id,
},
))
def _request_action(self, agent_info: AgentInfo) -> None:
"""
Request an act from the agent targetted here. If the
agent is found by the server, this request will be
forwarded.
"""
send_packet = Packet(
packet_type=PACKET_TYPE_REQUEST_ACTION,
sender_id=SYSTEM_CHANNEL_ID,
receiver_id=agent_info.agent.get_agent_id(),
data={},
)
channel_info = self.channels[agent_info.used_channel_id]
channel_info.channel.send(send_packet)
def agent_id(self, new_agent_id: str):
"""
We want to be able to display these labels to the
frontend users, so when these are updated by a
world we forward that to the frontend
"""
packaged_act = Packet(
packet_type=PACKET_TYPE_UPDATE_AGENT_STATUS,
sender_id="mephisto",
receiver_id=self.__mephisto_agent_id,
data={"state": {"agent_display_name": new_agent_id}},
)
self.mephisto_agent.observe(packaged_act)
self.__agent_id = new_agent_id
def run_assignment(self, assignment: "Assignment",
agents: List["Agent"]) -> None:
"""
ParlAI runners will initialize a task world, then run them to completion
if possible
"""
for agent in agents:
assert agent is not None, "task was not fully assigned"
opt: Dict[str, Any] = self.shared_state.world_opt
parlai_agents = [MephistoAgentWrapper(a) for a in agents]
try:
world = self.parlai_world_module.make_world(
opt,
parlai_agents,
initialization_data=assignment.get_assignment_data(
)) # type: ignore
except TypeError:
# make_world doesn't ask for initialization_data
world = self.parlai_world_module.make_world(
opt, parlai_agents) # type: ignore
world_id = self.get_world_id("assignment", assignment.db_id)
self.id_to_worlds[world_id] = world
while not world.episode_done(
) and assignment.db_id in self.running_assignments:
world.parley()
# TODO(WISH) it would be nice to have individual agents be able to submit their
# final things without needing to wait for their partner, such
# as if one needs to rate and the other doesn't
world.shutdown()
if hasattr(world, "prep_save_data"):
for idx in range(len(parlai_agents)):
agents[idx].observe(
Packet(
packet_type=PACKET_TYPE_AGENT_ACTION,
sender_id="mephisto",
receiver_id=agents[idx].db_id,
data={
"id":
"SUBMIT_WORLD_DATA",
"WORLD_DATA":
world.prep_save_data([parlai_agents[idx]]),
"text":
"",
},
))
def run_onboarding(self, agent: "OnboardingAgent") -> None:
"""
ParlAI Onboarding will initialize an onboarding
world, then run it to completion if possible
"""
opt: Dict[str, Any] = self.shared_state.onboarding_world_opt
parlai_agent = MephistoAgentWrapper(agent)
try:
world = self.parlai_world_module.make_onboarding_world(
opt,
parlai_agent,
initialization_data=self.get_init_data_for_agent(agent),
) # type: ignore
except TypeError:
# make_world doesn't ask for initialization_data
world = self.parlai_world_module.make_onboarding_world(opt, parlai_agent) # type: ignore
world_id = self.get_world_id("onboard", agent.get_agent_id())
self.id_to_worlds[world_id] = world
while (
not world.episode_done()
and agent.get_agent_id() in self.running_onboardings
):
world.parley()
world.shutdown()
if hasattr(world, "prep_save_data"):
agent.observe(
Packet(
packet_type=PACKET_TYPE_AGENT_ACTION,
sender_id="mephisto",
receiver_id=agent.db_id,
data={
"id": "SUBMIT_WORLD_DATA",
"WORLD_DATA": world.prep_save_data([parlai_agent]),
"text": "",
},
)
)
# Mark the agent as done, then wait for the incoming submit action
agent.mark_done()
while not agent.has_action.is_set():
done_act = agent.act()
if done_act is not None:
break
time.sleep(0.3)
def enqueue_agent_details(self, request_id: str,
additional_data: Dict[str, Any]):
"""
Synchronous method to enqueue a message sending the given agent details
"""
base_data = {"request_id": request_id}
for key, val in additional_data.items():
base_data[key] = val
self.message_queue.put(
Packet(
packet_type=PACKET_TYPE_AGENT_DETAILS,
subject_id=self.request_id_to_channel_id[request_id],
data=base_data,
))
self.process_outgoing_queue(self.message_queue)
self.log_metrics_for_packet(self.request_id_to_packet[request_id])
del self.request_id_to_channel_id[request_id]
del self.request_id_to_packet[request_id]
def _request_status_update(self) -> None:
"""
Check last round of statuses, then request
an update from the server on all agent's current status
"""
if time.time() - self.last_status_check < STATUS_CHECK_TIME:
return
self.last_status_check = time.time()
for channel_id, channel_info in self.channels.items():
send_packet = Packet(
packet_type=PACKET_TYPE_REQUEST_AGENT_STATUS,
sender_id=SYSTEM_CHANNEL_ID,
receiver_id=channel_id,
data={},
)
channel_info.channel.send(send_packet)
def _send_status_update(self, agent_info: AgentInfo) -> None:
"""
Handle telling the frontend agent about a change in their
active status. (Pushing a change in AgentState)
"""
send_packet = Packet(
packet_type=PACKET_TYPE_UPDATE_AGENT_STATUS,
sender_id=SYSTEM_CHANNEL_ID,
receiver_id=agent_info.agent.get_agent_id(),
data={
"status": agent_info.agent.db_status,
"state": {
"done_text":
STATUS_TO_TEXT_MAP.get(agent_info.agent.db_status),
"task_done":
agent_info.agent.db_status ==
AgentState.STATUS_PARTNER_DISCONNECT,
},
},
)
channel_info = self.channels[agent_info.used_channel_id]
channel_info.channel.send(send_packet)
def _assign_unit_to_agent(self, packet: Packet, channel_info: ChannelInfo,
units: List["Unit"]):
"""Handle creating an agent for the specific worker to register an agent"""
crowd_data = packet.data["provider_data"]
task_run = channel_info.job.task_runner.task_run
crowd_provider = channel_info.job.provider
worker_id = crowd_data["worker_id"]
worker = Worker(self.db, worker_id)
logger.debug(f"Worker {worker_id} is being assigned one of "
f"{len(units)} units.")
reserved_unit = None
while len(units) > 0 and reserved_unit is None:
unit = units.pop(0)
reserved_unit = task_run.reserve_unit(unit)
if reserved_unit is None:
self.message_queue.append(
Packet(
packet_type=PACKET_TYPE_PROVIDER_DETAILS,
sender_id=SYSTEM_CHANNEL_ID,
receiver_id=channel_info.channel_id,
data={
"request_id": packet.data["request_id"],
"agent_id": None
},
))
else:
agent = crowd_provider.AgentClass.new_from_provider_data(
self.db, worker, unit, crowd_data)
logger.debug(f"Created agent {agent}, {agent.db_id}.")
self.message_queue.append(
Packet(
packet_type=PACKET_TYPE_PROVIDER_DETAILS,
sender_id=SYSTEM_CHANNEL_ID,
receiver_id=channel_info.channel_id,
data={
"request_id": packet.data["request_id"],
"agent_id": agent.get_agent_id(),
},
))
agent_info = AgentInfo(agent=agent,
used_channel_id=channel_info.channel_id)
self.agents[agent.get_agent_id()] = agent_info
self.agents_by_registration_id[
crowd_data["agent_registration_id"]] = agent_info
# Launch individual tasks
if not channel_info.job.task_runner.is_concurrent:
unit_thread = threading.Thread(
target=self._launch_and_run_unit,
args=(unit, agent_info, channel_info.job.task_runner),
name=f"Unit-thread-{unit.db_id}",
)
agent_info.assignment_thread = unit_thread
unit_thread.start()
else:
# See if the concurrent unit is ready to launch
assignment = unit.get_assignment()
agents = assignment.get_agents()
if None in agents:
agent.update_status(AgentState.STATUS_WAITING)
return # need to wait for all agents to be here to launch
# Launch the backend for this assignment
agent_infos = [
self.agents[a.db_id] for a in agents if a is not None
]
assign_thread = threading.Thread(
target=self._launch_and_run_assignment,
args=(assignment, agent_infos,
channel_info.job.task_runner),
name=f"Assignment-thread-{assignment.db_id}",
)
for agent_info in agent_infos:
agent_info.agent.update_status(AgentState.STATUS_IN_TASK)
agent_info.assignment_thread = assign_thread
assign_thread.start()
def _register_agent(self, packet: Packet, channel_info: ChannelInfo):
"""Process an agent registration packet to register an agent"""
# First see if this is a reconnection
crowd_data = packet.data["provider_data"]
agent_registration_id = crowd_data["agent_registration_id"]
logger.debug(
f"Incoming request to register agent {agent_registration_id}.")
if agent_registration_id in self.agents_by_registration_id:
agent = self.agents_by_registration_id[agent_registration_id].agent
# Update the source channel, in case it has changed
self.agents[
agent.get_agent_id()].used_channel_id = channel_info.channel_id
self.message_queue.append(
Packet(
packet_type=PACKET_TYPE_PROVIDER_DETAILS,
sender_id=SYSTEM_CHANNEL_ID,
receiver_id=channel_info.channel_id,
data={
"request_id": packet.data["request_id"],
"agent_id": agent.get_agent_id(),
},
))
logger.debug(
f"Found existing agent_registration_id {agent_registration_id}, "
f"reconnecting to agent {agent.get_agent_id()}.")
return
# Process a new agent
task_runner = channel_info.job.task_runner
task_run = task_runner.task_run
worker_id = crowd_data["worker_id"]
worker = Worker(self.db, worker_id)
# get the list of tentatively valid units
units = task_run.get_valid_units_for_worker(worker)
if len(units) == 0:
self.message_queue.append(
Packet(
packet_type=PACKET_TYPE_PROVIDER_DETAILS,
sender_id=SYSTEM_CHANNEL_ID,
receiver_id=channel_info.channel_id,
data={
"request_id": packet.data["request_id"],
"agent_id": None
},
))
logger.debug(
f"Found existing agent_registration_id {agent_registration_id}, "
f"had no valid units.")
return
# If there's onboarding, see if this worker has already been disqualified
worker_id = crowd_data["worker_id"]
worker = Worker(self.db, worker_id)
blueprint = task_run.get_blueprint(args=task_runner.args)
if isinstance(blueprint,
OnboardingRequired) and blueprint.use_onboarding:
if worker.is_disqualified(blueprint.onboarding_qualification_name):
self.message_queue.append(
Packet(
packet_type=PACKET_TYPE_PROVIDER_DETAILS,
sender_id=SYSTEM_CHANNEL_ID,
receiver_id=channel_info.channel_id,
data={
"request_id": packet.data["request_id"],
"agent_id": None,
},
))
logger.debug(
f"Worker {worker_id} is already disqualified by onboarding "
f"qual {blueprint.onboarding_qualification_name}.")
return
elif not worker.is_qualified(
blueprint.onboarding_qualification_name):
# Send a packet with onboarding information
onboard_data = blueprint.get_onboarding_data(worker.db_id)
onboard_agent = OnboardingAgent.new(self.db, worker, task_run)
onboard_agent.state.set_init_state(onboard_data)
agent_info = AgentInfo(agent=onboard_agent,
used_channel_id=channel_info.channel_id)
onboard_id = onboard_agent.get_agent_id()
# register onboarding agent
self.agents[onboard_id] = agent_info
self.onboarding_packets[onboard_id] = packet
self.message_queue.append(
Packet(
packet_type=PACKET_TYPE_PROVIDER_DETAILS,
sender_id=SYSTEM_CHANNEL_ID,
receiver_id=channel_info.channel_id,
data={
"request_id": packet.data["request_id"],
"agent_id": onboard_id,
"onboard_data": onboard_data,
},
))
logger.debug(
f"Worker {worker_id} is starting onboarding thread with "
f"onboarding agent id {onboard_id}.")
# Create an onboarding thread
onboard_thread = threading.Thread(
target=self._launch_and_run_onboarding,
args=(agent_info, channel_info.job.task_runner),
name=f"Onboard-thread-{onboard_id}",
)
onboard_agent.update_status(AgentState.STATUS_ONBOARDING)
agent_info.assignment_thread = onboard_thread
onboard_thread.start()
return
# Not onboarding, so just register directly
self._assign_unit_to_agent(packet, channel_info, units)
