add components to output artifacts from final eval

Signed-off-by: Michael Clifford <[email protected]>

pipeline.py

@@ -78,7 +78,9 @@ def ilab_pipeline_wrapper(mock: List[Literal[MOCKED_STAGES]]):
         )
         from utils import (
             model_to_pvc_op,
+            pvc_to_mmlu_branch_op,
             pvc_to_model_op,
+            pvc_to_mt_bench_branch_op,
             pvc_to_mt_bench_op,
         )
 
@@ -452,6 +454,30 @@ def pipeline(
             mount_path="/output",
         )
 
+        output_mt_bench_branch_task = pvc_to_mt_bench_branch_op(
+            pvc_path="/output/mt_bench_branch/mt_bench_branch_data.json",
+        )
+        output_mt_bench_branch_task.after(final_eval_task)
+        output_mt_bench_branch_task.set_caching_options(False)
+
+        mount_pvc(
+            task=output_mt_bench_branch_task,
+            pvc_name=output_pvc_task.output,
+            mount_path="/output",
+        )
+
+        output_mmlu_branch_task = pvc_to_mmlu_branch_op(
+            pvc_path="/output/mmlu_branch/mmlu_branch_data.json",
+        )
+        output_mmlu_branch_task.after(final_eval_task)
+        output_mmlu_branch_task.set_caching_options(False)
+
+        mount_pvc(
+            task=output_mmlu_branch_task,
+            pvc_name=output_pvc_task.output,
+            mount_path="/output",
+        )
+
         sdg_pvc_delete_task = DeletePVC(pvc_name=sdg_input_pvc_task.output)
         sdg_pvc_delete_task.after(final_eval_task)
 

pipeline.yaml

@@ -329,6 +329,18 @@ components:
           defaultValue: /model
           isOptional: true
           parameterType: STRING
+  comp-pvc-to-mmlu-branch-op:
+    executorLabel: exec-pvc-to-mmlu-branch-op
+    inputDefinitions:
+      parameters:
+        pvc_path:
+          parameterType: STRING
+    outputDefinitions:
+      artifacts:
+        mmlu_branch_output:
+          artifactType:
+            schemaTitle: system.Artifact
+            schemaVersion: 0.0.1
   comp-pvc-to-model-op:
     executorLabel: exec-pvc-to-model-op
     inputDefinitions:
@@ -341,6 +353,18 @@ components:
           artifactType:
             schemaTitle: system.Model
             schemaVersion: 0.0.1
+  comp-pvc-to-mt-bench-branch-op:
+    executorLabel: exec-pvc-to-mt-bench-branch-op
+    inputDefinitions:
+      parameters:
+        pvc_path:
+          parameterType: STRING
+    outputDefinitions:
+      artifacts:
+        mt_bench_branch_output:
+          artifactType:
+            schemaTitle: system.Artifact
+            schemaVersion: 0.0.1
   comp-pvc-to-mt-bench-op:
     executorLabel: exec-pvc-to-mt-bench-op
     inputDefinitions:
@@ -733,6 +757,14 @@ deploymentSpec:
         - /bin/sh
         - -c
         image: registry.access.redhat.com/ubi9/toolbox
+    exec-pvc-to-mmlu-branch-op:
+      container:
+        args:
+        - cp -r {{$.inputs.parameters['pvc_path']}} {{$.outputs.artifacts['mmlu_branch_output'].path}}
+        command:
+        - /bin/sh
+        - -c
+        image: registry.access.redhat.com/ubi9/toolbox
     exec-pvc-to-model-op:
       container:
         args:
@@ -741,6 +773,14 @@ deploymentSpec:
         - /bin/sh
         - -c
         image: registry.access.redhat.com/ubi9/toolbox
+    exec-pvc-to-mt-bench-branch-op:
+      container:
+        args:
+        - cp -r {{$.inputs.parameters['pvc_path']}} {{$.outputs.artifacts['mt_bench_branch_output'].path}}
+        command:
+        - /bin/sh
+        - -c
+        image: registry.access.redhat.com/ubi9/toolbox
     exec-pvc-to-mt-bench-op:
       container:
         args:
