How do I use an input argument to my Python script as the output artifact in Kubeflow Pipeline

Question

I'm currently using Kubeflow Pipeline to train a number of models. The crux of the issue is that I want to use a saved checkpoint as the input to a subsequent model.

I was originally running training using a Python script inside a project repository and am using the Docker image of this repository as an entire component. Let's call this image model_training:latest. The component using this image is implemented using a YAML file in the pipeline.

Within my image there's a Python script that I'll call train.py. train.py receives a number of arguments, among them one called pretrained_path. If this is True, then the model will be loaded from that path.

My YAML file looks something like this:

name: model-train
inputs:
  - {name: random_seed, type: Integer}
  - {name: pretrained_path, type: String}
outputs:
  - {name: pretrained_path, type: String}
implementation:
  container:
    image: model_training:latest
    command:
      - accelerate
      - launch
      - --multi_gpu
      - --num_processes
      - "2"
      - /repo_name/train.py
    args:
      - --random_seed
      - {inputValue: random_seed}
      - --pretrained_path
      - {outputPath: pretrained_path}

The working directory for this component was set to /repo_name.

When I compile the pipeline using this YAML file for the component, the KFP UI is telling me that pretrained_path isn't an output parameter or artifact, but an "Input Artifact." What does this mean, and how would I tell the component that I need to output pretrained_path?

How do I use an input argument to my Python script as the output artifact in Kubeflow Pipeline

0 Answers0