Add support for JSON data bases for graph data structures.

CHANGELOG.md

@@ -2,6 +2,12 @@
 
 ## [Unreleased]
 
+### Added
+
+- JSON-backed storage can emit binary tensor files in rank-qualified case
+  directories or self-contained base64 manifests, with a separate manifest for
+  each domain and rank (#205).
+
 ### Changed
 
 - Workflow environments can now use active system Flux Python bindings instead

src/AMSlib/AMS.cpp

@@ -12,6 +12,7 @@
 #endif
 #include <unistd.h>
 
+#include <experimental/filesystem>
 #include <fstream>
 #include <nlohmann/json.hpp>
 #include <regex>
@@ -34,6 +35,8 @@ using namespace ams;
 namespace
 {
 
+namespace fs = std::experimental::filesystem;
+
 struct AMSAbstractModel {
 public:
   std::string SPath;

src/AMSlib/CMakeLists.txt

@@ -6,7 +6,7 @@
 # handle sources and headers
 set(AMS_LIB_SRC wf/debug.cpp wf/logger.cpp wf/utils.cpp wf/SmallVector.cpp ml/surrogate.cpp wf/basedb.cpp AMSTensor.cpp AMSGraph.cpp wf/interface.cpp wf/resource_manager.cpp ml/Model.cpp ml/AbstractModel.cpp AMS.cpp)
 
-  list(APPEND AMS_LIB_SRC wf/hdf5db.cpp)
+  list(APPEND AMS_LIB_SRC wf/hdf5db.cpp wf/jsondb.cpp)
 
 if (ENABLE_RMQ)
   list(APPEND AMS_LIB_SRC wf/rmqdb.cpp)

src/AMSlib/include/AMSGraph.hpp

@@ -17,6 +17,8 @@ class AMSTensorFieldMap
   AMSTensorMap fields_;
 
 public:
+  using const_iterator = AMSTensorMap::const_iterator;
+
   // Explicit named tensor field store. There is intentionally no operator[]:
   // use set()/insert() to create fields and at()/find() to read them so missing
   // lookups never create invalid/default AMSTensors.
@@ -38,6 +40,11 @@ class AMSTensorFieldMap
   AMSTensor& insert(std::string name, AMSTensor tensor);
   AMSTensor& set(std::string name, AMSTensor tensor);
 
+  const_iterator begin() const noexcept { return fields_.begin(); }
+  const_iterator end() const noexcept { return fields_.end(); }
+  const_iterator cbegin() const noexcept { return fields_.cbegin(); }
+  const_iterator cend() const noexcept { return fields_.cend(); }
+
   bool empty() const noexcept { return fields_.empty(); }
   std::size_t size() const noexcept { return fields_.size(); }
   void clear() noexcept { fields_.clear(); }

src/AMSlib/include/AMSTypes.hpp

@@ -20,6 +20,6 @@ typedef enum {
 
 typedef enum { AMS_UBALANCED = 0, AMS_BALANCED } AMSExecPolicy;
 
-typedef enum { AMS_NONE = 0, AMS_HDF5, AMS_RMQ } AMSDBType;
+typedef enum { AMS_NONE = 0, AMS_HDF5, AMS_RMQ, AMS_JSON } AMSDBType;
 
 }  // namespace ams

src/AMSlib/wf/basedb.cpp

@@ -1,6 +1,10 @@
+#include "wf/basedb.hpp"
+
+#include <memory>
 #include <string>
 
 #include "AMS.h"
+#include "wf/jsondb.hpp"
 
 namespace ams
 {
@@ -17,6 +21,8 @@ AMSDBType getDBType(std::string type)
     return AMSDBType::AMS_HDF5;
   } else if (type.compare("rmq") == 0) {
     return AMSDBType::AMS_RMQ;
+  } else if (type.compare("json") == 0) {
+    return AMSDBType::AMS_JSON;
   }
   return AMSDBType::AMS_NONE;
 }
@@ -30,10 +36,52 @@ std::string getDBTypeAsStr(AMSDBType type)
       return "hdf5";
     case AMSDBType::AMS_RMQ:
       return "rmq";
+    case AMSDBType::AMS_JSON:
+      return "json";
   }
   return "Unknown";
 }
 
+std::shared_ptr<BaseDB> DBManager::createDB(std::string& domainName,
+                                            AMSDBType dbType,
+                                            uint64_t rId)
+{
+  AMS_DBG(DBManager, "Instantiating data base");
+
+  if ((dbType == AMSDBType::AMS_HDF5 || dbType == AMSDBType::AMS_JSON) &&
+      !fs_interface.isConnected()) {
+    THROW(std::runtime_error,
+          "File System is not configured, Please specify output directory");
+  } else if (dbType == AMSDBType::AMS_RMQ && !rmq_interface.isConnected()) {
+    THROW(std::runtime_error, "Rabbit MQ data base is not configured");
+  }
+
+  switch (dbType) {
+#ifdef __AMS_ENABLE_HDF5__
+    case AMSDBType::AMS_HDF5:
+      return std::make_shared<hdf5DB>(fs_interface.path(), domainName, rId);
+#endif
+#ifdef __AMS_ENABLE_RMQ__
+    case AMSDBType::AMS_RMQ:
+      return std::make_shared<RabbitMQDB>(rmq_interface,
+                                          domainName,
+                                          rId,
+                                          updateSurrogate);
+#endif
+    case AMSDBType::AMS_JSON: {
+      // JSONDB needs json_mode configuration - get from environment or default
+      const char* json_mode_env = std::getenv("AMS_JSON_MODE");
+      std::string json_mode = json_mode_env ? json_mode_env : "binary";
+      return std::make_shared<JSONDB>(fs_interface.path(),
+                                      domainName,
+                                      rId,
+                                      json_mode);
+    }
+    default:
+      return nullptr;
+  }
+  return nullptr;
+}
 
 }  // namespace db
 }  // namespace ams

