def _retrieve_capacity_blocks_from_fleet_config()

in src/slurm_plugin/capacity_block_manager.py [0:0]

• 17 lines of code
• 6 McCabe index (conditional complexity)

    def _retrieve_capacity_blocks_from_fleet_config(self):
        """
        Collect list of capacity reservation target from all queues/compute-resources in the fleet config.

        Fleet config json has the following format:
        {
            "my-queue": {
                "my-compute-resource": {
                   "Api": "create-fleet",
                    "CapacityType": "on-demand|spot|capacity-block",
                    "AllocationStrategy": "lowest-price|capacity-optimized",
                    "Instances": [
                        { "InstanceType": "p4d.24xlarge" }
                    ],
                    "MaxPrice": "",
                    "Networking": {
                        "SubnetIds": ["subnet-123456"]
                    },
                    "CapacityReservationId": "id"
                }
            }
        }
        """
        capacity_blocks: Dict[str, CapacityBlock] = {}
        logger.info("Retrieving Capacity Blocks from fleet configuration.")

        for queue_name, queue_config in self._fleet_config.items():
            for compute_resource_name, compute_resource_config in queue_config.items():
                if self._is_compute_resource_associated_to_capacity_block(compute_resource_config):
                    capacity_block_id = self._capacity_reservation_id_from_compute_resource_config(
                        compute_resource_config
                    )
                    # retrieve existing CapacityBlock if exists or create a new one.
                    capacity_block = capacity_blocks.get(
                        capacity_block_id, CapacityBlock(capacity_block_id=capacity_block_id)
                    )
                    capacity_block.add_compute_resource(
                        queue_name=queue_name, compute_resource_name=compute_resource_name
                    )
                    capacity_blocks.update({capacity_block_id: capacity_block})

        return capacity_blocks