Fix: Set OMP_NUM_THREADS by default in Elastic

[fellhorn]

Why are the changes needed?

torchrun sets the environment variable OMP_NUM_THREADS automatically if not specified (code). If it is not set, we saw some special cases in which our experiments were stuck for several minutes with high CPU load.

Also see https://pytorch.org/tutorials/recipes/recipes/tuning_guide.html#utilize-openmp for more details.

To make Elastic tasks behave the same as executions started with torchrun, I would propose copying this behavior over to flytekit.

What changes were proposed in this pull request?

See above

How was this patch tested?

1. Added unit tests
2. Tested with internal workflows to check the behavior of OMP_NUM_THREADS.

Check all the applicable boxes

• I updated the documentation accordingly.
• All new and existing tests passed.
• All commits are signed-off.

Related PRs

Docs link

[thomasjpfan]

Although this is changing behavior of current code, I think it's a better default.

For my curiosity, how are coming across this issue with PyTorch? (PyTorch's DataLoader also sets the number of threads to one here so it does not oversubscribe with num_workers.)

Can we add to Elastic's docstring, how a user should set OMP_NUM_THREADS? Do you see users doing this:

# my_workflow.py
import os

os.environ["OMP_NUM_THREADS"] = "2"

@task(task_config=Elastic(...))
def my_task():
    ....

[kumare3]

@fellhorn i like @thomasjpfan's question. @thomasjpfan it should be set using
'Pyflyte run --env'

[fellhorn]

@fellhorn i like @thomasjpfan's question. @thomasjpfan it should be set using 'Pyflyte run --env'

Good point, I added some information now in 95a10f9:

Like torchrun, this plugin sets the environment variable OMP_NUM_THREADS to 1 if it is not set.
Please see https://pytorch.org/tutorials/recipes/recipes/tuning_guide.html for potential performance improvements.
To change OMP_NUM_THREADS, specify it in the environment dict of the flytekit task decorator or via pyflyte run --env.

[fellhorn]

Although this is changing behavior of current code, I think it's a better default.

For my curiosity, how are coming across this issue with PyTorch? (PyTorch's DataLoader also sets the number of threads to one here so it does not oversubscribe with num_workers.)

Can we add to Elastic's docstring, how a user should set OMP_NUM_THREADS? Do you see users doing this:

# my_workflow.py
import os

os.environ["OMP_NUM_THREADS"] = "2"

@task(task_config=Elastic(...))
def my_task():
    ....

@thomasjpfan regarding your other question: OpenMP is e.g. also used for torch.jit.script inference, where no default value for the number of threads is set.