Fixing Worker Pool race condition for assignments

[JackUrb]

Overview

An awkward race condition existed that allowed Mephisto to attempt to launch an Assignment before the WorkerPool had actually registered all of the Agents for that Assignment into its local self.agents dictionary. This could occur if the last _assign_unit_to_agent thread for an Assignment ran during the await-wrapped get_init_data_for_agent call of the second to last _assign_unit_to_agent for the same Assignment. Then the second to last thread would resume during the last call's get_init_data_for_agent. The outcome is that loop.run_in_executor(None, assignment.get_agents) would return all the Agents, but one wouldn't yet be added to the WorkerPool's tracking.

Fix Details

Rather than just using the None check for the Assignment.get_agents() call, we do a get check for the WorkerPool's agents as well, then bail early if any haven't been set.

Testing:

Automated testing

Would be nice if there was an asyncio testing library that intentionally took unusual routes (like by adding asyncio.sleep of random durations to all await calls) to catch this kind of thing sooner...

[codecov-commenter]

Codecov Report

Merging #875 (3e5af73) into main (ebac0c8) will increase coverage by 0.02%.
The diff coverage is 50.00%.

@@            Coverage Diff             @@
##             main     #875      +/-   ##
==========================================
+ Coverage   64.60%   64.62%   +0.02%     
==========================================
  Files         108      108              
  Lines        9301     9307       +6     
==========================================
+ Hits         6009     6015       +6     
  Misses       3292     3292              

Impacted Files	Coverage Δ
mephisto/operations/worker_pool.py	80.73% <50.00%> (-0.42%)	⬇️
mephisto/data_model/unit.py	77.59% <0.00%> (-0.55%)	⬇️
mephisto/utils/metrics.py	29.10% <0.00%> (-0.22%)	⬇️
...ractions/blueprints/mixins/screen_task_required.py	74.74% <0.00%> (+0.25%)	⬆️
mephisto/abstractions/architects/mock_architect.py	90.84% <0.00%> (+2.61%)	⬆️

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update ebac0c8...3e5af73. Read the comment docs.

[chiehminwei]

Thanks so much for the fix! LGTM