src/AMSlib/wf/basedb.hpp

@@ -31,7 +31,14 @@
 #include "wf/resource_manager.hpp"
 #include "wf/utils.hpp"
 
-namespace fs = std::experimental::filesystem;
+// Forward declarations for graph types
+namespace ams
+{
+struct AMSHomogeneousGraph;
+struct AMSHomogeneousGraphFields;
+struct AMSHeterogeneousGraph;
+struct AMSHeterogeneousGraphFields;
+}  // namespace ams
 
 #ifdef __AMS_ENABLE_HDF5__
 #include <H5Ipublic.h>
@@ -124,6 +131,35 @@ class BaseDB
   virtual void store(ArrayRef<torch::Tensor> Inputs,
                      ArrayRef<torch::Tensor> Outputs) = 0;
 
+  /**
+   * @brief Store graph data with outputs/targets for training.
+   * Default implementation throws - only backends that support graphs override this.
+   * @param[in] graph The homogeneous graph containing input features
+   * @param[in] outputs The graph fields containing output/target data
+   */
+  virtual void store(
+      [[maybe_unused]] const ams::AMSHomogeneousGraph& graph,
+      [[maybe_unused]] const ams::AMSHomogeneousGraphFields& outputs)
+  {
+    THROW(std::runtime_error,
+          (this->type() + " database does not support graph storage").c_str());
+  }
+
+  /**
+   * @brief Store heterogeneous graph data with outputs/targets for training.
+   * Default implementation throws - only backends that support graphs override this.
+   * @param[in] graph The heterogeneous graph containing input features
+   * @param[in] outputs The graph fields containing output/target data
+   */
+  virtual void store(
+      [[maybe_unused]] const ams::AMSHeterogeneousGraph& graph,
+      [[maybe_unused]] const ams::AMSHeterogeneousGraphFields& outputs)
+  {
+    THROW(std::runtime_error,
+          (this->type() + " database does not support heterogeneous graph "
+                          "storage")
+              .c_str());
+  }
 
   uint64_t getId() const { return id; }
 
