main: Reorganise so arg parsing and env settings come before tf imports

Author: lisavader

src/geneml/main.py

@@ -21,6 +21,18 @@
     run_model,
 )
 
+args = parse_args()
+params = build_params_namedtuple(args)
+
+if params.cpu_only:
+    os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
+
+os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
+
+# TF and any modules that use it need to be imported after environment variables are set
+import tensorflow as tf # noqa: E402, I001
+from geneml.model_loader import get_cached_gene_ml_model # noqa: E402, I001
+
 logger = logging.getLogger("geneml")
 
 
@@ -88,13 +100,6 @@ def process_contig(contig_id: str, seq: str, params: Params, tensorflow_thread_c
     """
     start_time = time.time()
 
-    import tensorflow as tf
-
-    from geneml.model_loader import get_cached_gene_ml_model
-
-    if params.cpu_only:
-        tf.config.set_visible_devices([], 'GPU')
-
     tf.config.threading.set_inter_op_parallelism_threads(tensorflow_thread_count)
     tf.config.threading.set_intra_op_parallelism_threads(tensorflow_thread_count)
 
@@ -262,13 +267,6 @@ def main() -> None:
     Returns:
         None.
     """
-    args = parse_args()
-    params = build_params_namedtuple(args)
-
-    if params.cpu_only:
-        os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
-        os.environ['TF_CPP_MIN_LOG_LEVEL'] = '1'
-
     logfile = ''.join([params.basepath, '.log'])
     setup_logger(logfile, debug = params.debug, verbose= params.verbose)
     write_setup_info(params)