cluster management fixes

Author: SamoraHunter

ml_grid/model_classes/H2OBaseClassifier.py

@@ -129,11 +129,48 @@ def __setstate__(self, state):
 
     def _ensure_h2o_is_running(self):
         """Safely checks for and initializes an H2O cluster if not running."""
-        cluster = h2o.cluster()
+        try:
+            cluster = h2o.cluster()
+        except Exception:
+            cluster = None
+
         show_progress = getattr(global_parameters, "h2o_show_progress", False)
 
-        if not (cluster and cluster.is_running()):
-            h2o.init()
+        is_healthy = False
+        if cluster and cluster.is_running():
+            is_healthy = True
+            try:
+                # Check if cluster has memory.
+                # total_mem is in bytes. If it's 0 or None, it's broken.
+                memory = None
+                try:
+                    memory = cluster.total_mem()
+                except Exception:
+                    try:
+                        memory = cluster.free_mem()
+                    except Exception:
+                        pass
+
+                if memory is not None and isinstance(memory, (int, float)):
+                    if memory < 1024 * 1024:  # < 1MB
+                        self.logger.warning(
+                            f"H2O cluster is running but reports {memory} memory. Treating as unhealthy."
+                        )
+                        is_healthy = False
+            except Exception as e:
+                self.logger.warning(f"H2O cluster check failed: {e}")
+
+        if not is_healthy:
+            # If it was running but unhealthy, try to shut it down first to clear state
+            if cluster and cluster.is_running():
+                try:
+                    self.logger.warning("Shutting down unhealthy H2O cluster...")
+                    cluster.shutdown()
+                except Exception:
+                    pass
+
+            self.logger.info("Initializing H2O cluster...")
+            h2o.init(strict_version_check=False)
             self._is_cluster_owner = True
 
         # Set progress bar visibility based on the global parameter

tests/test_h2o_classifiers.py

@@ -236,13 +236,13 @@ def test_h2o_gam_knot_distribution_error(h2o_session_fixture):
 
     # Ensure enough unique values survive the CV split to pass the cardinality check (>= 10)
     # We need > 10 unique values in the training fold.
-    # With 50/50 split, we need roughly > 20 unique values in total.
-    # We keep it skewed (mostly 0s) to trigger the quantile error.
-    skewed_vals = np.array([0] * 70 + list(range(1, 31)))
+    # We increase sample size to 200 to ensure stability of the split while maintaining skew.
+    # 140 zeros, 60 unique values (1..60).
+    skewed_vals = np.array([0] * 140 + list(range(1, 61)))
     np.random.shuffle(skewed_vals)
 
-    X = pd.DataFrame({"feature1": np.random.rand(100), "feature_skewed": skewed_vals})
-    y = pd.Series(np.random.randint(0, 2, 100), name="outcome")
+    X = pd.DataFrame({"feature1": np.random.rand(200), "feature_skewed": skewed_vals})
+    y = pd.Series(np.random.randint(0, 2, 200), name="outcome")
 
     estimator = H2O_GAM_class(
         X=X, y=y, parameter_space_size="small"