@@ -175,29 +211,29 @@ class FileDB : public BaseDB
          uint64_t rId)
       : BaseDB(rId)
   {
-    fs::path Path(path);
+    std::experimental::filesystem::path Path(path);
     std::error_code ec;
 
-    if (!fs::exists(Path, ec)) {
+    if (!std::experimental::filesystem::exists(Path, ec)) {
       std::cerr << "[ERROR]: Path:'" << path << "' does not exist\n";
       exit(-1);
     }
 
     checkError(ec);
 
-    if (!fs::is_directory(Path, ec)) {
+    if (!std::experimental::filesystem::is_directory(Path, ec)) {
       std::cerr << "[ERROR]: Path:'" << path << "' is a file NOT a directory\n";
       exit(-1);
     }
 
-    Path = fs::absolute(Path);
+    Path = std::experimental::filesystem::absolute(Path);
     fp = Path.string();
 
     // We can now create the filename
     std::string dbfn(fn + "_");
     dbfn += std::to_string(rId) + suffix;
-    Path /= fs::path(dbfn);
-    this->fn = fs::absolute(Path).string();
+    Path /= std::experimental::filesystem::path(dbfn);
+    this->fn = std::experimental::filesystem::absolute(Path).string();
     AMS_DBG(DB, "File System DB writes to file {}", this->fn)
   }
 
@@ -1559,8 +1595,9 @@ class RMQInterface
     flush(100, 100);
     _publishingManager->stop();
     auto size = MessagesBuffer::getInstance().size();
-    if (size != 0)
+    if (size != 0) {
       AMS_DBG(RMQInterface, "Rank {} did not ack {} messages", _rId, size)
+    }
   }
 
   ~RMQInterface()
@@ -1668,10 +1705,10 @@ class FilesystemInterface
   bool connect(std::string& path)
   {
     connected = true;
-    fs::path Path(path);
+    std::experimental::filesystem::path Path(path);
     std::error_code ec;
 
-    if (!fs::exists(Path, ec)) {
+    if (!std::experimental::filesystem::exists(Path, ec)) {
       THROW(std::runtime_error,
             ("Path: :'" + path + "' does not exist").c_str());
       exit(-1);
@@ -1760,37 +1797,10 @@ class DBManager
   * @param[in] rId a unique Id for each process taking part in a distributed
   * execution (rank-id)
   */
+  // Declared here, implemented in basedb.cpp to avoid including jsondb.hpp in header
   std::shared_ptr<BaseDB> createDB(std::string& domainName,
                                    AMSDBType dbType,
-                                   uint64_t rId = 0)
-  {
-
-    AMS_DBG(DBManager, "Instantiating data base");
-
-    if ((dbType == AMSDBType::AMS_HDF5) && !fs_interface.isConnected()) {
-      THROW(std::runtime_error,
-            "File System is not configured, Please specify output directory");
-    } else if (dbType == AMSDBType::AMS_RMQ && !rmq_interface.isConnected()) {
-      THROW(std::runtime_error, "Rabbit MQ data base is not configured");
-    }
-
-    switch (dbType) {
-#ifdef __AMS_ENABLE_HDF5__
-      case AMSDBType::AMS_HDF5:
-        return std::make_shared<hdf5DB>(fs_interface.path(), domainName, rId);
-#endif
-#ifdef __AMS_ENABLE_RMQ__
-      case AMSDBType::AMS_RMQ:
-        return std::make_shared<RabbitMQDB>(rmq_interface,
-                                            domainName,
-                                            rId,
-                                            updateSurrogate);
-#endif
-      default:
-        return nullptr;
-    }
-    return nullptr;
-  }
+                                   uint64_t rId = 0);
 
   /**
   * @brief get a data base object referred by this string.
@@ -1904,10 +1914,10 @@ class DBManager
                           std::string& routing_key,
                           bool update_surrogate)
   {
-    fs::path Path(rmq_cert);
+    std::experimental::filesystem::path Path(rmq_cert);
     std::error_code ec;
     AMS_CWARNING(AMS,
-                 !fs::exists(Path, ec),
+                 !std::experimental::filesystem::exists(Path, ec),
                  "Certificate file '{}' for RMQ server does not exist. AMS "
                  "will "
                  "try to connect without it.",

src/AMSlib/wf/hdf5db.cpp

@@ -13,6 +13,7 @@
 #include <torch/torch.h>
 #include <torch/types.h>
 
+#include <experimental/filesystem>
 #include <stdexcept>
 
 #include "ArrayRef.hpp"
@@ -22,6 +23,11 @@
 using namespace ams::db;
 using namespace ams;
 
+namespace
+{
+namespace fs = std::experimental::filesystem;
+}
+
 static std::string SmallVectorToString(ams::MutableArrayRef<hsize_t> shape)
 {
   std::ostringstream oss;

src/AMSlib/wf/interface.cpp

@@ -490,26 +490,6 @@ void callAMS(ams::AMSWorkflow* executor,
   executor->evaluate(Physics, tins, tinouts, touts);
 }
 
-// ============================================================================
-// Graph-based callApplication overloads
-// ============================================================================
-
-void callApplication(ams::HomogeneousGraphDomainFn CallBack,
-                     const ams::AMSHomogeneousGraph& graph,
-                     ams::AMSHomogeneousGraphFields& outputs)
-{
-  // Directly invoke the user's physics callback with graph-native types
-  CallBack(graph, outputs);
-}
-
-void callApplication(ams::HeterogeneousGraphDomainFn CallBack,
-                     const ams::AMSHeterogeneousGraph& graph,
-                     ams::AMSHeterogeneousGraphFields& outputs)
-{
-  // Directly invoke the user's physics callback with graph-native types
-  CallBack(graph, outputs);
-}
-
 // ============================================================================
 // Graph surrogate execution (in ams namespace for friend access)
 // ============================================================================
@@ -675,31 +655,15 @@ void callAMS(ams::AMSWorkflow* executor,
              const ams::AMSHomogeneousGraph& graph_input,
              ams::AMSHomogeneousGraphFields& outputs)
 {
-  // Try graph surrogate execution first
-  bool surrogate_used = tryGraphSurrogate(executor, graph_input, outputs);
-
-  // If surrogate succeeded, we're done
-  if (surrogate_used) {
-    return;
-  }
-
-  // Otherwise, fallback to original physics computation
-  callApplication(Physics, graph_input, outputs);
+  // Delegate to public evaluate method (mirrors tensor pattern)
+  executor->evaluate(Physics, graph_input, outputs);
 }
 
 void callAMS(ams::AMSWorkflow* executor,
              ams::HeterogeneousGraphDomainFn Physics,
              const ams::AMSHeterogeneousGraph& graph_input,
              ams::AMSHeterogeneousGraphFields& outputs)
 {
-  // Try graph surrogate execution first
-  bool surrogate_used = tryGraphSurrogate(executor, graph_input, outputs);
-
-  // If surrogate succeeded, we're done
-  if (surrogate_used) {
-    return;
-  }
-
-  // Otherwise, fallback to original physics computation
-  callApplication(Physics, graph_input, outputs);
+  // Delegate to public evaluate method (mirrors tensor pattern)
+  executor->evaluate(Physics, graph_input, outputs);
 }