Address edge case and remove client-side naming restrictions

Signed-off-by: droctothorpe <[email protected]> Co-authored-by: zazulam <[email protected]> Co-authored-by: MonicaZhang1 <[email protected]>
kubeflow · Aug 1, 2024 · daa8c9f · daa8c9f
1 parent 1162bf9
commit daa8c9f
Show file tree

Hide file tree

Showing 8 changed files with 204 additions and 225 deletions.
diff --git a/backend/src/v2/compiler/argocompiler/argo.go b/backend/src/v2/compiler/argocompiler/argo.go
@@ -268,7 +268,17 @@ func ExtractBaseComponentName(componentName string) string {
 	baseComponentName := componentName
 	componentNameArray := strings.Split(componentName, "-")
 
-	if _, err := strconv.Atoi(componentNameArray[len(componentNameArray)-1]); err == nil {
+	if i, err := strconv.Atoi(componentNameArray[len(componentNameArray)-1]); err == nil {
+		// This conditional addresses a very rare edge case that can surface if
+		// and only if an end user has two component function names that are
+		// identical except that one of them ends in _1, for example, foo, and
+		// foo_1. In that context, ExtractBaseName will tell the foo_1 function
+		// to use the foo component's logic. Ultimately, the way that the IR
+		// compiler handles duplicate task instances can probably be optimized.
+		// When that happens, we can remove this conditional.
+		if i == 1 {
+			return baseComponentName
+		}
 		baseComponentName = strings.Join(componentNameArray[:len(componentNameArray)-1], "-")
 	}
 

diff --git a/sdk/python/kfp/compiler/compiler_test.py b/sdk/python/kfp/compiler/compiler_test.py
@@ -129,7 +129,7 @@ def comp():
 
 
 @dsl.component
-def return_one() -> int:
+def return_1() -> int:
     return 1
 
 
@@ -3369,43 +3369,43 @@ def test_cpu_memory_optional(self):
 
         @dsl.pipeline
         def simple_pipeline():
-            return_one()
-            return_one().set_cpu_limit('5')
-            return_one().set_memory_limit('50G')
-            return_one().set_cpu_request('2').set_cpu_limit(
+            return_1()
+            return_1().set_cpu_limit('5')
+            return_1().set_memory_limit('50G')
+            return_1().set_cpu_request('2').set_cpu_limit(
                 '5').set_memory_request('4G').set_memory_limit('50G')
 
         dict_format = json_format.MessageToDict(simple_pipeline.pipeline_spec)
 
         self.assertNotIn(
             'resources', dict_format['deploymentSpec']['executors']
-            ['exec-return-one']['container'])
+            ['exec-return-1']['container'])
 
         self.assertEqual(
-            5, dict_format['deploymentSpec']['executors']['exec-return-one-2']
+            5, dict_format['deploymentSpec']['executors']['exec-return-1-2']
             ['container']['resources']['cpuLimit'])
         self.assertNotIn(
             'memoryLimit', dict_format['deploymentSpec']['executors']
-            ['exec-return-one-2']['container']['resources'])
+            ['exec-return-1-2']['container']['resources'])
 
         self.assertEqual(
-            50, dict_format['deploymentSpec']['executors']['exec-return-one-3']
+            50, dict_format['deploymentSpec']['executors']['exec-return-1-3']
             ['container']['resources']['memoryLimit'])
         self.assertNotIn(
             'cpuLimit', dict_format['deploymentSpec']['executors']
-            ['exec-return-one-3']['container']['resources'])
+            ['exec-return-1-3']['container']['resources'])
 
         self.assertEqual(
-            2, dict_format['deploymentSpec']['executors']['exec-return-one-4']
+            2, dict_format['deploymentSpec']['executors']['exec-return-1-4']
             ['container']['resources']['cpuRequest'])
         self.assertEqual(
-            5, dict_format['deploymentSpec']['executors']['exec-return-one-4']
+            5, dict_format['deploymentSpec']['executors']['exec-return-1-4']
             ['container']['resources']['cpuLimit'])
         self.assertEqual(
-            4, dict_format['deploymentSpec']['executors']['exec-return-one-4']
+            4, dict_format['deploymentSpec']['executors']['exec-return-1-4']
             ['container']['resources']['memoryRequest'])
         self.assertEqual(
-            50, dict_format['deploymentSpec']['executors']['exec-return-one-4']
+            50, dict_format['deploymentSpec']['executors']['exec-return-1-4']
             ['container']['resources']['memoryLimit'])
 
 

diff --git a/sdk/python/kfp/dsl/component_factory.py b/sdk/python/kfp/dsl/component_factory.py
@@ -66,13 +66,6 @@ class ComponentInfo():
 
 def _python_function_name_to_component_name(name):
     name_with_spaces = re.sub(' +', ' ', name.replace('_', ' ')).strip(' ')
-    name_list = name_with_spaces.split(' ')
-
-    if name_list[-1].isdigit():
-        raise ValueError(
-            f'Invalid function name "{name}". The function name must not end in `_<int>`.'
-        )
-
     return name_with_spaces[0].upper() + name_with_spaces[1:]
 
 

diff --git a/sdk/python/kfp/dsl/component_factory_test.py b/sdk/python/kfp/dsl/component_factory_test.py
@@ -174,30 +174,6 @@ def comp(Output: OutputPath(str), text: str) -> str:
                 pass
 
 
-class TestPythonFunctionName(unittest.TestCase):
-
-    def test_invalid_function_name(self):
-
-        with self.assertRaisesRegex(
-                ValueError,
-                r'Invalid function name "comp_2". The function name must not end in `_<int>`.'
-        ):
-
-            @component
-            def comp_2(text: str) -> str:
-                pass
-
-    def test_valid_function_name(self):
-
-        @component
-        def comp_v2(text: str) -> str:
-            pass
-
-        @component
-        def comp_(text: str) -> str:
-            pass
-
-
 class TestExtractComponentInterfaceListofArtifacts(unittest.TestCase):
 
     def test_python_component_input(self):

diff --git a/sdk/python/test_data/pipelines/if_elif_else_complex.py b/sdk/python/test_data/pipelines/if_elif_else_complex.py
@@ -18,7 +18,7 @@
 
 
 @dsl.component
-def int_0_to_9999_func() -> int:
+def int_0_to_9999() -> int:
     import random
     return random.randint(0, 9999)
 
@@ -49,7 +49,7 @@ def lucky_number_pipeline(add_drumroll: bool = True,
                           repeat_if_lucky_number: bool = True,
                           trials: List[int] = [1, 2, 3]):
     with dsl.ParallelFor(trials) as trial:
-        int_task = int_0_to_9999_func().set_caching_options(False)
+        int_task = int_0_to_9999().set_caching_options(False)
         with dsl.If(add_drumroll == True):
             with dsl.If(trial == 3):
                 print_and_return(text='Adding drumroll on last trial!')