updated bedrock module to latest API (remains to be tested)

mochi-hpc · Sep 20, 2024 · 6f30698 · 6f30698
1 parent 5823a5b
commit 6f30698
Show file tree

Hide file tree

Showing 3 changed files with 78 additions and 81 deletions.
diff --git a/spack.yaml b/spack.yaml
@@ -10,7 +10,7 @@ spack:
   - fmt
   - tclap
   - mochi-bedrock-module-api
-  - mochi-abt-io
+  - mochi-abt-io+bedrock
   - mochi-remi
   - py-mochi-margo
   - py-configspace
@@ -34,3 +34,5 @@ spack:
       require: "@1.1.4:"
     nlohmann-json-schema-validator:
       require: "@2.3.0:"
+    mochi-bedrock-module-api:
+      require: "@0.2.0:"
diff --git a/src/BedrockModule.cpp b/src/BedrockModule.cpp
@@ -1,98 +1,91 @@
 /*
- * (C) 2023 The University of Chicago
+ * (C) 2024 The University of Chicago
  *
  * See COPYRIGHT in top-level directory.
  */
-#include "warabi/Client.hpp"
+
 #include "warabi/Provider.hpp"
-#include "warabi/TargetHandle.hpp"
-#include <bedrock/AbstractServiceFactory.hpp>
+#include <bedrock/AbstractComponent.hpp>
 
 namespace tl = thallium;
 
