Example tests to run in parallel

[avital]

Right now we run each example's tests sequentially in tests/run_all_tests.sh (this is due to pytest-dev/pytest#3151). Because of that, we can't run different example tests in parallel (though pytest can do this within a single run).

Consider either finding a way to work around the pytest "can't define two tests with the same filename in different directories", or if that's really not possible, at least try to find a way to rewrite tests/run_all_tests.sh to run in parallel but safely (e.g. can we run pytest multiple times -- does it keep any state directory that we need to isolate?)

[avital]

/cc @Marvin182 @andsteing

[avital]

Marking as "pull requests welcome"

[andsteing]

Current timings from test when run on Github:

- flax/ in 61.14s (0:01:01)
- examples/pixelcnn in 10.56s
- examples/seq2seq in 20.55s
- examples/mnist/ in 9.11s
- examples/ppo/ in 20.88s
- examples/graph/ in 5.15s
- examples/sst2/ in 22.40s
- examples/lm1b/ in 42.18s
- examples/imagenet/ in 141.12s (0:02:21)
- linen_examples/pixelcnn/ in 30.65s
- linen_examples/wmt/ in 94.05s (0:01:34)
- linen_examples/seq2seq/ in 25.71s
- linen_examples/mnist/ in 10.00s
- linen_examples/nlp_seq/ in 2.28s
- linen_examples/imagenet/ in 330.18s (0:05:30)

[andsteing]

@avital is there a way to see Github VM CPU usage during the build workflow? How many cores do we have?

When running locally, expensive tests like wmt and imagenet that make most of the test time max out a single core, so parallelizing would only make sense if we have multiple cores on Github.

[avital]

My guess is that the default is one core, but I'm not sure. Looks like we can use "self-hosted runners" if we want to.

@jheek knows more than me about GitHub actions, maybe he knows how we can find the VM specs for the free GitHub action build machine

[jheek]

The specs of the machine seem to be 2 Cores with 7Gb ram.
We could start using pytest -n auto to automatically parallelize to the number of CPUs available.

[jheek]

I also noticed we are spending most time compiling. I'll try disabling the heavy LLVM optimization in the longer tests

[jheek]

Filed #741 for disabling LLVM optimization in some tests

[marcvanzee]

Closing this since I believe that after #2458 we run our tests in parallel. @cgarciae please let me know if you think I am wrong!