Extract current running stage from dvc

Question

I'm conducting an experiment using 'dvc repro -f', where multiple stages are executed according to the dvc.yaml configuration. For instance:

Stages:
 Training:
   foreach:
    -cycle: 0
    -cycle: 1
    -cycle: 2
   do:
    cmd:
     python train.py
 Selection:
   foreach:
    -cycle: 0
    -cycle: 1
    -cycle: 2
   do:
    cmd:
     python train.py

During the execution of each stage, such as the training stage at cycle 0, I aim to extract its stage name, like 'Training_0', within a Python program and during selection stage it should be Selection_0. I'm seeking a method to extract this information either while the stage is being executed or just before its execution begins. i tried using the dvc.api but the api does not return the current stage that is running. How can I achieve this?

Answer 1

You can pass it to the train.py as an arg:

stages:
 Training:
   foreach:
    -cycle: 0
    -cycle: 1
    -cycle: 2
   do:
    cmd:
     python train.py "Training-${item.cycle}"
 Selection:
   foreach:
    -cycle: 0
    -cycle: 1
    -cycle: 2
   do:
    cmd:
     python train.py "Selection-${item.cycle}"

Would something like this work for you? It would be a nice addition to have this in the dvc.api and probably as an ENV variable. It feels it should be straightforward to add. Feel free to help us do this - all contributions are welcome and we can help :)