Make sure that perf_device_trigger always try to assign different shards on same bot

Bug: 855302
Cq-Include-Trybots: luci.chromium.try:android_optional_gpu_tests_rel;luci.chromium.try:linux_optional_gpu_tests_rel;luci.chromium.try:mac_optional_gpu_tests_rel;luci.chromium.try:win_optional_gpu_tests_rel
Change-Id: Ifedb4069548c3c1706ae5b95e3ad922522b89178
Reviewed-on: https://chromium-review.googlesource.com/1111461Reviewed-by: Emily Hanley <eyaich@chromium.org>
Commit-Queue: Ned Nguyen <nednguyen@google.com>
Cr-Commit-Position: refs/heads/master@{#569586}

testing/trigger_scripts/perf_device_trigger.py

@@ -128,19 +128,21 @@ class PerfDeviceTriggerer(base_test_triggerer.BaseTestTriggerer):
                        'pool. Contact labs to rack in more hardware')
 
     shard_to_bot_assignment_map = {}
-    unallocated_bots = set(self._eligible_bots_by_ids.values())
+    unallocated_bots_by_ids = self._eligible_bots_by_ids.copy()
     for shard_index in xrange(args.shards):
       bot_id = self._query_swarming_for_last_shard_id(shard_index)
-      if bot_id:
+      if bot_id and bot_id in unallocated_bots_by_ids:
         bot = self._eligible_bots_by_ids[bot_id]
         shard_to_bot_assignment_map[shard_index] = bot
-        unallocated_bots.discard(bot)
+        unallocated_bots_by_ids.pop(bot_id)
       else:
         shard_to_bot_assignment_map[shard_index] = None
 
     # Now create sets of remaining healthy and bad bots
-    unallocated_healthy_bots = {b for b in unallocated_bots if b.is_alive()}
-    unallocated_bad_bots = {b for b in unallocated_bots if not b.is_alive()}
+    unallocated_healthy_bots = {
+        b for b in unallocated_bots_by_ids.values() if b.is_alive()}
+    unallocated_bad_bots = {
+        b for b in unallocated_bots_by_ids.values() if not b.is_alive()}
 
     # Try assigning healthy bots for new shards first.
     for shard_index, bot in sorted(shard_to_bot_assignment_map.iteritems()):

testing/trigger_scripts/perf_device_trigger_unittest.py

@@ -264,5 +264,21 @@ class UnitTest(unittest.TestCase):
     self.assertIn(expected_task_assignment.get(1), new_healthy_bots)
     self.assertIn(expected_task_assignment.get(2), new_healthy_bots)
 
+  def test_previously_duplicate_task_assignemnts(self):
+    triggerer = self.setup_and_trigger(
+        previous_task_assignment_map={0: 'build3', 1: 'build3', 2: 'build5',
+                                      3: 'build6'},
+        alive_bots=['build3', 'build4', 'build5', 'build7'],
+        dead_bots=['build1', 'build6'])
+    expected_task_assignment = self.get_triggered_shard_to_bot(
+        triggerer, num_shards=3)
+
+    # Test that the new assignment will add a new bot to avoid
+    # assign 'build3' to both shard 0 & shard 1 as before.
+    # It also replaces the dead 'build6' bot.
+    self.assertEquals(set(expected_task_assignment.values()),
+        {'build3', 'build4', 'build5', 'build7'})
+
+
 if __name__ == '__main__':
   unittest.main()