Improve logic and update tests

Signed-off-by: droctothorpe <[email protected]> Co-authored-by: zazulam <[email protected]>
kubeflow · Aug 5, 2024 · f962ef3 · f962ef3
1 parent 65ce19f
commit f962ef3
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 14 deletions.
diff --git a/backend/src/v2/compiler/argocompiler/argo.go b/backend/src/v2/compiler/argocompiler/argo.go
@@ -278,14 +278,25 @@ func (c *workflowCompiler) extractFunctionName(componentName string) string {
 	}
 	executorLabel := c.spec.Components[componentName].GetExecutorLabel()
 	log.Debug("executorLabel: ", executorLabel)
-	for executorName, executorValue := range c.executors {
-		log.Debug("executorName: ", executorName)
-		if executorName == executorLabel {
-			args := executorValue.GetContainer().Args
-			componentFunctionName := args[len(args)-1]
-			log.Debug("componentFunctionName: ", componentFunctionName)
-
-			return componentFunctionName
+	// There are more nested conditionals here than we would prefer, but we
+	// don't want to make any assumptions about the presence of specific fields
+	// in the IR.
+	if c.executors != nil {
+		for executorName, executorValue := range c.executors {
+			log.Debug("executorName: ", executorName)
+			if executorName == executorLabel {
+				args := executorValue.GetContainer().GetArgs()
+				if args != nil {
+					if len(args) > 1 {
+						penultimateArg := args[len(args)-2]
+						if penultimateArg == "--function_to_execute" {
+							componentFunctionName := args[len(args)-1]
+							log.Debug("componentFunctionName: ", componentFunctionName)
+							return componentFunctionName
+						}
+					}
+				}
+			}
 		}
 	}
 

diff --git a/backend/src/v2/compiler/argocompiler/testdata/create_mount_delete_dynamic_pvc.yaml b/backend/src/v2/compiler/argocompiler/testdata/create_mount_delete_dynamic_pvc.yaml
@@ -8,9 +8,9 @@ spec:
     parameters:
     - name: kubernetes-comp-comp
       value: '{"pvcMount":[{"mountPath":"/data","taskOutputParameter":{"outputParameterKey":"name","producerTask":"createpvc"}}]}'
-    - name: components-comp-comp
+    - name: components-comp
       value: '{"executorLabel":"exec-comp"}'
-    - name: implementations-comp-comp
+    - name: implementations-comp
       value: '{"args":["--executor_input","{{$}}","--function_to_execute","comp"],"command":["sh","-c","\nif
         ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip || python3
         -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1
@@ -197,11 +197,11 @@ spec:
       - arguments:
           parameters:
           - name: component
-            value: '{{workflow.parameters.components-comp-comp}}'
+            value: '{{workflow.parameters.components-comp}}'
           - name: task
             value: '{"cachingOptions":{"enableCache":true},"componentRef":{"name":"comp-comp"},"dependentTasks":["createpvc"],"taskInfo":{"name":"comp"}}'
           - name: container
-            value: '{{workflow.parameters.implementations-comp-comp}}'
+            value: '{{workflow.parameters.implementations-comp}}'
           - name: parent-dag-id
             value: '{{inputs.parameters.parent-dag-id}}'
           - name: kubernetes-config
@@ -222,11 +222,11 @@ spec:
       - arguments:
           parameters:
           - name: component
-            value: '{{workflow.parameters.components-comp-comp}}'
+            value: '{{workflow.parameters.components-comp}}'
           - name: task
             value: '{"cachingOptions":{"enableCache":true},"componentRef":{"name":"comp-comp-2"},"dependentTasks":["comp","createpvc"],"taskInfo":{"name":"comp-2"}}'
           - name: container
-            value: '{{workflow.parameters.implementations-comp-comp}}'
+            value: '{{workflow.parameters.implementations-comp}}'
           - name: parent-dag-id
             value: '{{inputs.parameters.parent-dag-id}}'
           - name: kubernetes-config