dvc.yaml

# dvc parameters listed under `vars` will not be tracked
#  which is used to list settings that do not affect the
#  training or evaluation of the model, e.g. which logger
#  has been used.
vars:
  - logger: mlflow
stages:
  mllam_version:
    cmd: source machines/environment.sh; python -c "from neural_lam import __version__; print(__version__)" > version.mllam.txt
  prepare_dataset:
    cmd: sbatch -W machines/slurm.mllam-data-prep.sh data/datastore.yaml --output data/datastore.zarr
    deps:
    - data/datastore.yaml
    outs:
    - data/datastore.zarr
  create_graph:
    cmd: python -m neural_lam.create_graph --config_path data/config.yaml ${graph-defaults}
    params:
    - data/config.yaml:
      - datastore
    deps:
    - data/datastore.zarr
    outs:
    - data/graph/
  train:
    cmd:
    - sbatch -W ${slurm-defaults} machines/slurm.neural-lam.sh --config_path ./data/config.yaml --logger ${logger} --logger_project ${logger_project} ${baseline-defaults}
    deps:
    - version.mllam.txt
    - data/graph/
    - data/config.yaml
    - data/datastore.yaml
    outs:
    - saved_models
  evaluate:
    cmd:
    - sbatch -W ${slurm-eval} machines/slurm.neural-lam.sh --config_path ./data/config.yaml --eval val --load
      ./saved_models/*/last.ckpt --logger ${logger} ${logger_project} ${baseline-eval}
    deps:
    - version.mllam.txt
    - data/graph/
    - data/config.yaml
    - data/datastore.yaml
    - saved_models
    outs:
    - mlruns