examples/demos/simple_pipeline/gpu-example-pipeline.yaml

apiVersion: argoproj.io/v1alpha1
kind: Workflow
metadata:
  generateName: pipeline-gpu-example-
spec:
  arguments:
    parameters:
    - name: learning-rate
      value: '0.1'
    - name: num-layers
      value: '2'
    - name: optimizer
      value: ftrl
  entrypoint: pipeline-gpu-example
  serviceAccountName: pipeline-runner
  templates:
  - dag:
      tasks:
      - arguments:
          parameters:
          - name: training-output
            value: '{{tasks.training.outputs.parameters.training-output}}'
        dependencies:
        - training
        name: postprocessing
        template: postprocessing
      - arguments:
          parameters:
          - name: learning-rate
            value: '{{inputs.parameters.learning-rate}}'
          - name: num-layers
            value: '{{inputs.parameters.num-layers}}'
          - name: optimizer
            value: '{{inputs.parameters.optimizer}}'
        name: training
        template: training
    inputs:
      parameters:
      - name: learning-rate
      - name: num-layers
      - name: optimizer
    name: pipeline-gpu-example
  - container:
      args:
      - echo "{{inputs.parameters.training-output}}"
      command:
      - sh
      - -c
      image: library/bash:4.4.23
    inputs:
      parameters:
      - name: training-output
    name: postprocessing
    outputs:
      artifacts:
      - name: mlpipeline-ui-metadata
        path: /mlpipeline-ui-metadata.json
        s3:
          accessKeySecret:
            key: accesskey
            name: mlpipeline-minio-artifact
          bucket: mlpipeline
          endpoint: minio-service.kubeflow:9000
          insecure: true
          key: runs/{{workflow.uid}}/{{pod.name}}/mlpipeline-ui-metadata.tgz
          secretKeySecret:
            key: secretkey
            name: mlpipeline-minio-artifact
      - name: mlpipeline-metrics
        path: /mlpipeline-metrics.json
        s3:
          accessKeySecret:
            key: accesskey
            name: mlpipeline-minio-artifact
          bucket: mlpipeline
          endpoint: minio-service.kubeflow:9000
          insecure: true
          key: runs/{{workflow.uid}}/{{pod.name}}/mlpipeline-metrics.tgz
          secretKeySecret:
            key: secretkey
            name: mlpipeline-minio-artifact
  - container:
      args:
      - --batch-size
      - '64'
      - --lr
      - '{{inputs.parameters.learning-rate}}'
      - --num-layers
      - '{{inputs.parameters.num-layers}}'
      - --optimizer
      - '{{inputs.parameters.optimizer}}'
      command:
      - python
      - /mxnet/example/image-classification/train_mnist.py
      image: katib/mxnet-mnist-example
      resources:
        limits:
          nvidia.com/gpu: 1
    inputs:
      parameters:
      - name: learning-rate
      - name: num-layers
      - name: optimizer
    name: training
    outputs:
      artifacts:
      - name: mlpipeline-ui-metadata
        path: /mlpipeline-ui-metadata.json
        s3:
          accessKeySecret:
            key: accesskey
            name: mlpipeline-minio-artifact
          bucket: mlpipeline
          endpoint: minio-service.kubeflow:9000
          insecure: true
          key: runs/{{workflow.uid}}/{{pod.name}}/mlpipeline-ui-metadata.tgz
          secretKeySecret:
            key: secretkey
            name: mlpipeline-minio-artifact
      - name: mlpipeline-metrics
        path: /mlpipeline-metrics.json
        s3:
          accessKeySecret:
            key: accesskey
            name: mlpipeline-minio-artifact
          bucket: mlpipeline
          endpoint: minio-service.kubeflow:9000
          insecure: true
          key: runs/{{workflow.uid}}/{{pod.name}}/mlpipeline-metrics.tgz
          secretKeySecret:
            key: secretkey
            name: mlpipeline-minio-artifact
      parameters:
      - name: training-output
        valueFrom:
          path: /etc/timezone