From 4780697747274035f86999756e186247fe6baed8 Mon Sep 17 00:00:00 2001
From: Jammy2211 <JNightingale2211@gmail.com>
Date: Mon, 6 Apr 2026 22:00:50 +0100
Subject: [PATCH] perf: skip VRAM estimation, I/O, model.info, and result.info
 in test mode 2+

When PYAUTOFIT_TEST_MODE >= 2:
- print_vram_use() returns immediately (saves ~16s)
- pre_fit_output() and post_fit_output() are skipped (saves ~4s)
- _fit_bypass_test_mode() skips all file I/O (saves ~6s)
- model.info returns a stub (saves ~7s for MGE models)
- result_info_from() returns a stub (saves ~5s)

Also adds PYAUTO_DISABLE_JAX=1 env var to force use_jax=False
in Analysis base class, avoiding JAX JIT overhead in smoke tests.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 autofit/mapper/prior_model/abstract.py       |  4 +++
 autofit/non_linear/analysis/analysis.py      |  8 +++++
 autofit/non_linear/search/abstract_search.py | 32 +++++++++-----------
 autofit/text/text_util.py                    |  5 +++
 4 files changed, 31 insertions(+), 18 deletions(-)

diff --git a/autofit/mapper/prior_model/abstract.py b/autofit/mapper/prior_model/abstract.py
index bce4b0d3d..2e1f9f3ce 100644
--- a/autofit/mapper/prior_model/abstract.py
+++ b/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())
 
diff --git a/autofit/non_linear/analysis/analysis.py b/autofit/non_linear/analysis/analysis.py
index 794d32361..572acab7e 100644
--- a/autofit/non_linear/analysis/analysis.py
+++ b/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
diff --git a/autofit/non_linear/search/abstract_search.py b/autofit/non_linear/search/abstract_search.py
index 6bda83675..8e9f0a43b 100644
--- a/autofit/non_linear/search/abstract_search.py
+++ b/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
diff --git a/autofit/text/text_util.py b/autofit/text/text_util.py
index ea649c242..5f2945d90 100644
--- a/autofit/text/text_util.py
+++ b/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"):