-class WarabiFactory : public bedrock::AbstractServiceFactory {
-
-    public:
+class WarabiComponent : public bedrock::AbstractComponent {
 
-    WarabiFactory() {}
+    std::unique_ptr<warabi::Provider> m_provider;
 
-    void *registerProvider(const bedrock::FactoryArgs &args) override {
-        remi_client_t remi_cl = nullptr;
-        remi_provider_t remi_pr = nullptr;
-        auto it = args.dependencies.find("remi_client");
-        if(it != args.dependencies.end() && it->second.dependencies.size() != 0) {
-            remi_cl = it->second.dependencies[0]->getHandle<remi_client_t>();
-        }
-        it = args.dependencies.find("remi_provider");
-        if(it != args.dependencies.end() && it->second.dependencies.size() != 0) {
-            remi_pr = it->second.dependencies[0]->getHandle<remi_provider_t>();
-        }
-        auto provider = new warabi::Provider(args.mid, args.provider_id,
-                args.config, tl::pool(args.pool), remi_cl, remi_pr);
-        return static_cast<void *>(provider);
-    }
+    public:
 
-    void deregisterProvider(void *p) override {
-        auto provider = static_cast<warabi::Provider *>(p);
-        delete provider;
+    WarabiComponent(const tl::engine& engine,
+                   uint16_t  provider_id,
+                   const std::string& config,
+                   const tl::pool& pool,
+                   remi_client_t remi_cl = nullptr,
+                   remi_provider_t remi_pr = nullptr)
+    : m_provider{
+        std::make_unique<warabi::Provider>(
+            engine, provider_id, config, pool, remi_cl, remi_pr)}
+    {}
+
+    void* getHandle() override {
+        return static_cast<void*>(m_provider.get());
     }
 
-    std::string getProviderConfig(void *p) override {
-        auto provider = static_cast<warabi::Provider *>(p);
-        return provider->getConfig();
+    std::string getConfig() override {
+        return m_provider->getConfig();
     }
 
-    void migrateProvider(
-        void* p, const char* dest_addr,
-        uint16_t dest_provider_id,
-        const char* options_json, bool remove_source) override {
-        auto provider = static_cast<warabi::Provider *>(p);
-        if(!remove_source) {
-            throw warabi::Exception{
-                "Migrating a provider without removing the source is not supported"};
+    static std::shared_ptr<bedrock::AbstractComponent>
+        Register(const bedrock::ComponentArgs& args) {
+            tl::pool pool;
+            auto it = args.dependencies.find("pool");
+            if(it != args.dependencies.end() && !it->second.empty()) {
+                pool = it->second[0]->getHandle<tl::pool>();
+            }
+            remi_client_t remi_sender = nullptr;
+            it = args.dependencies.find("remi_sender");
+            if(it != args.dependencies.end() && !it->second.empty()) {
+                auto component = it->second[0]->getHandle<bedrock::ComponentPtr>();
+                remi_sender = static_cast<remi_client_t>(component->getHandle());
+            }
+            remi_provider_t remi_receiver = nullptr;
+            it = args.dependencies.find("remi_receiver");
+            if(it != args.dependencies.end() && !it->second.empty()) {
+                auto component = it->second[0]->getHandle<bedrock::ComponentPtr>();
+                remi_receiver = static_cast<remi_provider_t>(component->getHandle());
+            }
+            return std::make_shared<WarabiComponent>(
+                args.engine, args.provider_id, args.config, pool,
+                remi_sender, remi_receiver);
         }
-        provider->migrateTarget(dest_addr, dest_provider_id, options_json);
-    }
-
-    void *initClient(const bedrock::FactoryArgs& args) override {
-        return static_cast<void *>(new warabi::Client(args.mid));
-    }
-
-    void finalizeClient(void *client) override {
-        delete static_cast<warabi::Client *>(client);
-    }
-
-    std::string getClientConfig(void* c) override {
-        auto client = static_cast<warabi::Client*>(c);
-        return client->getConfig();
-    }
 
-    void *createProviderHandle(void *c, hg_addr_t address,
-            uint16_t provider_id) override {
-        auto client = static_cast<warabi::Client *>(c);
-        auto addr_str = static_cast<std::string>(tl::endpoint{client->engine(), address, false});
-        auto th = new warabi::TargetHandle{
-            client->makeTargetHandle(addr_str, provider_id)};
-        return static_cast<void *>(th);
-    }
-
-    void destroyProviderHandle(void *providerHandle) override {
-        auto th = static_cast<warabi::TargetHandle *>(providerHandle);
-        delete th;
-    }
-
-    const std::vector<bedrock::Dependency> &getProviderDependencies() override {
-        static const std::vector<bedrock::Dependency> no_dependency = {
-            {"remi_client", "remi", BEDROCK_OPTIONAL},
-            {"remi_provider", "remi", BEDROCK_OPTIONAL}
-        };
-        return no_dependency;
-    }
-
-    const std::vector<bedrock::Dependency> &getClientDependencies() override {
-        static const std::vector<bedrock::Dependency> no_dependency;
-        return no_dependency;
-    }
+    static std::vector<bedrock::Dependency>
+        GetDependencies(const bedrock::ComponentArgs& args) {
+            (void)args;
+            std::vector<bedrock::Dependency> dependencies{
+                bedrock::Dependency{
+                    /* name */ "pool",
+                    /* type */ "pool",
+                    /* is_required */ false,
+                    /* is_array */ false,
+                    /* is_updatable */ false
+                },
+                bedrock::Dependency{
+                    /* name */ "remi_sender",
+                    /* type */ "remi_sender",
+                    /* is_required */ false,
+                    /* is_array */ false,
+                    /* is_updatable */ false
+                },
+                bedrock::Dependency{
+                    /* name */ "remi_receiver",
+                    /* type */ "remi_receiver",
+                    /* is_required */ false,
+                    /* is_array */ false,
+                    /* is_updatable */ false
+                }
+            };
+            return dependencies;
+        }
 };
 
-BEDROCK_REGISTER_MODULE_FACTORY(warabi, WarabiFactory)
+BEDROCK_REGISTER_COMPONENT_TYPE(warabi, WarabiComponent)
diff --git a/tests/spack.yaml b/tests/spack.yaml
@@ -3,7 +3,7 @@ spack:
   - cmake
   - pkgconfig
   - mochi-thallium
-  - mochi-abt-io
+  - mochi-abt-io+bedrock
   - nlohmann-json
   - nlohmann-json-schema-validator
   - spdlog
@@ -34,6 +34,8 @@ spack:
       require: "@2.3.0:"
     py-configspace:
       require: "@1.1.4:"
+    mochi-bedrock-module-api:
+      require: "@0.2.0:"
   mirrors:
     mochi-buildcache:
       url: oci://ghcr.io/mochi-hpc/mochi-spack-buildcache