Disable JITs in JAX and PyTorch to enable effective profiling (#892)

* Disable JIT in PyTorch and JAX to improve profiling (but don't override existing environment variables).

* Disable JIT.

* Documented.

Author: emeryberger

scalene/scalene_preload.py

@@ -19,6 +19,16 @@ def get_preload_environ(args: argparse.Namespace) -> Dict[str, str]:
             )
         }
 
+        # Disable JITting in PyTorch and JAX to improve profiling,
+        # unless the environment variables are already set.
+        # JAX_DISABLE_JIT: https://jax.readthedocs.io/en/latest/debugging/flags.html#id1
+        # PYTORCH_JIT: https://pytorch.org/docs/stable/jit.html#disable-jit-for-debugging
+        jit_flags = [ ('JAX_DISABLE_JIT', '1'), # truthy => disable JIT
+                      ('PYTORCH_JIT', '0') ]    # falsy => disable JIT
+        for name, val in jit_flags:
+            if name not in os.environ:
+                env[name] = val
+
         # Set environment variables for loading the Scalene dynamic library,
         # which interposes on allocation and copying functions.
         if sys.platform == "darwin":