From 0534d50af35df43452c2083876991e764f0672ea Mon Sep 17 00:00:00 2001
From: "Alinson S. Xavier" <git@axavier.org>
Date: Mon, 26 Feb 2024 16:41:50 -0600
Subject: [PATCH] BasicCollector: Do not crash on exception

---
 miplearn/collectors/basic.py | 98 +++++++++++++++++++-----------------
 miplearn/io.py               |  5 +-
 2 files changed, 56 insertions(+), 47 deletions(-)

diff --git a/miplearn/collectors/basic.py b/miplearn/collectors/basic.py
index 697fbb3..9ed6b03 100644
--- a/miplearn/collectors/basic.py
+++ b/miplearn/collectors/basic.py
@@ -9,6 +9,7 @@ import sys
 from io import StringIO
 from os.path import exists
 from typing import Callable, List, Any
+import traceback
 
 from ..h5 import H5File
 from ..io import _RedirectOutput, gzip, _to_h5_filename
@@ -29,52 +30,57 @@ class BasicCollector:
         verbose: bool = False,
     ) -> None:
         def _collect(data_filename: str) -> None:
-            h5_filename = _to_h5_filename(data_filename)
-            mps_filename = h5_filename.replace(".h5", ".mps")
-
-            if exists(h5_filename):
-                # Try to read optimal solution
-                mip_var_values = None
-                try:
-                    with H5File(h5_filename, "r") as h5:
-                        mip_var_values = h5.get_array("mip_var_values")
-                except:
-                    pass
-
-                if mip_var_values is None:
-                    print(f"Removing empty/corrupted h5 file: {h5_filename}")
-                    os.remove(h5_filename)
-                else:
-                    return
-
-            with H5File(h5_filename, "w") as h5:
-                streams: List[Any] = [StringIO()]
-                if verbose:
-                    streams += [sys.stdout]
-                with _RedirectOutput(streams):
-                    # Load and extract static features
-                    model = build_model(data_filename)
-                    model.extract_after_load(h5)
-
-                    if not self.skip_lp:
-                        # Solve LP relaxation
-                        relaxed = model.relax()
-                        relaxed.optimize()
-                        relaxed.extract_after_lp(h5)
-
-                    # Solve MIP
-                    model.optimize()
-                    model.extract_after_mip(h5)
-
-                    if self.write_mps:
-                        # Add lazy constraints to model
-                        model._lazy_enforce_collected()
-
-                        # Save MPS file
-                        model.write(mps_filename)
-                        gzip(mps_filename)
-
-                h5.put_scalar("mip_log", streams[0].getvalue())
+            try:
+                h5_filename = _to_h5_filename(data_filename)
+                mps_filename = h5_filename.replace(".h5", ".mps")
+
+                if exists(h5_filename):
+                    # Try to read optimal solution
+                    mip_var_values = None
+                    try:
+                        with H5File(h5_filename, "r") as h5:
+                            mip_var_values = h5.get_array("mip_var_values")
+                    except:
+                        pass
+
+                    if mip_var_values is None:
+                        print(f"Removing empty/corrupted h5 file: {h5_filename}")
+                        os.remove(h5_filename)
+                    else:
+                        return
+
+                with H5File(h5_filename, "w") as h5:
+                    streams: List[Any] = [StringIO()]
+                    if verbose:
+                        streams += [sys.stdout]
+                    with _RedirectOutput(streams):
+                        # Load and extract static features
+                        model = build_model(data_filename)
+                        model.extract_after_load(h5)
+
+                        if not self.skip_lp:
+                            # Solve LP relaxation
+                            relaxed = model.relax()
+                            relaxed.optimize()
+                            relaxed.extract_after_lp(h5)
+
+                        # Solve MIP
+                        model.optimize()
+                        model.extract_after_mip(h5)
+
+                        if self.write_mps:
+                            # Add lazy constraints to model
+                            model._lazy_enforce_collected()
+
+                            # Save MPS file
+                            model.write(mps_filename)
+                            gzip(mps_filename)
+
+                    h5.put_scalar("mip_log", streams[0].getvalue())
+            except:
+                print(f"Error processing: data_filename")
+                traceback.print_exc()
+
 
         if n_jobs > 1:
             p_umap(
diff --git a/miplearn/io.py b/miplearn/io.py
index ccad431..95792a1 100644
--- a/miplearn/io.py
+++ b/miplearn/io.py
@@ -87,7 +87,10 @@ def read_pkl_gz(filename: str) -> Any:
 def _to_h5_filename(data_filename: str) -> str:
     output = f"{data_filename}.h5"
     output = output.replace(".gz.h5", ".h5")
+    output = output.replace(".csv.h5", ".h5")
+    output = output.replace(".jld2.h5", ".h5")
     output = output.replace(".json.h5", ".h5")
+    output = output.replace(".lp.h5", ".h5")
+    output = output.replace(".mps.h5", ".h5")
     output = output.replace(".pkl.h5", ".h5")
-    output = output.replace(".jld2.h5", ".h5")
     return output