Move env set-up from test/run_tests.sh to individual .py files

[zhanyong-wan]

🚀 Feature

Currently the python test files cannot be run directly - they need to be invoked via test/run_tests.sh, which sets up the test environment before running the python file, and the environment varies from test to test.

This makes debugging complex. If we move the set-up logic from .sh to .py, it will make debugging and running the tests easier. It also makes the test files self-sufficient.

[yaoshiang]

This would also enable integration with VSCode's visual python debugging of unit tests.

[ysiraichi]

Could you give an example?

[zhanyong-wan]

Basically, we need to:

1. reverse engineer what test/run_tests.sh does before invoking a test*.py file.
2. replicate the logic inside the test*.py file itself.

[ysiraichi]

For example, test_operations.py is run with the following configurations:

• Normal (no special env variable is set)
• XLA_DISABLE_FUNCTIONALIZATION=1
• XLA_USE_EAGER_DEBUG_MODE=1
• XLA_EXPERIMENTAL="nonzero:masked_select:masked_scatter:nms"

How would we change this test?

[zhanyong-wan]

Here's an idea:

1. keep test_operations.py for the normal case.
2. add test_operations_no_functionalization.py, which imports test_operations, sets the env var, and runs the tests.
3. similarly, add test_operations_eager_debug.py and test_operations_masked.py.

Now, each of the 4 *.py files will be self sufficient (not depending on any .sh set-up).

Some refactoring may be needed to make this happen.

[ysiraichi]

This is an interesting idea. However, wouldn't it generate an unnecessarily large amount of files? Instead, what do you think of doing something like PyTorch does for testing dynamic shapes?

In summary, instead of having one test_*_no_functionalization.py for each file we run, we would have one test_no_functionalization.py file that would call the other tests from the other files with its environment set up. Then, we wouldn't have to create a new file for each combination we want to run.

While we are at it, maybe it's a good opportunity to turn many of these environment variable configurations into their own configuration variable inside PyTorch/XLA. This would make testing more convenient (but, perhaps, bc-breaking).

[zhanyong-wan]

Tests need to be obviously correct. Therefore I want the test set up logic to be as straightforward as possible. If we follow the standard python unittest pattern, it will be very easy to figure out which test cases there are and how to select which ones to run. For example, the standard pattern of test_foo.py FooTest.test_abc will run the test_abc test case in test class FooTest defined in test_foo.py. Deviating from this pattern will increase the effort needed by a developer to debug a test failure.

If we follow this pytorch example, can we still easily find out which test cases the test_*.py file contains and easily filter them? If yes, I'd be fine with it.