[SPH] allow stopping evolve until according to walltime

Author: tdavidcl

src/shammodels/sph/include/shammodels/sph/Model.hpp

@@ -956,8 +956,9 @@ namespace shammodels::sph {
             solver.print_timestep_logs();
         }
 
-        inline bool evolve_until(Tscal target_time, i32 niter_max) {
-            return solver.evolve_until(target_time, niter_max);
+        inline EvolveUntilResults evolve_until(
+            Tscal target_time, i32 niter_max, f64 max_global_walltime = -1) {
+            return solver.evolve_until(target_time, niter_max, max_global_walltime);
         }
 
         private:

src/shammodels/sph/include/shammodels/sph/Solver.hpp

@@ -31,6 +31,7 @@
 #include "shamsys/legacy/log.hpp"
 #include "shamtree/TreeTraversalCache.hpp"
 #include <functional>
+#include <limits>
 #include <memory>
 #include <optional>
 #include <stdexcept>
@@ -51,6 +52,14 @@ namespace shammodels::sph {
         inline f64 tcompute_max() { return shamalgs::collective::allreduce_max(tcompute); }
     };
 
+    struct EvolveUntilResults {
+        bool reach_target_time;
+        bool reach_niter_max;
+        bool reach_max_walltime;
+
+        i32 iter_count;
+    };
+
     /**
      * @brief The shamrock SPH model
      *
@@ -236,7 +245,16 @@ namespace shammodels::sph {
             return solver_config.get_dt_sph();
         }
 
-        inline bool evolve_until(Tscal target_time, i32 niter_max) {
+        inline EvolveUntilResults evolve_until(
+            Tscal target_time, i32 niter_max, f64 max_walltime = -1) {
+
+            f64 start_wall_time
+                = shamalgs::collective::allreduce_max(shambase::details::get_wtime());
+
+            const bool walltime_limit_active = (max_walltime != -1);
+            i32 next_walltime_check_iter
+                = walltime_limit_active ? 1 : std::numeric_limits<i32>::max();
+
             auto step = [&]() {
                 Tscal dt = solver_config.get_dt_sph();
                 Tscal t  = solver_config.get_time();
@@ -260,13 +278,81 @@ namespace shammodels::sph {
 
                 if ((iter_count >= niter_max) && (niter_max != -1)) {
                     logger::info_ln("SPH", "stopping evolve until because of niter =", iter_count);
-                    return false;
+                    return {
+                        .reach_target_time  = false,
+                        .reach_niter_max    = true,
+                        .reach_max_walltime = false,
+                        .iter_count         = iter_count,
+                    };
+                }
+
+                if (walltime_limit_active && iter_count >= next_walltime_check_iter) {
+                    f64 global_walltime
+                        = shamalgs::collective::allreduce_max(shambase::details::get_wtime());
+
+                    if ((max_walltime != -1) && (global_walltime >= max_walltime)) {
+                        logger::info_ln(
+                            "SPH",
+                            shambase::format(
+                                "stopping evolve until because of "
+                                "max_walltime = {:.2f}s > {:.2f}s",
+                                global_walltime,
+                                max_walltime));
+                        return {
+                            .reach_target_time  = false,
+                            .reach_niter_max    = false,
+                            .reach_max_walltime = true,
+                            .iter_count         = iter_count,
+                        };
+                    }
+
+                    f64 sec_per_iter
+                        = (global_walltime - start_wall_time) / static_cast<f64>(iter_count);
+
+                    auto get_remaining_iters = [&](f64 delta_walltime, f64 factor) -> i32 {
+                        if (sec_per_iter > 0) {
+                            f64 tmp = factor * delta_walltime / sec_per_iter;
+                            if (tmp > std::numeric_limits<i32>::max()) {
+                                return std::numeric_limits<i32>::max();
+                            }
+                            return static_cast<i32>(tmp);
+                        }
+                        return 1000; // default to 1000 iterations if sec_per_iter is 0
+                    };
+
+                    i32 iters_to_next_check = std::numeric_limits<i32>::max();
+
+                    i32 next_walltime_type = 0; // 0 for walltime, 1 for global walltime
+
+                    if (max_walltime != -1) {
+                        i32 iters_to_limit
+                            = get_remaining_iters(max_walltime - global_walltime, 0.1);
+                        if (iters_to_limit < iters_to_next_check) {
+                            next_walltime_type  = 1;
+                            iters_to_next_check = iters_to_limit;
+                        }
+                    }
+
+                    next_walltime_check_iter = iter_count + std::max(1, iters_to_next_check);
+
+                    logger::info_ln(
+                        "SPH",
+                        shambase::format(
+                            "next walltime check in {:.2f}s (niter = {}) global walltime = {:.2f}s",
+                            iters_to_next_check * sec_per_iter,
+                            iters_to_next_check,
+                            global_walltime));
                 }
             }
 
             print_timestep_logs();
 
-            return true;
+            return {
+                .reach_target_time  = true,
+                .reach_niter_max    = false,
+                .reach_max_walltime = false,
+                .iter_count         = iter_count,
+            };
         }
     };
 

src/shammodels/sph/src/pySPHModel.cpp

@@ -734,12 +734,13 @@ void add_instance(py::module &m, std::string name_config, std::string name_model
         .def("evolve_once", &T::evolve_once)
         .def(
             "evolve_until",
-            [](T &self, f64 target_time, i32 niter_max) {
-                return self.evolve_until(target_time, niter_max);
+            [](T &self, f64 target_time, i32 niter_max, f64 max_walltime) {
+                return self.evolve_until(target_time, niter_max, max_walltime);
             },
             py::arg("target_time"),
             py::kw_only(),
-            py::arg("niter_max") = -1)
+            py::arg("niter_max")    = -1,
+            py::arg("max_walltime") = -1)
         .def(
             "set_dt",
             [](T &self, f64 dt) {
@@ -1649,6 +1650,21 @@ ON_PYTHON_INIT {
 
     py::module msph = m.def_submodule("model_sph", "Shamrock sph solver");
 
+    py::class_<EvolveUntilResults>(m, "EvolveUntilResults")
+        .def_readwrite("reach_target_time", &EvolveUntilResults::reach_target_time)
+        .def_readwrite("reach_niter_max", &EvolveUntilResults::reach_niter_max)
+        .def_readwrite("reach_max_walltime", &EvolveUntilResults::reach_max_walltime)
+        .def_readwrite("iter_count", &EvolveUntilResults::iter_count)
+        .def("__repr__", [](const EvolveUntilResults &self) {
+            return shambase::format(
+                "EvolveUntilResults(reach_target_time={}, reach_niter_max={}, "
+                "reach_max_walltime={}, iter_count={})",
+                self.reach_target_time,
+                self.reach_niter_max,
+                self.reach_max_walltime,
+                self.iter_count);
+        });
+
     using namespace shammodels::sph;
 
     add_instance<f64_3, shammath::M4>(msph, "SPHModel_f64_3_M4_SolverConfig", "SPHModel_f64_3_M4");

src/shampylib/src/pyShamsys.cpp

@@ -16,6 +16,7 @@
 #include "shambase/exception.hpp"
 #include "shambase/numeric_limits.hpp"
 #include "shambase/stacktrace.hpp"
+#include "shamalgs/collective/reduction.hpp"
 #include "shambindings/pybindaliases.hpp"
 #include "shampylib/pyNodeInstance.hpp"
 #include "shamrock/experimental_features.hpp"
@@ -156,6 +157,14 @@ ON_PYTHON_INIT {
         dump profiling data
     )pbdoc");
 
+    m.def("get_wtime", []() {
+        return shambase::details::get_wtime();
+    });
+
+    m.def("get_wtime_sync", []() {
+        return shamalgs::collective::allreduce_max(shambase::details::get_wtime());
+    });
+
     py::module sys_module = m.def_submodule("sys", "system handling part of shamrock");
     sys_module.def("signal_handler", &shamsys::details::signal_callback_handler);