Add custom allocation option

uclahs-cds · Nov 27, 2024 · f47ee5c · f47ee5c
1 parent c612346
commit f47ee5c
Showing 1 changed file with 17 additions and 1 deletion.
diff --git a/config/resource_handler/resource_handler.config b/config/resource_handler/resource_handler.config
@@ -142,6 +142,14 @@ class PipelineAllocation {
         }
     }
 
+    private void apply_custom_allocations(Map custom_allocations) {
+        custom_allocations.each { process, custom_allocation ->
+            custom_allocation.each { resource_type, allocation ->
+                this.processed_resources[process][resource_type] = this.system_resources.check_limits(allocation, resource_type)
+            }
+        }
+    }
+
     private void select_allocation_profile(Map params) {
         String profile_tag = null
 
@@ -256,13 +264,21 @@ resource_handler {
         }
     }
 
-    handle_resources = { Object resource_file, Map current_params=params ->
+    handle_resources = { Object resource_file, Map customized_allocations=[:], Map current_params=params ->
+        // Load base.config by default for all pipelines
+        includeConfig "${projectDir}/config/base.config"
+
         def allocation_handler = new PipelineAllocation(resource_file, current_params)
         def system_resources = allocation_handler.get_system_resources()
 
         // Set params for limits for each resources
         resource_handler.set_resource_limit_params(system_resources)
 
+        // Apply custom allocations if given
+        if (custom_allocations) {
+            allocation_handler.apply_custom_allocations(custom_allocations)
+        }
+
         // Apply base resource updates if given
         if (params.containsKey('base_resource_update') && params.base_resource_update) {
             allocation_handler.apply_base_updates(params.base_resource_update)