@@ -1888,6 +1928,20 @@ root:
                 producerTask: importer
         taskInfo:
           name: model-to-pvc-op
+      pvc-to-mmlu-branch-op:
+        cachingOptions: {}
+        componentRef:
+          name: comp-pvc-to-mmlu-branch-op
+        dependentTasks:
+        - createpvc-3
+        - run-final-eval-op
+        inputs:
+          parameters:
+            pvc_path:
+              runtimeValue:
+                constant: /output/mmlu_branch/mmlu_branch_data.json
+        taskInfo:
+          name: pvc-to-mmlu-branch-op
       pvc-to-model-op:
         cachingOptions: {}
         componentRef:
@@ -1902,6 +1956,20 @@ root:
                 constant: /output/phase_2/model/hf_format/candidate_model
         taskInfo:
           name: pvc-to-model-op
+      pvc-to-mt-bench-branch-op:
+        cachingOptions: {}
+        componentRef:
+          name: comp-pvc-to-mt-bench-branch-op
+        dependentTasks:
+        - createpvc-3
+        - run-final-eval-op
+        inputs:
+          parameters:
+            pvc_path:
+              runtimeValue:
+                constant: /output/mt_bench_branch/mt_bench_branch_data.json
+        taskInfo:
+          name: pvc-to-mt-bench-branch-op
       pvc-to-mt-bench-op:
         cachingOptions: {}
         componentRef:
@@ -2336,12 +2404,24 @@ platforms:
             taskOutputParameter:
               outputParameterKey: name
               producerTask: createpvc-2
+        exec-pvc-to-mmlu-branch-op:
+          pvcMount:
+          - mountPath: /output
+            taskOutputParameter:
+              outputParameterKey: name
+              producerTask: createpvc-3
         exec-pvc-to-model-op:
           pvcMount:
           - mountPath: /output
             taskOutputParameter:
               outputParameterKey: name
               producerTask: createpvc-3
+        exec-pvc-to-mt-bench-branch-op:
+          pvcMount:
+          - mountPath: /output
+            taskOutputParameter:
+              outputParameterKey: name
+              producerTask: createpvc-3
         exec-pvc-to-mt-bench-op:
           pvcMount:
           - mountPath: /output

utils/__init__.py

@@ -2,13 +2,17 @@
 from .components import (
     ilab_importer_op,
     model_to_pvc_op,
+    pvc_to_mmlu_branch_op,
     pvc_to_model_op,
+    pvc_to_mt_bench_branch_op,
     pvc_to_mt_bench_op,
 )
 
 __all__ = [
     "model_to_pvc_op",
     "pvc_to_mt_bench_op",
+    "pvc_to_mt_bench_branch_op",
+    "pvc_to_mmlu_branch_op",
     "pvc_to_model_op",
     "ilab_importer_op",
     "faked",

utils/components.py

@@ -15,6 +15,26 @@ def pvc_to_mt_bench_op(mt_bench_output: dsl.Output[dsl.Artifact], pvc_path: str)
     )
 
 
+@dsl.container_component
+def pvc_to_mt_bench_branch_op(
+    mt_bench_branch_output: dsl.Output[dsl.Artifact], pvc_path: str
+):
+    return dsl.ContainerSpec(
+        TOOLBOX_IMAGE,
+        ["/bin/sh", "-c"],
+        [f"cp -r {pvc_path} {mt_bench_branch_output.path}"],
+    )
+
+
+@dsl.container_component
+def pvc_to_mmlu_branch_op(mmlu_branch_output: dsl.Output[dsl.Artifact], pvc_path: str):
+    return dsl.ContainerSpec(
+        TOOLBOX_IMAGE,
+        ["/bin/sh", "-c"],
+        [f"cp -r {pvc_path} {mmlu_branch_output.path}"],
+    )
+
+
 @dsl.container_component
 def pvc_to_model_op(model: dsl.Output[dsl.Model], pvc_path: str):
     return dsl.ContainerSpec(