pipeline.yml

$schema: https://azuremlschemas.azureedge.net/latest/pipelineJob.schema.json
type: pipeline

description: Pipeline using distributed job to train model based on cifar-10 dataset

display_name: cifar-10-pipeline-example
experiment_name: cifar-10-pipeline-example
jobs:
  get_data:
    type: command
    command: >-
      wget https://azuremlexamples.blob.core.windows.net/datasets/cifar-10-python.tar.gz;
      tar -xvzf cifar-10-python.tar.gz -C ${{outputs.cifar}};
      rm cifar-10-python.tar.gz;
    compute: azureml:gpu-cluster
    environment: azureml://registries/azureml/environments/sklearn-1.1/labels/latest
    outputs:
      cifar:
        type: uri_folder
        mode: upload
  train_model:
    type: command
    command: >-
      python main.py
      --data-dir ${{inputs.cifar}}
      --epochs ${{inputs.epochs}}
      --model-dir ${{outputs.model_dir}}
    code: src/train-model
    inputs:
      epochs: 1
      cifar: ${{parent.jobs.get_data.outputs.cifar}}
    outputs:
      model_dir:
        type: uri_folder
        mode: upload
    environment: azureml://registries/azureml/environments/acpt-pytorch-2.2-cuda12.1/labels/latest
    compute: azureml:gpu-cluster
    distribution:
      type: pytorch
      process_count_per_instance: 1
    resources:
      instance_count: 2
  eval_model:
    type: command
    command: >-
      python main.py
      --data-dir ${{inputs.cifar}}
      --model-dir ${{inputs.model_dir}}/model
    code: src/eval-model
    environment: azureml://registries/azureml/environments/acpt-pytorch-2.2-cuda12.1/labels/latest
    compute: azureml:gpu-cluster
    distribution:
      type: pytorch
      process_count_per_instance: 1
    resources:
      instance_count: 2
    inputs:
      cifar: ${{parent.jobs.get_data.outputs.cifar}}
      model_dir: ${{parent.jobs.train_model.outputs.model_dir}}