From 12acc05bf0d4462fd15cdeaf6ea33c4f117b47c1 Mon Sep 17 00:00:00 2001 From: Anushan Fernando Date: Tue, 15 Oct 2024 07:00:30 -0700 Subject: [PATCH] Set env var for torax sim no compile test in test. PiperOrigin-RevId: 686089551 --- torax/tests/sim_no_compile.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/torax/tests/sim_no_compile.py b/torax/tests/sim_no_compile.py index 223d0ae8..6f600fee 100644 --- a/torax/tests/sim_no_compile.py +++ b/torax/tests/sim_no_compile.py @@ -14,6 +14,8 @@ """Tests that TORAX can be run with compilation disabled.""" +import copy +import os from typing import Optional, Sequence from absl.testing import absltest @@ -28,6 +30,17 @@ class SimTest(sim_test_case.SimTestCase): """No-compilation integration tests for torax.sim.""" + def setUp(self): + self.original_flag = copy.deepcopy(os.environ) + compilation_flag = os.environ.get('TORAX_COMPILATION_ENABLED', None) + if compilation_flag is None or compilation_flag in ['1', 'True', 'true']: + os.environ['TORAX_COMPILATION_ENABLED'] = 'False' + super().setUp() + + def tearDown(self): + os.environ = self.original_flag + super().tearDown() + @parameterized.named_parameters( ( 'test_implicit_optimizer_no_compile',