def test_onboarding_agents(self) -> None:
"""Ensure that the db can create and manipulate onboarding agents"""
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)
task = task_run.get_task()
worker_name, worker_id = get_test_worker(db)
onboarding_agent_id = db.new_onboarding_agent(worker_id, task.db_id,
task_run_id, "mock")
self.assertIsNotNone(onboarding_agent_id)
onboarding_agent = OnboardingAgent(db, onboarding_agent_id)
self.assertIsInstance(onboarding_agent, OnboardingAgent)
found_agents = db.find_onboarding_agents(worker_id=worker_id)
self.assertEqual(len(found_agents), 1)
self.assertIsInstance(found_agents[0], OnboardingAgent)
found_agent = found_agents[0]
self.assertEqual(found_agent.db_id, onboarding_agent_id)
self.assertEqual(found_agent.get_status(), AgentState.STATUS_NONE)
def find_onboarding_agents(
self,
status: Optional[str] = None,
worker_id: Optional[str] = None,
task_id: Optional[str] = None,
task_run_id: Optional[str] = None,
task_type: Optional[str] = None,
) -> List[OnboardingAgent]:
"""
Try to find any onboarding agent that matches the above. When called with no arguments,
return all onboarding agents.
"""
with self.table_access_condition:
conn = self._get_connection()
c = conn.cursor()
c.execute(
"""
SELECT * from onboarding_agents
WHERE (?1 IS NULL OR status = ?1)
AND (?2 IS NULL OR worker_id = ?2)
AND (?3 IS NULL OR task_id = ?3)
AND (?4 IS NULL OR task_run_id = ?4)
AND (?5 IS NULL OR task_type = ?5)
""",
(
status,
nonesafe_int(worker_id),
nonesafe_int(task_id),
nonesafe_int(task_run_id),
task_type,
),
)
rows = c.fetchall()
return [
OnboardingAgent(self, str(r["onboarding_agent_id"]), row=r)
for r in rows
]
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)
async def register_agent(self, crowd_data: Dict[str, Any],
worker: "Worker", request_id: str):
"""Process an agent registration packet to register an agent, returning the agent_id"""
# Process a new agent
logger.debug(f"Registering agent {crowd_data}, {request_id}")
live_run = self.get_live_run()
loop = live_run.loop_wrap.loop
task_run = live_run.task_run
agent_registration_id = crowd_data["agent_registration_id"]
# get the list of tentatively valid units
with EXTERNAL_FUNCTION_LATENCY.labels(
function="get_valid_units_for_worker").time():
units = task_run.get_valid_units_for_worker(worker)
if len(units) == 0:
AGENT_DETAILS_COUNT.labels(response="no_available_units").inc()
live_run.client_io.enqueue_agent_details(
request_id,
AgentDetails(
worker_id=worker.db_id,
failure_reason=WorkerFailureReasons.NO_AVAILABLE_UNITS,
).to_dict(),
)
logger.debug(
f"agent_registration_id {agent_registration_id}, had no valid units."
)
return
with EXTERNAL_FUNCTION_LATENCY.labels(
function="filter_units_for_worker").time():
units = await loop.run_in_executor(
None,
partial(live_run.task_runner.filter_units_for_worker, units,
worker),
)
# If there's onboarding, see if this worker has already been disqualified
blueprint = live_run.blueprint
if isinstance(blueprint,
OnboardingRequired) and blueprint.use_onboarding:
qual_name = blueprint.onboarding_qualification_name
assert (
qual_name
is not None), "Cannot be using onboarding and have a null qual"
if worker.is_disqualified(qual_name):
AGENT_DETAILS_COUNT.labels(response="not_qualified").inc()
live_run.client_io.enqueue_agent_details(
request_id,
AgentDetails(
worker_id=worker.db_id,
failure_reason=WorkerFailureReasons.NOT_QUALIFIED,
).to_dict(),
)
logger.debug(
f"Worker {worker.db_id} is already disqualified by onboarding "
f"qual {qual_name}.")
return
elif not worker.is_qualified(qual_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)
live_run.client_io.associate_agent_with_registration(
onboard_agent.get_agent_id(),
request_id,
crowd_data["agent_registration_id"],
)
onboard_agent.state.set_init_state(onboard_data)
onboard_agent.set_live_run(live_run)
onboard_id = onboard_agent.get_agent_id()
# register onboarding agent
self.onboarding_agents[onboard_id] = onboard_agent
self.onboarding_infos[onboard_id] = OnboardingInfo(
crowd_data=crowd_data,
request_id=request_id,
)
ONBOARDING_OUTCOMES.labels(outcome="launched").inc()
ACTIVE_ONBOARDINGS.inc()
AGENT_DETAILS_COUNT.labels(
response="assigned_onboarding").inc()
live_run.client_io.enqueue_agent_details(
request_id,
AgentDetails(
worker_id=worker.db_id,
agent_id=onboard_id,
init_task_data=onboard_data,
).to_dict(),
)
logger.info(f"{worker} is starting onboarding thread with "
f"onboarding {onboard_agent}.")
async def cleanup_onboarding():
del self.onboarding_agents[onboard_id]
del self.onboarding_infos[onboard_id]
ACTIVE_ONBOARDINGS.dec()
# Run the onboarding
live_run.task_runner.execute_onboarding(
onboard_agent, cleanup_onboarding)
return
if isinstance(blueprint,
ScreenTaskRequired) and blueprint.use_screening_task:
if (blueprint.worker_needs_screening(worker)
and blueprint.should_generate_unit()):
with EXTERNAL_FUNCTION_LATENCY.labels(
function="get_screening_unit_data").time():
screening_data = await loop.run_in_executor(
None, blueprint.get_screening_unit_data)
if screening_data is not None:
launcher = live_run.task_launcher
assert (
launcher is not None
), "LiveTaskRun must have launcher to use screening tasks"
with EXTERNAL_FUNCTION_LATENCY.labels(
function="launch_screening_unit").time():
screen_unit = await loop.run_in_executor(
None,
partial(
launcher.launch_screening_unit,
screening_data,
),
)
units = [screen_unit]
else:
AGENT_DETAILS_COUNT.labels(
response="no_available_units").inc()
live_run.client_io.enqueue_agent_details(
request_id,
AgentDetails(
worker_id=worker.db_id,
failure_reason=WorkerFailureReasons.
NO_AVAILABLE_UNITS,
).to_dict(),
)
logger.debug(
f"No screening units left for {agent_registration_id}."
)
return
if isinstance(blueprint, UseGoldUnit) and blueprint.use_golds:
if blueprint.should_produce_gold_for_worker(worker):
with EXTERNAL_FUNCTION_LATENCY.labels(
function="get_gold_unit_data_for_worker").time():
gold_data = await loop.run_in_executor(
None,
partial(blueprint.get_gold_unit_data_for_worker,
worker))
if gold_data is not None:
launcher = live_run.task_launcher
gold_unit = await loop.run_in_executor(
None,
partial(
launcher.launch_gold_unit,
gold_data,
),
)
units = [gold_unit]
else:
AGENT_DETAILS_COUNT.labels(
response="no_available_units").inc()
live_run.client_io.enqueue_agent_details(
request_id,
AgentDetails(
worker_id=worker.db_id,
failure_reason=WorkerFailureReasons.
NO_AVAILABLE_UNITS,
).to_dict(),
)
logger.debug(
f"No gold units left for {agent_registration_id}...")
return
# Not onboarding, so just register directly
await self._assign_unit_to_agent(crowd_data, worker, request_id, units)