workflow.yaml

apiVersion: argoproj.io/v1alpha1
kind: Workflow
metadata:
  generateName:  mlapp-
  labels:
    workflow: mlapp 
spec:
  entrypoint: mlapp 
  templates:
  - name: mlapp 
    dag:
      tasks:
      - name: preprocess
        template: preprocess

      - name: model1
        dependencies: [preprocess]
        template: train
        arguments:
          artifacts:
          - name: dataset
            from: "{{tasks.preprocess.outputs.artifacts.dataset}}"

  - name: preprocess
    container:
      image: analytics-kubelet-055005d.private-us-east-1.github.net:30550/mlapp:latest
      name: preprocess
      command: ["python", "/src/preprocess.py"]
      env:
        - name: DASK_SCHEDULER_ADDRESS
          value: my-dask-scheduler:8786
    outputs:
     artifacts:
     - name: dataset
       path: /data

  - name: train
    inputs:
      artifacts:
      - name: dataset
        path: /data
    outputs:
     artifacts:
     - name: model
       path: /output
    container:
      image: analytics-kubelet-055005d.private-us-east-1.github.net:30550/mlapp-gpu:latest
      name: trainer
      resources:
        requests:
          nvidia.com/gpu: 1 # requesting 1 GPU
        limits:
          nvidia.com/gpu: 1
      command: ["python", "/src/train.py"]