Can xm.xla_device still work at 8-core even when already called?

[tmabraham]

So I am aware that xm.xla_device() needs to be called in the spawned function and if it is called earlier (before spawning of the multiple processes) then the system assumes there is only one device. I am wondering if it is possible to have a feature where it is possible to call xm.xla_device() before spawning (for example to do single-core experiments) and in the spawned processes (to do multiple core experiments. Like can there be a mode for xm.xla_device() indicating single-core vs multiple-core. I ask because, currently it is impossible to demonstrate single-core and multiple-core functionality in the same notebook/program, and it would be nice to do so for tutorial purposes.

[dlibenzi]

Cannot.
The issue is creating all the environment linked to a computation client, and forking.
Not worth the added complexity.
Just set nprocs=1 and you get single core.
With nprocs=1 there is not even a fork happening, just a function call.
And has the advantage for the code to be the same.
IMHO we should not teach single vs. multi core, as this might lead to users having to rewrite code.
Single core as particular case of N core is the way to lay down the teaching.

[tmabraham]

Thank you for your response.

IMHO we should not teach single vs. multi core, as this might lead to users having to rewrite code.
Single core as particular case of N core is the way to lay down the teaching.

I think the confusion when teaching single core vs multi-core, is that the package usage is slightly different. This is what is discussed in the official documentation (here). If we want to highlight both of these usages (ex: my kernel here) then it is impossible to do so in the same environment. I guess maybe there should be less focus in the documentation regarding a difference between single core and multiple core training. Instead, maybe the documentation should focus on just the multi-core training and mention single core as a special example as you mention. I might also rewrite my kernel accordingly.

[dlibenzi]

Yeah, I know our documentation has been written in that way, and should probably be revised to avoid users investing code in single core, which needs some rewriting for multi core.
Everything would still be fine if we did not have to use fork as spawn method, due to Colab limitations.
Nevertheless, the correct teaching should not differentiate single vs multi core.

[tmabraham]

@dlibenzi Thank you for the clarification!

[dlibenzi]

There was an issue (just a Colab warning of using sys.exit()) with nprocs=1 which has been fixed with #2281