check that requested cpu/memory less than or equals according limits

Signed-off-by: ntny <[email protected]>
kubeflow · Sep 25, 2024 · 0d429d2 · 0d429d2
1 parent 880e46d
commit 0d429d2
Show file tree

Hide file tree

Showing 2 changed files with 59 additions and 0 deletions.
diff --git a/sdk/python/kfp/dsl/pipeline_task.py b/sdk/python/kfp/dsl/pipeline_task.py
@@ -312,6 +312,21 @@ def set_caching_options(self, enable_caching: bool) -> 'PipelineTask':
         self._task_spec.enable_caching = enable_caching
         return self
 
+    def ensure_resource_requests_meet_limits(self) -> None:
+        resources = self.container_spec.resources
+        if (resources.memory_request is not None
+                and resources.memory_limit is not None
+                and resources.memory_request > resources.memory_limit):
+            raise ValueError(f'Requested memory: {resources.memory_request} cannot be greater than memory limit: {resources.memory_limit}. '
+                             'Check the set_memory_request and set_memory_limit parameters.')
+        if (resources.cpu_request is not None
+                and resources.cpu_limit is not None
+                and resources.cpu_request > resources.cpu_limit):
+            raise ValueError(
+                f'Requested cpu: {resources.cpu_request} cannot be greater than cpu limit: {resources.cpu_limit}. '
+                'Check the set_cpu_request and set_cpu_limit parameters.')
+
+
     def _ensure_container_spec_exists(self) -> None:
         """Ensures that the task has a container spec."""
         caller_method_name = inspect.stack()[1][3]
@@ -372,6 +387,8 @@ def set_cpu_request(
             self.container_spec.resources = structures.ResourceSpec(
                 cpu_request=cpu)
 
+        self.ensure_resource_requests_meet_limits()
+
         return self
 
     @block_if_final()
@@ -400,6 +417,8 @@ def set_cpu_limit(
             self.container_spec.resources = structures.ResourceSpec(
                 cpu_limit=cpu)
 
+        self.ensure_resource_requests_meet_limits()
+
         return self
 
     @block_if_final()
@@ -503,6 +522,8 @@ def set_memory_request(
             self.container_spec.resources = structures.ResourceSpec(
                 memory_request=memory)
 
+        self.ensure_resource_requests_meet_limits()
+
         return self
 
     @block_if_final()
@@ -530,6 +551,8 @@ def set_memory_limit(
             self.container_spec.resources = structures.ResourceSpec(
                 memory_limit=memory)
 
+        self.ensure_resource_requests_meet_limits()
+
         return self
 
     @block_if_final()

diff --git a/sdk/python/kfp/dsl/pipeline_task_test.py b/sdk/python/kfp/dsl/pipeline_task_test.py
@@ -231,6 +231,42 @@ def test_set_accelerator_limit(self, limit, expected_limit):
         self.assertEqual(expected_limit,
                          task.container_spec.resources.accelerator_count)
 
+    @parameterized.parameters(
+        {
+            'memory': '2M',
+            'limit': '1M',
+        },
+    )
+    def test_set_memory_request_greater_than_limit_should_raise(self, memory: str, limit: str):
+        task = pipeline_task.PipelineTask(
+            component_spec=structures.ComponentSpec.from_yaml_documents(
+                V2_YAML),
+            args={'input1': 'value'},
+        )
+        with self.assertRaisesRegex(
+                ValueError,
+                r'Requested memory: 2M cannot be greater than memory limit: 1M. '
+                'Check the set_memory_request and set_memory_limit parameters.'):
+            task.set_memory_request(memory).set_memory_limit(limit)
+
+    @parameterized.parameters(
+        {
+            'memory': '2',
+            'limit': '1',
+        },
+    )
+    def test_set_cpu_request_greater_than_limit_should_raise(self, memory: str, limit: str):
+        task = pipeline_task.PipelineTask(
+            component_spec=structures.ComponentSpec.from_yaml_documents(
+                V2_YAML),
+            args={'input1': 'value'},
+        )
+        with self.assertRaisesRegex(
+                ValueError,
+                r'Requested cpu: 2 cannot be greater than cpu limit: 1. '
+                'Check the set_cpu_request and set_cpu_limit parameters.'):
+            task.set_cpu_request(memory).set_cpu_limit(limit)
+
     @parameterized.parameters(
         {
             'memory': '1E',