Merge pull request #1184 from rhayes777/feature/smoke-test-fast-pyautofit

perf: Skip VRAM, I/O, model.info in test mode 2+ and add PYAUTO_DISABLE_JAX

Author: Jammy2211

autofit/mapper/prior_model/abstract.py

@@ -1778,6 +1778,10 @@ def info(self) -> str:
         parameter of the overall model.
         This information is extracted from each priors *model_info* property.
         """
+        from autofit.non_linear.test_mode import test_mode_level
+
+        if test_mode_level() >= 2:
+            return f"Total Free Parameters = {self.prior_count}\n\n[test mode — info skipped]"
 
         formatter = TextFormatter(line_length=info_whitespace())
 

autofit/non_linear/analysis/analysis.py

@@ -35,6 +35,9 @@ class Analysis(ABC):
     def __init__(
         self, use_jax : bool = False, **kwargs
     ):
+        import os
+        if os.environ.get("PYAUTO_DISABLE_JAX") == "1":
+            use_jax = False
 
         self._use_jax = use_jax
         self.kwargs = kwargs
@@ -325,6 +328,11 @@ def print_vram_use(self, model, batch_size : int) -> str:
         batch_size
             The batch size to profile, which is the number of model evaluations JAX will perform simultaneously.
         """
+        from autofit.non_linear.test_mode import test_mode_level
+
+        if test_mode_level() >= 2:
+            return
+
         if not self._use_jax:
             print("use_jax=False for this analysis, therefore does not use GPU and VRAM use cannot be profiled.")
             return

autofit/non_linear/search/abstract_search.py

@@ -497,11 +497,13 @@ class represented by model M and gives a score for their fitness.
         analysis = analysis.modify_before_fit(paths=self.paths, model=model)
         model.unfreeze()
 
-        self.pre_fit_output(
-            analysis=analysis,
-            model=model,
-            info=info,
-        )
+        mode = test_mode_level()
+        if mode < 2:
+            self.pre_fit_output(
+                analysis=analysis,
+                model=model,
+                info=info,
+            )
 
         if not self.paths.is_complete:
             result = self.start_resume_fit(
@@ -514,13 +516,14 @@ class represented by model M and gives a score for their fitness.
                 model=model,
             )
 
-        analysis = analysis.modify_after_fit(
-            paths=self.paths, model=model, result=result
-        )
+        if mode < 2:
+            analysis = analysis.modify_after_fit(
+                paths=self.paths, model=model, result=result
+            )
 
-        self.post_fit_output(
-            search_internal=result.search_internal,
-        )
+            self.post_fit_output(
+                search_internal=result.search_internal,
+            )
 
         gc.collect()
 
@@ -856,8 +859,6 @@ def _fit_bypass_test_mode(
         )
 
         samples_summary = samples.summary()
-        self.paths.save_samples_summary(samples_summary=samples_summary)
-        self.paths.save_samples(samples=samples)
 
         result = analysis.make_result(
             samples_summary=samples_summary,
@@ -866,13 +867,8 @@ def _fit_bypass_test_mode(
             search_internal=None,
         )
 
-        analysis.save_results(paths=self.paths, result=result)
-        analysis.save_results_combined(paths=self.paths, result=result)
-
         model.unfreeze()
 
-        self.paths.completed()
-
         return result
 
     @staticmethod

autofit/text/text_util.py

@@ -55,6 +55,11 @@ def result_info_from(samples) -> str:
     Output the full model.results file, which include the most-likely model, most-probable model at 1 and 3
     sigma confidence and information on the maximum log likelihood.
     """
+    from autofit.non_linear.test_mode import test_mode_level
+
+    if test_mode_level() >= 2:
+        return "[test mode — result info skipped]"
+
     results = []
 
     if hasattr(samples, "log_evidence"):