Cannot set XLA flags on docker

[neel04]

I had a RESOURCE_EXHAUSTED: Compilation failure: Aborting compilation early because it's unlikely to have enough device memory that I wanted to debug. It advised setting this flag:

--xla_tpu_impure_oom_fast_exit_threshold=-1

but despite setting this flag using -e option, -e xla_tpu_impure_oom_fast_exit_threshold=-1 it still doesn't provide any of the promised "detailed logging info".

Anyone have any ideas? Seems like its just not detecting the flag. I've tried export it in the bash script as well, but that's outside the docker container so doesn't work 🤷

[jakevdp]

Try using -e XLA_FLAGS="--xla_tpu_impure_oom_fast_exit_threshold=-1"; XLA_FLAGS is the name of the environment variable you can use to specify flags you'd like to pass to the xla compiler.

[neel04]

Thanks! though in this case, it complained the flag wasn't found - so the correct one was -e LIBTPU_INIT_ARGS="--xla_tpu_impure_oom_fast_exit_threshold=-1"