How to use interpolation at hydra.job.env_set

[YoshikiKubotani]

Hi, hydra lovers. I have one question.

I am now working on the shared server that has multiple devices.
And I would like to limit the number of used devices by setting CUDA_VISIBLE_DEVICES.

So far, I set this environment variable using os.environ['CUDA_VISIBLE_DEVICES'], but I found on official page that hydra can set environment variables from hydra.job.env_set in a config file.

Then, I tried running using a config file as follows.

device: 1

hydra:
  run:
    dir: outputs/${device}
  job:
    env_set:
      CUDA_VISIBLE_DEVICES: "${device}"

But it didn't work, and I got an error below.

Interpolation key 'device' not found
    full_key: hydra.job.env_set.CUDA_VISIBLE_DEVICES
    reference_type=Dict[str, str]
    object_type=dict

I confirmed that I could run a code by setting "1" explicitly instead of "${device}".
I thought we could use interpolation in the same way when setting hydra.run.dir as the example above, but is my understanding wrong?

[Jasha10]

Hi @YoshikiKubotani,

Some parts of the top-level hydra config are computed early on, before the user job's config is computed.
As an example of why this is necessary, the hydra.searchpath key will influence how the user's job config is composed, so it's necessary to compute hydra.searchpath before beginning composition of the job's config.

My suspicion is that Hydra is trying to read from the hydra.job.env_set variable before the job config has been computed (and thus before the device key has been set to 1).

[YoshikiKubotani]

Hi @Jasha10

Thank you for your reply!
Let me make the point clear; do you mean I can't achieve what I want to do with the current hydra implementation due to its resolving order?
If so, do you have any ideas for a workaround?

[Jasha10]

do you mean I can't achieve what I want to do with the current hydra implementation due to its resolving order? If so, do you have any ideas for a workaround?

Yes, Hydra's resolving order means that you will need to use a workaround.
Please give this a try:

# main.yaml
device: ${hydra:job.env_set.CUDA_VISIBLE_DEVICES}

hydra:
  run:
    dir: outputs/${device}
  job:
    env_set:
      CUDA_VISIBLE_DEVICES: 1

# app.py
import os
import hydra

@hydra.main(config_name="main", config_path=".", version_base=None)
def app(cfg):
    print(f"{cfg=}")
    print(f"{cfg.device=}")
    print(f"{os.environ['CUDA_VISIBLE_DEVICES']=}")

app()

$ python app.py
cfg={'device': '${hydra:job.env_set.CUDA_VISIBLE_DEVICES}'}
cfg.device='1'
os.environ['CUDA_VISIBLE_DEVICES']='1'

This uses the hydra: resolver.

[YoshikiKubotani]

@Jasha10
It worked without any errors!
I did not have the idea of the reversed reference. Thank you very much!!