Ensure auto_nbands handler doesn't always terminate jobs

src/custodian/vasp/handlers.py

@@ -725,15 +725,31 @@ def correct(self, directory="./"):
             self.error_count["algo_tet"] += 1
 
         if "auto_nbands" in self.errors and (nbands := self._get_nbands_from_outcar(directory)):
-            nelect = load_outcar(os.path.join(directory, "OUTCAR")).nelect
-            if nelect and nbands > 2 * nelect:
+            outcar = load_outcar(os.path.join(directory, "OUTCAR"))
+
+            if (nelect := outcar.nelect) and (nbands > 2 * nelect):
                 self.error_count["auto_nbands"] += 1
                 warnings.warn(
                     "NBANDS seems to be too high. The electronic structure may be inaccurate. "
                     "You may want to rerun this job with a smaller number of cores.",
                     UserWarning,
                 )
 
+            elif nbands := vi["INCAR"].get("NBANDS"):
+                kpar = vi["INCAR"].get("KPAR", 1)
+                ncore = vi["INCAR"].get("NCORE", 1)
+                # If the user set an NBANDS that isn't compatible with parallelization settings,
+                # increase NBANDS to ensure correct task distribution and issue a UserWarning.
+                # The number of ranks per band is (number of MPI ranks) / (KPAR * NCORE)
+                if (ranks := outcar.run_stats.get("cores")) and (rem_bands := nbands % (ranks // (kpar * ncore))) != 0:
+                    actions.append({"dict": "INCAR", "action": {"_set": {"NBANDS": nbands + rem_bands}}})
+                    warnings.warn(
+                        f"Your NBANDS={nbands} setting was incompatible with your parallelization "
+                        f"settings, KPAR={kpar}, NCORE={ncore}, over {ranks} ranks. "
+                        f"The number of bands has been decreased accordingly to {nbands + rem_bands}.",
+                        UserWarning,
+                    )
+
         VaspModder(vi=vi, directory=directory).apply_actions(actions)
         return {"errors": list(self.errors), "actions": actions}