10_mlops/components/train_model_caip.yaml

name: train_model_caip
description: Trains an ML model on CAIP
inputs:
- {name: input_topdir, type: GCSPath, description: 'Top-level directory for TF records'}
- {name: region, type: str, description: 'Region'}
- {name: job_dir, type: GCSPath, description: 'Top-level output directory'}
outputs:
- {name: trained_model, type: GCSPath, description: 'location of trained model'}
implementation:
  container:
    image: gcr.io/ai-analytics-solutions/practical-ml-vision-book:latest
    command: [
        "bash", "/src/practical-ml-vision-book/10_mlops/components/train_model_caip.sh", 
    ]
    args: [
        {inputValue: input_topdir},
        {inputValue: region},
        {inputValue: job_dir},
        {outputPath: trained_model},
    ]