address pr comments

Signed-off-by: zazulam <[email protected]>

backend/src/v2/compiler/argocompiler/argo.go

@@ -262,15 +262,14 @@ func (c *workflowCompiler) argumentsPlaceholder(componentName string) (string, e
 	return workflowParameter(componentName), nil
 }
 
-// extractBaseComponentName removes the iteration suffix that the IR compiler
+// ExtractBaseComponentName removes the iteration suffix that the IR compiler
 // adds to the component name.
 func ExtractBaseComponentName(componentName string) string {
 	baseComponentName := componentName
 	componentNameArray := strings.Split(componentName, "-")
 
 	if _, err := strconv.Atoi(componentNameArray[len(componentNameArray)-1]); err == nil {
 		baseComponentName = strings.Join(componentNameArray[:len(componentNameArray)-1], "-")
-
 	}
 
 	return baseComponentName

backend/src/v2/compiler/argocompiler/argo_test.go

@@ -154,6 +154,21 @@ func Test_extractBaseComponentName(t *testing.T) {
 			componentName:    "component",
 			expectedBaseName: "component",
 		},
+		{
+			name:             "Last char is int",
+			componentName:    "component-v2",
+			expectedBaseName: "component-v2",
+		},
+		{
+			name:             "Multiple dashes, ends with int",
+			componentName:    "service-api-v2",
+			expectedBaseName: "service-api-v2",
+		},
+		{
+			name:             "Multiple dashes and ints",
+			componentName:    "module-1-2-3",
+			expectedBaseName: "module-1-2",
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {

sdk/python/kfp/compiler/compiler_test.py

@@ -129,7 +129,7 @@ def comp():
 
 
 @dsl.component
-def return_1() -> int:
+def return_one() -> int:
     return 1
 
 

sdk/python/kfp/dsl/component_factory.py

@@ -66,6 +66,13 @@ 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:]
 
 

sdk/python/kfp/dsl/component_factory_test.py

@@ -174,6 +174,30 @@ 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):