Add extra resource specifications

eurodatacube · Nov 23, 2023 · 18273d2 · 18273d2
1 parent 3fcddec
commit 18273d2
Show file tree

Hide file tree

Showing 4 changed files with 44 additions and 18 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,4 +1,5 @@
 ## 1.1.9
+* Add extra resource specifications
 * Remove kernel autoselect based on ancient images
 * Refactor output notebook handling
 

diff --git a/pygeoapi_kubernetes_papermill/notebook.py b/pygeoapi_kubernetes_papermill/notebook.py
@@ -209,6 +209,12 @@ def __init__(self, processor_def: dict) -> None:
         self.run_as_user: Optional[int] = processor_def["run_as_user"]
         self.run_as_group: Optional[int] = processor_def["run_as_group"]
         self.conda_store_groups: List[str] = processor_def["conda_store_groups"]
+        self.extra_resource_limits: Dict[str, str] = processor_def[
+            "extra_resource_limits"
+        ]
+        self.extra_resource_requests: Dict[str, str] = processor_def[
+            "extra_resource_requests"
+        ]
 
     def create_job_pod_spec(
         self,
@@ -320,7 +326,7 @@ def create_job_pod_spec(
             ],
             working_dir=str(CONTAINER_HOME),
             volume_mounts=extra_config.volume_mounts,
-            resources=_resource_requirements(requested),
+            resources=self._resource_requirements(requested),
             env=[
                 # this is provided in jupyter worker containers and we also use it
                 # for compatibility checks
@@ -435,6 +441,24 @@ def _image(self, requested_image: Optional[str]) -> str:
 
         return image
 
+    def _resource_requirements(self, requested: RequestParameters):
+        return k8s_client.V1ResourceRequirements(
+            limits=drop_none_values(
+                {
+                    "cpu": requested.cpu_limit,
+                    "memory": requested.mem_limit,
+                }
+                | self.extra_resource_limits
+            ),
+            requests=drop_none_values(
+                {
+                    "cpu": requested.cpu_requests,
+                    "memory": requested.mem_requests,
+                }
+                | self.extra_resource_requests
+            ),
+        )
+
     def affinity(self, requested_node_purpose: Optional[str]) -> k8s_client.V1Affinity:
         if node_purpose := requested_node_purpose:
             if not re.fullmatch(
@@ -601,23 +625,6 @@ def working_dir(notebook_path: PurePath) -> PurePath:
     return abs_notebook_path.parent
 
 
-def _resource_requirements(requested: RequestParameters):
-    return k8s_client.V1ResourceRequirements(
-        limits=drop_none_values(
-            {
-                "cpu": requested.cpu_limit,
-                "memory": requested.mem_limit,
-            }
-        ),
-        requests=drop_none_values(
-            {
-                "cpu": requested.cpu_requests,
-                "memory": requested.mem_requests,
-            }
-        ),
-    )
-
-
 def home_volume_config(home_volume_claim_name: str) -> ExtraConfig:
     return ExtraConfig(
         volumes=[

diff --git a/tests/test_kubernetes_manager.py b/tests/test_kubernetes_manager.py
@@ -197,6 +197,8 @@ def papermill_processor() -> PapermillNotebookKubernetesProcessor:
             "run_as_user": None,
             "run_as_group": None,
             "conda_store_groups": [],
+            "extra_resource_limits": {},
+            "extra_resource_requests": {},
         }
     )
 

diff --git a/tests/test_notebook_processor.py b/tests/test_notebook_processor.py
@@ -86,6 +86,8 @@ def _create_processor(def_override=None) -> PapermillNotebookKubernetesProcessor
             "run_as_user": None,
             "run_as_group": None,
             "conda_store_groups": [],
+            "extra_resource_limits": {},
+            "extra_resource_requests": {},
             **(def_override if def_override else {}),
         }
     )
@@ -613,3 +615,17 @@ def test_extra_volumes_are_added_on_request(create_pod_kwargs):
     assert "/mnt/my" in [
         m.mount_path for m in job_pod_spec.pod_spec.containers[0].volume_mounts
     ]
+
+
+def test_extra_requirements_are_added(create_pod_kwargs):
+    processor = _create_processor(
+        {
+            "extra_resource_requests": {"ice/cream": 1},
+            "extra_resource_limits": {"ice/cream": 3},
+        }
+    )
+    job_pod_spec = processor.create_job_pod_spec(**create_pod_kwargs)
+
+    resources = job_pod_spec.pod_spec.containers[0].resources
+    assert resources.requests["ice/cream"] == 1
+    assert resources.limits["ice/cream"] == 3