Merge pull request STEllAR-GROUP#6617 from STEllAR-GROUP/fixing_6616

Adding support for the Flux job scheduling environment

Author: hkaiser

libs/core/batch_environments/CMakeLists.txt

@@ -1,4 +1,4 @@
-# Copyright (c) 2019-2022 The STE||AR-Group
+# Copyright (c) 2019-2025 The STE||AR-Group
 #
 # SPDX-License-Identifier: BSL-1.0
 # Distributed under the Boost Software License, Version 1.0. (See accompanying
@@ -9,6 +9,7 @@ list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake")
 set(batch_environments_headers
     hpx/batch_environments/alps_environment.hpp
     hpx/batch_environments/batch_environment.hpp
+    hpx/batch_environments/flux_environment.hpp
     hpx/batch_environments/pjm_environment.hpp
     hpx/batch_environments/pbs_environment.hpp
     hpx/batch_environments/slurm_environment.hpp
@@ -25,8 +26,8 @@ set(batch_environments_compat_headers
 # cmake-format: on
 
 set(batch_environments_sources
-    alps_environment.cpp batch_environment.cpp pbs_environment.cpp
-    pjm_environment.cpp slurm_environment.cpp
+    alps_environment.cpp batch_environment.cpp flux_environment.cpp
+    pbs_environment.cpp pjm_environment.cpp slurm_environment.cpp
 )
 
 include(HPX_AddModule)

libs/core/batch_environments/include/hpx/batch_environments/flux_environment.hpp

@@ -0,0 +1,40 @@
+//  Copyright (c) 2007-2025 Hartmut Kaiser
+//
+//  SPDX-License-Identifier: BSL-1.0
+//  Distributed under the Boost Software License, Version 1.0. (See accompanying
+//  file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)
+
+#pragma once
+
+#include <hpx/config.hpp>
+
+#include <cstddef>
+#include <string>
+
+namespace hpx::util::batch_environments {
+
+    struct flux_environment
+    {
+        HPX_CORE_EXPORT flux_environment();
+
+        constexpr bool valid() const noexcept
+        {
+            return valid_;
+        }
+
+        constexpr std::size_t node_num() const noexcept
+        {
+            return node_num_;
+        }
+
+        constexpr std::size_t num_localities() const noexcept
+        {
+            return num_localities_;
+        }
+
+    private:
+        std::size_t node_num_;
+        std::size_t num_localities_;
+        bool valid_;
+    };
+}    // namespace hpx::util::batch_environments

libs/core/batch_environments/src/batch_environment.cpp

@@ -1,4 +1,4 @@
-//  Copyright (c) 2007-2023 Hartmut Kaiser
+//  Copyright (c) 2007-2025 Hartmut Kaiser
 //  Copyright (c)      2013 Thomas Heller
 //
 //  SPDX-License-Identifier: BSL-1.0
@@ -9,6 +9,7 @@
 #include <hpx/asio/asio_util.hpp>
 #include <hpx/batch_environments/alps_environment.hpp>
 #include <hpx/batch_environments/batch_environment.hpp>
+#include <hpx/batch_environments/flux_environment.hpp>
 #include <hpx/batch_environments/pbs_environment.hpp>
 #include <hpx/batch_environments/pjm_environment.hpp>
 #include <hpx/batch_environments/slurm_environment.hpp>
@@ -90,6 +91,15 @@ namespace hpx::util {
             return;
         }
 
+        batch_environments::flux_environment const flux_env;
+        if (flux_env.valid())
+        {
+            batch_name_ = "FLUX";
+            num_localities_ = flux_env.num_localities();
+            node_num_ = flux_env.node_num();
+            return;
+        }
+
         batch_environments::slurm_environment const slurm_env(nodelist, debug);
         if (slurm_env.valid())
         {
@@ -200,12 +210,14 @@ namespace hpx::util {
 #if defined(HPX_HAVE_PARCELPORT_TCP)
             std::cerr << "Nodes from nodelist:" << std::endl;
             node_map_type::const_iterator const end = nodes_.end();
+            // clang-format off
             for (node_map_type::const_iterator it = nodes_.begin(); it != end;
                  ++it)
             {
                 std::cerr << (*it).second.first << ": " << (*it).second.second
                           << " (" << (*it).first << ")" << std::endl;
             }
+            // clang-format on
 #endif
         }
         HPX_UNUSED(nodes);

libs/core/batch_environments/src/flux_environment.cpp

@@ -0,0 +1,41 @@
+//  Copyright (c) 2007-2025 Hartmut Kaiser
+//
+//  SPDX-License-Identifier: BSL-1.0
+//  Distributed under the Boost Software License, Version 1.0. (See accompanying
+//  file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)
+
+#include <hpx/batch_environments/flux_environment.hpp>
+#include <hpx/util/from_string.hpp>
+
+#include <cstddef>
+#include <string>
+
+namespace hpx::util::batch_environments {
+
+    // FLUX_TASK_RANK: number of physical nodes
+    // FLUX_JOB_NNODES: number of nodes in the job
+    flux_environment::flux_environment()
+      : node_num_(static_cast<std::size_t>(-1))
+      , num_localities_(0)
+      , valid_(false)
+    {
+        char const* num_nodes = std::getenv("FLUX_JOB_NNODES");
+        valid_ = num_nodes != nullptr;
+        if (valid_)
+        {
+            // Get the number of localities
+            num_localities_ = from_string<std::size_t>(num_nodes);
+
+            // Initialize our node number, if available
+            char const* var = std::getenv("FLUX_TASK_RANK");
+            if (var != nullptr)
+            {
+                node_num_ = from_string<std::size_t>(var);
+            }
+            else
+            {
+                valid_ = false;
+            }
+        }
+    }
+}    // namespace hpx::util::batch_environments