Cache preferred instance groups

When creating unified job, stash the list of pk values from the
instance groups returned from preferred_instance_groups so that the
task management system does not need to call out to this method
repeatedly.

.preferred_instance_groups_cache is the new field

awx/main/scheduler/task_manager.py

@@ -512,6 +512,12 @@ class TaskManager(TaskBase):
             ScheduleTaskManager().schedule()
         from awx.main.tasks.system import handle_work_error, handle_work_success
 
+        # update capacity for control node and execution node
+        if task.controller_node:
+            self.instances[task.controller_node].consume_capacity(settings.AWX_CONTROL_NODE_TASK_IMPACT)
+        if task.execution_node:
+            self.instances[task.execution_node].consume_capacity(task.task_impact)
+
         dependent_tasks = dependent_tasks or []
 
         task_actual = {
@@ -598,7 +604,8 @@ class TaskManager(TaskBase):
                 continue
 
             found_acceptable_queue = False
-            preferred_instance_groups = task.preferred_instance_groups
+
+            preferred_instance_groups = self.instance_groups.get_instance_groups_from_task_cache(task)
 
             # Determine if there is control capacity for the task
             if task.capacity_type == 'control':
@@ -618,7 +625,6 @@ class TaskManager(TaskBase):
             # All task.capacity_type == 'control' jobs should run on control plane, no need to loop over instance groups
             if task.capacity_type == 'control':
                 task.execution_node = control_instance.hostname
-                control_instance.consume_capacity(control_impact)
                 execution_instance = self.instances[control_instance.hostname].obj
                 task.log_lifecycle("controller_node_chosen")
                 task.log_lifecycle("execution_node_chosen")
@@ -649,9 +655,7 @@ class TaskManager(TaskBase):
                         control_instance = execution_instance
                         task.controller_node = execution_instance.hostname
 
-                    control_instance.consume_capacity(settings.AWX_CONTROL_NODE_TASK_IMPACT)
                     task.log_lifecycle("controller_node_chosen")
-                    execution_instance.consume_capacity(task.task_impact)
                     task.log_lifecycle("execution_node_chosen")
                     logger.debug(
                         "Starting {} in group {} instance {} (remaining_capacity={})".format(

awx/main/scheduler/task_manager_models.py

@@ -67,6 +67,7 @@ class TaskManagerInstanceGroups:
     def __init__(self, instances_by_hostname=None, instance_groups=None, instance_groups_queryset=None):
         self.instance_groups = dict()
         self.controlplane_ig = None
+        self.pk_ig_map = dict()
 
         if instance_groups is not None:  # for testing
             self.instance_groups = instance_groups
@@ -81,6 +82,7 @@ class TaskManagerInstanceGroups:
                         instances_by_hostname[instance.hostname] for instance in instance_group.instances.all() if instance.hostname in instances_by_hostname
                     ],
                 )
+                self.pk_ig_map[instance_group.pk] = instance_group
 
     def get_remaining_capacity(self, group_name):
         instances = self.instance_groups[group_name]['instances']
@@ -121,3 +123,17 @@ class TaskManagerInstanceGroups:
                 elif i.capacity > largest_instance.capacity:
                     largest_instance = i
         return largest_instance
+
+    def get_instance_groups_from_task_cache(self, task):
+        igs = []
+        if task.preferred_instance_groups_cache:
+            for pk in task.preferred_instance_groups_cache:
+                ig = self.pk_ig_map.get(pk, None)
+                if ig:
+                    igs.append(ig)
+                else:
+                    logger.warn(f"Unknown instance group with pk {pk} for task {task}")
+        if len(igs) == 0:
+            logger.warn(f"No instance groups in cache exist, defaulting to global instance groups for task {task}")
+            return task.global_instance_groups
+        return igs