add option to specify computation backend through environment variable (#102)

* add option to specify computation backend through environment variable

* update readme

Author: pmeier

README.md

@@ -103,6 +103,13 @@ In fact, `ltt` is `pip` with a few added options:
   ltt install --pytorch-computation-backend=cu102 torch
   ```
 
+  Borrowing from the mutex packages that PyTorch provides for `conda` installations,
+  `--cpuonly` is available as shorthand for `--pytorch-computation-backend=cu102`.
+
+  In addition, the computation backend to be installed can also be set through the
+  `LTT_PYTORCH_COMPUTATION_BACKEND` environment variable. It will only be honored in
+  case no CLI option for the computation backend is specified.
+
 - By default, `ltt` installs stable PyTorch binaries. To install binaries from the
   nightly, test, or LTS channels pass the `--pytorch-channel` option:
 

light_the_torch/_patch.py

@@ -4,6 +4,7 @@
 import functools
 import itertools
 import optparse
+import os
 import re
 import sys
 import unittest.mock
@@ -145,6 +146,11 @@ def from_pip_argv(cls, argv: List[str]):
             }
         elif opts.cpuonly:
             cbs = {cb.CPUBackend()}
+        elif "LTT_PYTORCH_COMPUTATION_BACKEND" in os.environ:
+            cbs = {
+                cb.ComputationBackend.from_str(string.strip())
+                for string in os.environ["LTT_PYTORCH_COMPUTATION_BACKEND"].split(",")
+            }
         else:
             cbs = cb.detect_compatible_computation_backends()