From 41550b0659e726ef65dce46feb9790a6338f2a83 Mon Sep 17 00:00:00 2001 From: Yohann Dudouit Date: Mon, 1 Jun 2026 14:50:47 -0700 Subject: [PATCH 01/18] MAke diffusion MGN test not appear during build phase. --- tests/AMSlib/models/generate_mgn_graph_diffusion.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/AMSlib/models/generate_mgn_graph_diffusion.py b/tests/AMSlib/models/generate_mgn_graph_diffusion.py index 83f698e0..8f6a7ecb 100644 --- a/tests/AMSlib/models/generate_mgn_graph_diffusion.py +++ b/tests/AMSlib/models/generate_mgn_graph_diffusion.py @@ -56,6 +56,7 @@ import copy import json import os +import struct import sys from pathlib import Path from typing import Dict, Iterable, List, Optional, Tuple From bfff6ba779ff062c05465aedebd0787f347254e8 Mon Sep 17 00:00:00 2001 From: Yohann Dudouit Date: Fri, 21 Aug 2026 16:35:39 -0700 Subject: [PATCH 02/18] Add json Data Base. --- src/AMSlib/CMakeLists.txt | 2 +- src/AMSlib/include/AMSTypes.hpp | 2 +- src/AMSlib/wf/basedb.cpp | 47 ++ src/AMSlib/wf/basedb.hpp | 66 +- src/AMSlib/wf/interface.cpp | 24 +- src/AMSlib/wf/jsondb.cpp | 581 +++++++++++++ src/AMSlib/wf/jsondb.hpp | 191 +++++ src/AMSlib/wf/workflow.hpp | 89 +- .../test_graph_workflow_storage.cpp | 765 ++++++++++++++++++ 9 files changed, 1706 insertions(+), 61 deletions(-) create mode 100644 src/AMSlib/wf/jsondb.cpp create mode 100644 src/AMSlib/wf/jsondb.hpp create mode 100644 tests/AMSlib/ams_interface/test_graph_workflow_storage.cpp diff --git a/src/AMSlib/CMakeLists.txt b/src/AMSlib/CMakeLists.txt index df4229ee..5d378245 100644 --- a/src/AMSlib/CMakeLists.txt +++ b/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) diff --git a/src/AMSlib/include/AMSTypes.hpp b/src/AMSlib/include/AMSTypes.hpp index de9f5249..0ef4d87d 100644 --- a/src/AMSlib/include/AMSTypes.hpp +++ b/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 diff --git a/src/AMSlib/wf/basedb.cpp b/src/AMSlib/wf/basedb.cpp index 0edaa158..b5adaa1e 100644 --- a/src/AMSlib/wf/basedb.cpp +++ b/src/AMSlib/wf/basedb.cpp @@ -1,6 +1,9 @@ #include +#include #include "AMS.h" +#include "wf/basedb.hpp" +#include "wf/jsondb.hpp" namespace ams { @@ -17,6 +20,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 +35,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 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(fs_interface.path(), domainName, rId); +#endif +#ifdef __AMS_ENABLE_RMQ__ + case AMSDBType::AMS_RMQ: + return std::make_shared(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(fs_interface.path(), + domainName, + rId, + json_mode); + } + default: + return nullptr; + } + return nullptr; +} } // namespace db } // namespace ams diff --git a/src/AMSlib/wf/basedb.hpp b/src/AMSlib/wf/basedb.hpp index 33483eb3..d54980ef 100644 --- a/src/AMSlib/wf/basedb.hpp +++ b/src/AMSlib/wf/basedb.hpp @@ -33,6 +33,15 @@ 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 #include @@ -124,6 +133,32 @@ class BaseDB virtual void store(ArrayRef Inputs, ArrayRef 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(const ams::AMSHomogeneousGraph& graph, + 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(const ams::AMSHeterogeneousGraph& graph, + 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; } @@ -1760,37 +1795,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 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(fs_interface.path(), domainName, rId); -#endif -#ifdef __AMS_ENABLE_RMQ__ - case AMSDBType::AMS_RMQ: - return std::make_shared(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. diff --git a/src/AMSlib/wf/interface.cpp b/src/AMSlib/wf/interface.cpp index 71c474ea..2cb7fa61 100644 --- a/src/AMSlib/wf/interface.cpp +++ b/src/AMSlib/wf/interface.cpp @@ -675,16 +675,8 @@ 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, @@ -692,14 +684,6 @@ void callAMS(ams::AMSWorkflow* executor, 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); } diff --git a/src/AMSlib/wf/jsondb.cpp b/src/AMSlib/wf/jsondb.cpp new file mode 100644 index 00000000..52edd159 --- /dev/null +++ b/src/AMSlib/wf/jsondb.cpp @@ -0,0 +1,581 @@ +/* + * Copyright 2021-2023 Lawrence Livermore National Security, LLC and other + * AMSLib Project Developers + * + * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + */ + +#include "wf/jsondb.hpp" + +#include + +#include +#include +#include +#include + +#include "wf/debug.h" + +using namespace ams::db; +using namespace ams; + +// ---------------------------------------------------------------------- +// Helper functions +// ---------------------------------------------------------------------- + +namespace +{ + +// Check system endianness +bool isLittleEndian() +{ + uint32_t test = 0x01020304; + return (*reinterpret_cast(&test)) == 0x04; +} + +// Base64 encoding for pure JSON mode +std::string base64Encode(const uint8_t* data, size_t len) +{ + static const char* base64_chars = + "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"; + + std::string ret; + int i = 0; + uint8_t char_array_3[3]; + uint8_t char_array_4[4]; + + while (len--) { + char_array_3[i++] = *(data++); + if (i == 3) { + char_array_4[0] = (char_array_3[0] & 0xfc) >> 2; + char_array_4[1] = + ((char_array_3[0] & 0x03) << 4) + ((char_array_3[1] & 0xf0) >> 4); + char_array_4[2] = + ((char_array_3[1] & 0x0f) << 2) + ((char_array_3[2] & 0xc0) >> 6); + char_array_4[3] = char_array_3[2] & 0x3f; + + for (i = 0; i < 4; i++) ret += base64_chars[char_array_4[i]]; + i = 0; + } + } + + if (i) { + for (int j = i; j < 3; j++) char_array_3[j] = '\0'; + + char_array_4[0] = (char_array_3[0] & 0xfc) >> 2; + char_array_4[1] = + ((char_array_3[0] & 0x03) << 4) + ((char_array_3[1] & 0xf0) >> 4); + char_array_4[2] = + ((char_array_3[1] & 0x0f) << 2) + ((char_array_3[2] & 0xc0) >> 6); + + for (int j = 0; j < i + 1; j++) ret += base64_chars[char_array_4[j]]; + + while (i++ < 3) ret += '='; + } + + return ret; +} + +} // anonymous namespace + +// ---------------------------------------------------------------------- +// JSONDB Implementation +// ---------------------------------------------------------------------- + +JSONDB::JSONDB(std::string path, + std::string domain_name, + uint64_t rId, + std::string json_mode) + : FileDB(path, domain_name, "_jsondb.json", rId), + json_mode_(json_mode), + case_counter_(0), + finalized_(false) +{ + if (!isLittleEndian()) { + AMS_WARNING(JSONDB, + "System is not little-endian. Binary output may not be " + "compatible with Python loaders."); + } + + if (json_mode_ != "binary" && json_mode_ != "json") { + THROW(std::invalid_argument, + ("Invalid json_mode: " + json_mode_ + + ". Must be 'binary' or 'json'.") + .c_str()); + } + + AMS_DBG(JSONDB, + "Created JSONDB at '{}' with mode '{}'", + fp, + json_mode_); +} + +JSONDB::~JSONDB() +{ + if (!finalized_) { + try { + finalize(); + } catch (const std::exception& e) { + AMS_WARNING(JSONDB, + "Exception during automatic finalization: {}", + e.what()); + } + } +} + +std::string JSONDB::dtypeToString(AMSDType dtype) const +{ + switch (dtype) { + case AMS_SINGLE: + return "float32"; + case AMS_DOUBLE: + return "float64"; + case AMS_INT32: + return "int32"; + case AMS_INT64: + return "int64"; + default: + return "unknown"; + } +} + +std::string JSONDB::torchDTypeToString(torch::Dtype dtype) const +{ + if (dtype == torch::kFloat32 || dtype == torch::kFloat) + return "float32"; + if (dtype == torch::kFloat64 || dtype == torch::kDouble) + return "float64"; + if (dtype == torch::kInt32) + return "int32"; + if (dtype == torch::kInt64 || dtype == torch::kLong) + return "int64"; + return "unknown"; +} + +size_t JSONDB::dtypeSize(AMSDType dtype) const +{ + switch (dtype) { + case AMS_SINGLE: + return 4; + case AMS_DOUBLE: + return 8; + case AMS_INT32: + return 4; + case AMS_INT64: + return 8; + default: + return 0; + } +} + +size_t JSONDB::writeBinaryTensor(const AMSTensor& tensor, + const std::string& path) +{ + // Get tensor properties + const void* data = tensor.raw_data(); + size_t byte_size = tensor.elements() * tensor.element_size(); + AMSResourceType location = tensor.location(); + + // Handle GPU tensors - copy to CPU first + std::vector cpu_buffer; + if (location != AMSResourceType::AMS_HOST) { + AMS_WARNING( + JSONDB, + "GPU tensor detected. Copying to CPU for serialization (not yet " + "implemented - will fail)."); + THROW(std::runtime_error, + "GPU tensor serialization not yet implemented"); + // TODO: Implement cudaMemcpy/hipMemcpy here + } + + // Ensure contiguous layout + if (!tensor.contiguous()) { + AMS_WARNING(JSONDB, + "Non-contiguous tensor detected. Creating contiguous copy."); + // For non-contiguous, we need to iterate with strides + // For now, throw an error + THROW(std::runtime_error, + "Non-contiguous tensor serialization not yet implemented"); + } + + // Write binary file + fs::path full_path = fs::path(fp) / path; + fs::create_directories(full_path.parent_path()); + + std::ofstream file(full_path.string(), std::ios::binary); + if (!file.is_open()) { + THROW(std::runtime_error, + ("Failed to open file for writing: " + full_path.string()).c_str()); + } + + file.write(static_cast(data), byte_size); + file.close(); + + AMS_DBG(JSONDB, + "Wrote binary tensor to '{}' ({} bytes)", + path, + byte_size); + + return byte_size; +} + +size_t JSONDB::writeBinaryTensor(const torch::Tensor& tensor, + const std::string& path) +{ + // Ensure tensor is contiguous and on CPU + torch::Tensor cpu_tensor = tensor.contiguous().cpu(); + + size_t byte_size = cpu_tensor.nbytes(); + + // Write binary file + fs::path full_path = fs::path(fp) / path; + fs::create_directories(full_path.parent_path()); + + std::ofstream file(full_path.string(), std::ios::binary); + if (!file.is_open()) { + THROW(std::runtime_error, + ("Failed to open file for writing: " + full_path.string()).c_str()); + } + + file.write(static_cast(cpu_tensor.data_ptr()), byte_size); + file.close(); + + AMS_DBG(JSONDB, + "Wrote PyTorch tensor to '{}' ({} bytes)", + path, + byte_size); + + return byte_size; +} + +nlohmann::json JSONDB::encodeBase64Tensor(const AMSTensor& tensor) +{ + // For pure JSON mode, encode as base64 + const uint8_t* data = static_cast(tensor.raw_data()); + size_t byte_size = tensor.elements() * tensor.element_size(); + + // Handle GPU/non-contiguous tensors + if (tensor.location() != AMSResourceType::AMS_HOST || + !tensor.contiguous()) { + THROW(std::runtime_error, + "Base64 encoding only supports contiguous CPU tensors currently"); + } + + nlohmann::json result; + result["encoding"] = "base64"; + result["data"] = base64Encode(data, byte_size); + result["dtype"] = dtypeToString(tensor.dType()); + + // Add shape + auto shape_ref = tensor.shape(); + result["shape"] = std::vector(shape_ref.begin(), shape_ref.end()); + + return result; +} + +void JSONDB::validateEdgeIndex(const AMSTensor& edge_index, int64_t num_nodes) +{ + auto shape_ref = edge_index.shape(); + if (shape_ref.size() != 2 || shape_ref[0] != 2) { + std::ostringstream oss; + oss << "edge_index must have shape [2, E], got ["; + for (size_t i = 0; i < shape_ref.size(); ++i) { + oss << shape_ref[i]; + if (i < shape_ref.size() - 1) oss << ", "; + } + oss << "]"; + THROW(std::invalid_argument, oss.str().c_str()); + } + + // Check dtype is int64 + if (edge_index.dType() != AMS_INT64) { + THROW(std::invalid_argument, + "edge_index must have dtype int64"); + } + + // Validate indices are in range + const int64_t* indices = edge_index.data(); + int64_t num_edges = shape_ref[1]; + + for (int64_t i = 0; i < 2 * num_edges; ++i) { + if (indices[i] < 0 || indices[i] >= num_nodes) { + std::ostringstream oss; + oss << "edge_index contains out-of-range index: " << indices[i] + << " (num_nodes=" << num_nodes << ")"; + THROW(std::invalid_argument, oss.str().c_str()); + } + } + + // Check for self-loops + for (int64_t i = 0; i < num_edges; ++i) { + if (indices[i] == indices[i + num_edges]) { + std::ostringstream oss; + oss << "edge_index contains self-loop at edge " << i << ": " + << indices[i] << " -> " << indices[i + num_edges]; + THROW(std::invalid_argument, oss.str().c_str()); + } + } +} + +// ---------------------------------------------------------------------- +// Store methods +// ---------------------------------------------------------------------- + +void JSONDB::store(ArrayRef Inputs, + ArrayRef Outputs) +{ + // Create case directory + std::ostringstream case_name; + case_name << "case_" << std::setw(6) << std::setfill('0') << case_counter_; + std::string case_dir = case_name.str(); + + nlohmann::json case_json; + case_json["name"] = case_dir; + case_json["case_index"] = case_counter_; + + nlohmann::json tensors_json; + + // Store inputs + for (size_t i = 0; i < Inputs.size(); ++i) { + std::ostringstream tensor_name; + tensor_name << "input_" << i; + std::string name = tensor_name.str(); + + if (json_mode_ == "binary") { + std::string rel_path = case_dir + "/" + name + ".bin"; + size_t byte_size = writeBinaryTensor(Inputs[i], rel_path); + + auto sizes = Inputs[i].sizes(); + std::vector shape(sizes.begin(), sizes.end()); + + tensors_json[name] = nlohmann::json{{"path", rel_path}, + {"dtype", torchDTypeToString(Inputs[i].scalar_type())}, + {"shape", shape}, + {"byte_size", byte_size}}; + } else { + // Pure JSON mode - not yet fully implemented for torch tensors + THROW(std::runtime_error, + "Pure JSON mode not yet implemented for tensor storage"); + } + } + + // Store outputs + for (size_t i = 0; i < Outputs.size(); ++i) { + std::ostringstream tensor_name; + tensor_name << "output_" << i; + std::string name = tensor_name.str(); + + if (json_mode_ == "binary") { + std::string rel_path = case_dir + "/" + name + ".bin"; + size_t byte_size = writeBinaryTensor(Outputs[i], rel_path); + + auto sizes = Outputs[i].sizes(); + std::vector shape(sizes.begin(), sizes.end()); + + tensors_json[name] = nlohmann::json{{"path", rel_path}, + {"dtype", torchDTypeToString(Outputs[i].scalar_type())}, + {"shape", shape}, + {"byte_size", byte_size}}; + } + } + + case_json["tensors"] = tensors_json; + cases_.push_back(case_json); + + case_counter_++; + + AMS_DBG(JSONDB, + "Stored tensor data for case {} ({} inputs, {} outputs)", + case_dir, + Inputs.size(), + Outputs.size()); +} + +void JSONDB::store(const ams::AMSHomogeneousGraph& graph, + const ams::AMSHomogeneousGraphFields& outputs) +{ + // Create case directory + std::ostringstream case_name; + case_name << "step_" << std::setw(6) << std::setfill('0') << case_counter_; + std::string case_dir = case_name.str(); + + // Extract graph dimensions + auto node_shape = graph.node_features.shape(); + auto edge_shape = graph.edge_index.shape(); + + int64_t num_nodes = node_shape[0]; + int64_t num_edges = edge_shape[1]; + int64_t node_feature_dim = node_shape[1]; + int64_t edge_feature_dim = 0; + int64_t global_feature_dim = 0; + + if (graph.edge_features.elements() > 0) { + auto ef_shape = graph.edge_features.shape(); + edge_feature_dim = ef_shape[1]; + } + + if (graph.global_features.has_value() && + graph.global_features.value().elements() > 0) { + auto gf_shape = graph.global_features.value().shape(); + global_feature_dim = gf_shape[1]; + } + + // Validate edge_index + validateEdgeIndex(graph.edge_index, num_nodes); + + // Build case metadata + nlohmann::json case_json; + case_json["name"] = case_dir; + case_json["step_index"] = case_counter_; + case_json["num_nodes"] = num_nodes; + case_json["num_edges"] = num_edges; + case_json["node_feature_dim"] = node_feature_dim; + case_json["edge_feature_dim"] = edge_feature_dim; + case_json["global_feature_dim"] = global_feature_dim; + + nlohmann::json tensors_json; + + // Write node_features + if (json_mode_ == "binary") { + std::string rel_path = case_dir + "/node_features.bin"; + size_t byte_size = writeBinaryTensor(graph.node_features, rel_path); + + tensors_json["node_features"] = { + {"path", rel_path}, + {"dtype", dtypeToString(graph.node_features.dType())}, + {"shape", std::vector{num_nodes, node_feature_dim}}, + {"byte_size", byte_size}}; + } + + // Write edge_index + if (json_mode_ == "binary") { + std::string rel_path = case_dir + "/edge_index.bin"; + size_t byte_size = writeBinaryTensor(graph.edge_index, rel_path); + + tensors_json["edge_index"] = { + {"path", rel_path}, + {"dtype", "int64"}, + {"shape", std::vector{2, num_edges}}, + {"byte_size", byte_size}}; + } + + // Write edge_features + if (edge_feature_dim > 0) { + if (json_mode_ == "binary") { + std::string rel_path = case_dir + "/edge_features.bin"; + size_t byte_size = writeBinaryTensor(graph.edge_features, rel_path); + + tensors_json["edge_features"] = { + {"path", rel_path}, + {"dtype", dtypeToString(graph.edge_features.dType())}, + {"shape", std::vector{num_edges, edge_feature_dim}}, + {"byte_size", byte_size}}; + } + } + + // Write global_features + if (global_feature_dim > 0 && graph.global_features.has_value()) { + if (json_mode_ == "binary") { + std::string rel_path = case_dir + "/global_features.bin"; + size_t byte_size = + writeBinaryTensor(graph.global_features.value(), rel_path); + + tensors_json["global_features"] = { + {"path", rel_path}, + {"dtype", dtypeToString(graph.global_features.value().dType())}, + {"shape", std::vector{1, global_feature_dim}}, + {"byte_size", byte_size}}; + } + } + + // Write targets from outputs.node_fields + // For now, we look for specific known target names + // TODO: Make this more generic with iterator support in AMSTensorFieldMap + int target_dim = 0; + + // Check for "delta_u" target (heat_equation convention) + const AMSTensor* delta_u = outputs.node_fields.find("delta_u"); + if (delta_u != nullptr) { + auto t_shape = delta_u->shape(); + target_dim = (t_shape.size() > 1) ? t_shape[1] : 1; + + if (json_mode_ == "binary") { + std::string rel_path = case_dir + "/target_delta_u.bin"; + size_t byte_size = writeBinaryTensor(*delta_u, rel_path); + + tensors_json["target_delta_u"] = { + {"path", rel_path}, + {"dtype", dtypeToString(delta_u->dType())}, + {"shape", std::vector{num_nodes, target_dim}}, + {"byte_size", byte_size}}; + } + } + + // Add target_dim to case metadata + case_json["target_dim"] = target_dim; + + case_json["tensors"] = tensors_json; + cases_.push_back(case_json); + + case_counter_++; + + AMS_DBG(JSONDB, + "Stored graph data for step {} ({} nodes, {} edges)", + case_dir, + num_nodes, + num_edges); +} + +void JSONDB::store(const ams::AMSHeterogeneousGraph& graph, + const ams::AMSHeterogeneousGraphFields& outputs) +{ + // Heterogeneous graph storage not yet implemented + THROW(std::runtime_error, + "Heterogeneous graph storage not yet implemented in JSONDB"); +} + +void JSONDB::finalize() +{ + if (finalized_) { + AMS_DBG(JSONDB, "Manifest already finalized, skipping"); + return; + } + + // Build complete manifest + nlohmann::json manifest; + manifest["format_version"] = 1; + manifest["endianness"] = "little"; + + // Add metadata if set + if (!metadata_.is_null()) { + manifest["metadata"] = metadata_; + } + + // Add feature names if set + if (!feature_names_.is_null()) { + manifest["feature_names"] = feature_names_; + } + + // Add all cases + manifest["cases"] = cases_; + + // Write manifest.json + fs::path manifest_path = fs::path(fp) / "manifest.json"; + std::ofstream manifest_file(manifest_path.string()); + if (!manifest_file.is_open()) { + THROW(std::runtime_error, + ("Failed to open manifest file for writing: " + + manifest_path.string()) + .c_str()); + } + + manifest_file << std::setw(2) << manifest << std::endl; + manifest_file.close(); + + finalized_ = true; + + AMS_DBG(JSONDB, + "Finalized manifest with {} cases at '{}'", + cases_.size(), + manifest_path.string()); +} diff --git a/src/AMSlib/wf/jsondb.hpp b/src/AMSlib/wf/jsondb.hpp new file mode 100644 index 00000000..b68322f8 --- /dev/null +++ b/src/AMSlib/wf/jsondb.hpp @@ -0,0 +1,191 @@ +/* + * Copyright 2021-2023 Lawrence Livermore National Security, LLC and other + * AMSLib Project Developers + * + * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + */ + +#ifndef __AMS_JSON_DB__ +#define __AMS_JSON_DB__ + +#include + +#include +#include +#include +#include + +#include "AMSGraph.hpp" +#include "AMSTensor.hpp" +#include "wf/basedb.hpp" + +namespace fs = std::experimental::filesystem; + +namespace ams +{ +namespace db +{ + +/** + * @brief JSON database backend for storing both tensor and graph data. + * + * This class provides a JSON-based storage backend that can serialize: + * - Flat tensors (via store(ArrayRef, ArrayRef)) + * - Homogeneous graphs (via store(AMSHomogeneousGraph, AMSHomogeneousGraphFields)) + * - Heterogeneous graphs (via store(AMSHeterogeneousGraph, AMSHeterogeneousGraphFields)) + * + * Supports two modes: + * - "binary": Binary tensor files + JSON manifest (default, efficient) + * - "json": Pure JSON with base64-encoded binary data (human-readable) + * + * Output format is compatible with PyTorch Geometric data loaders. + */ +class JSONDB final : public FileDB +{ +private: + /** @brief JSON output mode: "binary" or "json" */ + std::string json_mode_; + /** @brief Case counter for generating unique identifiers */ + int case_counter_; + /** @brief Accumulated case metadata for manifest */ + std::vector cases_; + /** @brief Metadata from application (physics config, etc.) */ + nlohmann::json metadata_; + /** @brief Feature names for each tensor type */ + nlohmann::json feature_names_; + /** @brief Whether manifest has been finalized */ + bool finalized_; + + /** + * @brief Write a tensor to binary file + * @param[in] tensor The AMSTensor to write + * @param[in] path Relative path from output directory + * @return Actual byte size written + */ + size_t writeBinaryTensor(const AMSTensor& tensor, const std::string& path); + + /** + * @brief Write a PyTorch tensor to binary file + * @param[in] tensor The torch::Tensor to write + * @param[in] path Relative path from output directory + * @return Actual byte size written + */ + size_t writeBinaryTensor(const torch::Tensor& tensor, + const std::string& path); + + /** + * @brief Encode tensor as base64 JSON (for pure JSON mode) + * @param[in] tensor The AMSTensor to encode + * @return JSON object with encoded data + */ + nlohmann::json encodeBase64Tensor(const AMSTensor& tensor); + + /** + * @brief Validate edge_index tensor format + * @param[in] edge_index Edge connectivity tensor [2, E] + * @param[in] num_nodes Number of nodes in graph + */ + void validateEdgeIndex(const AMSTensor& edge_index, int64_t num_nodes); + + /** + * @brief Convert AMSDType to string + * @param[in] dtype The AMS data type + * @return String representation ("float32", "float64", "int64") + */ + std::string dtypeToString(AMSDType dtype) const; + + /** + * @brief Convert torch::Dtype to string + * @param[in] dtype The PyTorch data type + * @return String representation + */ + std::string torchDTypeToString(torch::Dtype dtype) const; + + /** + * @brief Get byte size for a data type + * @param[in] dtype The data type enum + * @return Size in bytes + */ + size_t dtypeSize(AMSDType dtype) const; + +public: + /** + * @brief Construct a JSON database + * @param[in] path Directory path for output files + * @param[in] domain_name Name of the domain (used in filenames) + * @param[in] rId Rank ID for distributed execution + * @param[in] json_mode "binary" (default) or "json" for output mode + */ + JSONDB(std::string path, + std::string domain_name, + uint64_t rId, + std::string json_mode = "binary"); + + /** + * @brief Destructor - finalizes manifest if not already done + */ + ~JSONDB(); + + // Delete copy/move constructors + JSONDB(const JSONDB&) = delete; + JSONDB& operator=(const JSONDB&) = delete; + + /** + * @brief Store tensor data (inputs and outputs) + * @param[in] Inputs Vector of input tensors + * @param[in] Outputs Vector of output tensors + */ + void store(ArrayRef Inputs, + ArrayRef Outputs) override; + + /** + * @brief Store homogeneous graph data with outputs + * @param[in] graph Input graph structure and features + * @param[in] outputs Output/target fields for training + */ + void store(const ams::AMSHomogeneousGraph& graph, + const ams::AMSHomogeneousGraphFields& outputs) override; + + /** + * @brief Store heterogeneous graph data with outputs + * @param[in] graph Input graph structure and features + * @param[in] outputs Output/target fields for training + */ + void store(const ams::AMSHeterogeneousGraph& graph, + const ams::AMSHeterogeneousGraphFields& outputs) override; + + /** + * @brief Finalize and write manifest.json + */ + void finalize(); + + /** + * @brief Set application metadata + * @param[in] metadata JSON object with application-specific config + */ + void setMetadata(const nlohmann::json& metadata) { metadata_ = metadata; } + + /** + * @brief Set feature names for documentation + * @param[in] feature_names JSON object mapping tensor types to name lists + */ + void setFeatureNames(const nlohmann::json& feature_names) + { + feature_names_ = feature_names; + } + + /** + * @brief Database type identifier + */ + std::string type() override { return "json"; } + + /** + * @brief Database type enum + */ + AMSDBType dbType() override { return AMSDBType::AMS_JSON; } +}; + +} // namespace db +} // namespace ams + +#endif // __AMS_JSON_DB__ diff --git a/src/AMSlib/wf/workflow.hpp b/src/AMSlib/wf/workflow.hpp index 6ea872ef..34d448fe 100644 --- a/src/AMSlib/wf/workflow.hpp +++ b/src/AMSlib/wf/workflow.hpp @@ -102,21 +102,37 @@ class AMSWorkflow void storeGraphData(const ams::AMSHomogeneousGraph& graph, const ams::AMSHomogeneousGraphFields& outputs) { - // TODO: Implement graph storage when database supports it - // For now, this is a no-op placeholder - (void)graph; - (void)outputs; - AMS_DBG(Workflow, "Graph storage not yet implemented (homogeneous)"); + if (!DB) { + AMS_WARNING(Workflow, "Cannot store graph data: database not initialized"); + return; + } + + try { + DB->store(graph, outputs); + AMS_DBG(Workflow, "Successfully stored homogeneous graph data"); + } catch (const std::exception& e) { + AMS_WARNING(Workflow, + "Failed to store homogeneous graph data: {}", + e.what()); + } } void storeGraphData(const ams::AMSHeterogeneousGraph& graph, const ams::AMSHeterogeneousGraphFields& outputs) { - // TODO: Implement graph storage when database supports it - // For now, this is a no-op placeholder - (void)graph; - (void)outputs; - AMS_DBG(Workflow, "Graph storage not yet implemented (heterogeneous)"); + if (!DB) { + AMS_WARNING(Workflow, "Cannot store graph data: database not initialized"); + return; + } + + try { + DB->store(graph, outputs); + AMS_DBG(Workflow, "Successfully stored heterogeneous graph data"); + } catch (const std::exception& e) { + AMS_WARNING(Workflow, + "Failed to store heterogeneous graph data: {}", + e.what()); + } } /** \brief Check if we can perform a surrogate model update. @@ -444,6 +460,59 @@ class AMSWorkflow CALIPER(CALI_MARK_END("AMSEvaluate");) } + // Graph-based evaluate methods (mirror tensor pattern) + void evaluate(HomogeneousGraphDomainFn CallBack, + const AMSHomogeneousGraph& graph_input, + AMSHomogeneousGraphFields& outputs) + { + CALIPER(CALI_MARK_BEGIN("AMSEvaluateGraph");) + + // Try surrogate first + bool surrogate_used = tryGraphSurrogate(this, graph_input, outputs); + if (surrogate_used) { + CALIPER(CALI_MARK_END("AMSEvaluateGraph");) + return; + } + + // Fallback to physics + CALIPER(CALI_MARK_BEGIN("PHYSICS MODULE");) + CallBack(graph_input, outputs); + CALIPER(CALI_MARK_END("PHYSICS MODULE");) + + // Store data after physics computation + if (DB) { + storeGraphData(graph_input, outputs); + } + + CALIPER(CALI_MARK_END("AMSEvaluateGraph");) + } + + void evaluate(HeterogeneousGraphDomainFn CallBack, + const AMSHeterogeneousGraph& graph_input, + AMSHeterogeneousGraphFields& outputs) + { + CALIPER(CALI_MARK_BEGIN("AMSEvaluateGraph");) + + // Try surrogate first + bool surrogate_used = tryGraphSurrogate(this, graph_input, outputs); + if (surrogate_used) { + CALIPER(CALI_MARK_END("AMSEvaluateGraph");) + return; + } + + // Fallback to physics + CALIPER(CALI_MARK_BEGIN("PHYSICS MODULE");) + CallBack(graph_input, outputs); + CALIPER(CALI_MARK_END("PHYSICS MODULE");) + + // Store data after physics computation + if (DB) { + storeGraphData(graph_input, outputs); + } + + CALIPER(CALI_MARK_END("AMSEvaluateGraph");) + } + std::string getDBName() { if (!DB) return ""; diff --git a/tests/AMSlib/ams_interface/test_graph_workflow_storage.cpp b/tests/AMSlib/ams_interface/test_graph_workflow_storage.cpp new file mode 100644 index 00000000..b6b5589d --- /dev/null +++ b/tests/AMSlib/ams_interface/test_graph_workflow_storage.cpp @@ -0,0 +1,765 @@ +/* + * Copyright 2021-2023 Lawrence Livermore National Security, LLC and other + * AMSLib Project Developers + * + * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + */ + +#include + +#include +#include +#include + +#include "AMS.h" +#include "AMSGraph.hpp" +#include "AMSTensor.hpp" +#include "nlohmann/json.hpp" + +using namespace ams; +namespace fs = std::filesystem; + +// Helper to create contiguous strides from shape +static std::vector contiguousStrides(const std::vector& shape) +{ + std::vector strides(shape.size(), 1); + int64_t stride = 1; + for (std::size_t i = shape.size(); i-- > 0;) { + strides[i] = stride; + stride *= shape[i]; + } + return strides; +} + +// Helper to create tensor with automatic strides +template +static AMSTensor makeTensor(std::vector shape) +{ + std::vector strides = contiguousStrides(shape); + return AMSTensor::create(shape, strides, AMSResourceType::AMS_HOST); +} + +CATCH_TEST_CASE( + "Homogeneous graph forced physics stores data and reveals native schema", + "[wf][graph][storage]") +{ + // Use unique temp directory + fs::path test_dir = + fs::temp_directory_path() / "ams_graph_schema_discovery_test"; + fs::remove_all(test_dir); + fs::create_directories(test_dir); + + // Initialize AMS (from A0 investigation: AMSInit before AMSConfigureFSDatabase) + AMSInit(); + AMSConfigureFSDatabase(AMSDBType::AMS_JSON, test_dir.string().c_str()); + + // Recorder configuration: threshold < 0 forces physics, store_data=true enables DB + AMSCAbstrModel recorder = + AMSRegisterAbstractModel("heat_graph_schema", -1.0, "", true); + AMSExecutor executor = AMSCreateExecutor(recorder, 0, 1); + + // Build small test graph + const int64_t N = 10; // nodes + const int64_t E = 9; // edges + + // Node features [N, 3] - float32 + auto node_feat = makeTensor({N, 3}); + float* nf = node_feat.data(); + for (int64_t i = 0; i < N * 3; i++) { + nf[i] = static_cast(i) * 0.1f; + } + + // Edge index [2, E] - int64 (canonical) + auto edge_idx = makeTensor({2, E}); + int64_t* ei = edge_idx.data(); + for (int64_t e = 0; e < E; e++) { + ei[e] = e; // source in first half + ei[E + e] = (e + 1) % N; // destination in second half + } + + // Edge features [E, 2] - float32 + auto edge_feat = makeTensor({E, 2}); + float* ef = edge_feat.data(); + for (int64_t e = 0; e < E * 2; e++) { + ef[e] = static_cast(e) * 0.01f; + } + + // Global features [1, 2] - float32 + auto global_feat = makeTensor({1, 2}); + float* gf = global_feat.data(); + gf[0] = 0.01f; // dt + gf[1] = 0.123f; // time_n + + // Construct graph via public API + AMSHomogeneousGraph graph(std::move(node_feat), + std::move(edge_idx), + std::move(edge_feat), + std::move(global_feat)); + + AMSHomogeneousGraphFields outputs; + + // Physics callback that returns float64 delta_u + int callback_count = 0; + HomogeneousGraphDomainFn physics = + [&](const AMSHomogeneousGraph& g, AMSHomogeneousGraphFields& o) { + callback_count++; + + // Return float64 delta_u [N, 1] (MFEM precision) + const int64_t num_nodes = g.node_features.shape()[0]; + auto delta = makeTensor({num_nodes, 1}); + double* data = delta.data(); + for (int64_t i = 0; i < num_nodes; i++) { + data[i] = static_cast(i) * 0.123; + } + + o.node_fields.insert("delta_u", std::move(delta)); + }; + + // Execute with forced physics + storage + AMSExecute(executor, physics, graph, outputs); + + // Verify physics ran + CATCH_REQUIRE(callback_count == 1); + CATCH_REQUIRE(outputs.node_fields.find("delta_u") != nullptr); + + // Destroy executor to flush manifest (but don't call AMSFinalize) + AMSDestroyExecutor(executor); + + // ======================================================================== + // INSPECT ACTUAL JSONDB OUTPUT - source of truth for schema + // ======================================================================== + + CATCH_REQUIRE(fs::exists(test_dir / "manifest.json")); + + std::ifstream manifest_file(test_dir / "manifest.json"); + nlohmann::json manifest; + manifest_file >> manifest; + + // Document native schema structure + std::cout << "\n=== NATIVE JSONDB SCHEMA (A1 Discovery) ===\n"; + std::cout << manifest.dump(2) << std::endl; + std::cout << "==========================================\n" << std::endl; + + // Verify essential structure exists + CATCH_REQUIRE(manifest.contains("format_version")); + CATCH_REQUIRE(manifest.contains("endianness")); + CATCH_REQUIRE(manifest.contains("cases")); + CATCH_REQUIRE(manifest["cases"].is_array()); + CATCH_REQUIRE(manifest["cases"].size() == 1); + + auto case0 = manifest["cases"][0]; + CATCH_REQUIRE(case0.contains("tensors")); + + // Verify all graph components stored + CATCH_REQUIRE(case0["tensors"].contains("node_features")); + CATCH_REQUIRE(case0["tensors"].contains("edge_index")); + CATCH_REQUIRE(case0["tensors"].contains("edge_features")); + + // CRITICAL: Verify global_features stored + CATCH_REQUIRE(case0["tensors"].contains("global_features")); + + // Verify output field stored + // Note: Field name may be "delta_u" or "target_delta_u" - discover actual + bool has_delta_u = case0["tensors"].contains("delta_u") || + case0["tensors"].contains("target_delta_u") || + case0["tensors"].contains("node:delta_u"); + CATCH_REQUIRE(has_delta_u); + + // Verify dtypes preserved + CATCH_REQUIRE(case0["tensors"]["node_features"]["dtype"].get() == + "float32"); + CATCH_REQUIRE(case0["tensors"]["edge_index"]["dtype"].get() == + "int64"); + CATCH_REQUIRE(case0["tensors"]["edge_features"]["dtype"].get() == + "float32"); + CATCH_REQUIRE(case0["tensors"]["global_features"]["dtype"].get() == + "float32"); + + // Verify float64 delta_u (CRITICAL for MFEM precision) + std::string delta_key = "delta_u"; + if (case0["tensors"].contains("target_delta_u")) { + delta_key = "target_delta_u"; + } else if (case0["tensors"].contains("node:delta_u")) { + delta_key = "node:delta_u"; + } + CATCH_REQUIRE(case0["tensors"][delta_key]["dtype"].get() == + "float64"); + + // Verify paths are relative to dataset root + std::string node_path = + case0["tensors"]["node_features"]["path"].get(); + CATCH_REQUIRE(!fs::path(node_path).is_absolute()); + CATCH_REQUIRE(fs::exists(test_dir / node_path)); + + // Note: Not destroying executor or calling AMSFinalize to avoid lifecycle + // issues between tests. Cleanup happens at process exit. + // fs::remove_all(test_dir); // Keep for manual inspection + + std::cout << "A1 schema discovery test PASSED. Review manifest output above " + "before proceeding to A2." + << std::endl; +} + +// ============================================================================ +// A2 Tests: Complete Storage Test Coverage +// ============================================================================ + +CATCH_TEST_CASE("store_data=false returns physics output with zero recorded cases", + "[wf][graph][storage]") +{ + fs::path test_dir = + fs::temp_directory_path() / "ams_graph_no_storage_test"; + fs::remove_all(test_dir); + fs::create_directories(test_dir); + + AMSInit(); + AMSConfigureFSDatabase(AMSDBType::AMS_JSON, test_dir.string().c_str()); + + // Register with store_data=false + AMSCAbstrModel no_store = + AMSRegisterAbstractModel("no_storage_domain", -1.0, "", false); + AMSExecutor executor = AMSCreateExecutor(no_store, 0, 1); + + // Build simple graph + auto node_feat = makeTensor({5, 2}); + auto edge_idx = makeTensor({2, 4}); + auto edge_feat = makeTensor({4, 1}); + auto global_feat = makeTensor({1, 2}); + + // Fill with identifiable values + float* nf = node_feat.data(); + for (int i = 0; i < 10; i++) nf[i] = static_cast(i) + 100.0f; + + // Fill edge index (4 edges for 5 nodes) + int64_t* ei = edge_idx.data(); + ei[0] = 0; ei[1] = 1; ei[2] = 2; ei[3] = 3; // sources + ei[4] = 1; ei[5] = 2; ei[6] = 3; ei[7] = 4; // destinations + + // Fill edge features + float* ef = edge_feat.data(); + for (int i = 0; i < 4; i++) ef[i] = static_cast(i) * 0.5f; + + // Fill global features + float* gf = global_feat.data(); + gf[0] = 0.01f; + gf[1] = 0.02f; + + AMSHomogeneousGraph graph(std::move(node_feat), + std::move(edge_idx), + std::move(edge_feat), + std::move(global_feat)); + + AMSHomogeneousGraphFields outputs; + + int callback_count = 0; + HomogeneousGraphDomainFn physics = + [&](const AMSHomogeneousGraph& g, AMSHomogeneousGraphFields& o) { + callback_count++; + const int64_t N = g.node_features.shape()[0]; + auto delta = makeTensor({N, 1}); + double* data = delta.data(); + for (int64_t i = 0; i < N; i++) data[i] = static_cast(i) + 200.0; + o.node_fields.insert("delta_u", std::move(delta)); + }; + + AMSExecute(executor, physics, graph, outputs); + + // Verify physics ran and returned output + CATCH_REQUIRE(callback_count == 1); + CATCH_REQUIRE(outputs.node_fields.find("delta_u") != nullptr); + + const auto& delta = outputs.node_fields.at("delta_u"); + CATCH_REQUIRE(delta.shape()[0] == 5); + CATCH_REQUIRE(delta.dType() == ams::AMS_DOUBLE); + + // Verify output values + const double* delta_data = delta.data(); + for (int i = 0; i < 5; i++) { + CATCH_REQUIRE(delta_data[i] == static_cast(i) + 200.0); + } + + // Note: Not destroying executor to avoid triggering AMSFinalize between tests + // The executor will be cleaned up at program exit + + // Verify NO manifest created (store_data=false) + CATCH_REQUIRE(!fs::exists(test_dir / "manifest.json")); + + fs::remove_all(test_dir); +} + +CATCH_TEST_CASE("Multiple calls accumulate cases with distinguishable values", + "[wf][graph][storage]") +{ + fs::path test_dir = + fs::temp_directory_path() / "ams_graph_accumulation_test"; + fs::remove_all(test_dir); + fs::create_directories(test_dir); + + AMSInit(); + AMSConfigureFSDatabase(AMSDBType::AMS_JSON, test_dir.string().c_str()); + + AMSCAbstrModel recorder = + AMSRegisterAbstractModel("accumulation_domain", -1.0, "", true); + AMSExecutor executor = AMSCreateExecutor(recorder, 0, 1); + + const int num_calls = 3; + int total_callbacks = 0; + + for (int call = 0; call < num_calls; call++) { + // Build graph with call-specific values + auto node_feat = makeTensor({4, 2}); + auto edge_idx = makeTensor({2, 3}); + auto edge_feat = makeTensor({3, 1}); + auto global_feat = makeTensor({1, 1}); + + // Fill node features with call-specific values + float* nf = node_feat.data(); + for (int i = 0; i < 8; i++) { + nf[i] = static_cast(call * 1000 + i); + } + + // Fill edge index with valid connectivity (3 edges for 4 nodes) + int64_t* ei = edge_idx.data(); + ei[0] = 0; ei[1] = 1; ei[2] = 2; // sources + ei[3] = 1; ei[4] = 2; ei[5] = 3; // destinations + + // Fill edge features + float* ef = edge_feat.data(); + for (int i = 0; i < 3; i++) { + ef[i] = static_cast(call * 100 + i); + } + + // Fill global features with call identifier + float* gf = global_feat.data(); + gf[0] = static_cast(call); + + AMSHomogeneousGraph graph(std::move(node_feat), + std::move(edge_idx), + std::move(edge_feat), + std::move(global_feat)); + + AMSHomogeneousGraphFields outputs; + + HomogeneousGraphDomainFn physics = + [&, call_id = call](const AMSHomogeneousGraph& g, + AMSHomogeneousGraphFields& o) { + total_callbacks++; + const int64_t N = g.node_features.shape()[0]; + auto delta = makeTensor({N, 1}); + double* data = delta.data(); + for (int64_t i = 0; i < N; i++) { + data[i] = static_cast(call_id * 100 + i); + } + o.node_fields.insert("delta_u", std::move(delta)); + }; + + AMSExecute(executor, physics, graph, outputs); + } + + CATCH_REQUIRE(total_callbacks == num_calls); + + // Destroy executor to flush manifest (but don't call AMSFinalize) + AMSDestroyExecutor(executor); + + // Verify manifest exists with correct case count + CATCH_REQUIRE(fs::exists(test_dir / "manifest.json")); + + std::ifstream manifest_file(test_dir / "manifest.json"); + nlohmann::json manifest; + manifest_file >> manifest; + + CATCH_REQUIRE(manifest["cases"].size() == num_calls); + + // Verify each case has distinguishable values + for (int call = 0; call < num_calls; call++) { + auto case_entry = manifest["cases"][call]; + CATCH_REQUIRE(case_entry["step_index"].get() == call); + + // Load and verify global feature (call-specific marker) + std::string global_path = case_entry["tensors"]["global_features"]["path"]; + fs::path full_path = test_dir / global_path; + CATCH_REQUIRE(fs::exists(full_path)); + + std::ifstream f(full_path, std::ios::binary); + float global_val; + f.read(reinterpret_cast(&global_val), sizeof(float)); + CATCH_REQUIRE(global_val == static_cast(call)); + + // Verify target output exists + CATCH_REQUIRE(case_entry["tensors"].contains("target_delta_u")); + } + + fs::remove_all(test_dir); +} + +CATCH_TEST_CASE("Heterogeneous graph typed storage", + "[.][wf][graph][storage][heterogeneous][future]") +{ + fs::path test_dir = + fs::temp_directory_path() / "ams_graph_heterogeneous_test"; + fs::remove_all(test_dir); + fs::create_directories(test_dir); + + AMSInit(); + AMSConfigureFSDatabase(AMSDBType::AMS_JSON, test_dir.string().c_str()); + + AMSCAbstrModel recorder = + AMSRegisterAbstractModel("hetero_domain", -1.0, "", true); + AMSExecutor executor = AMSCreateExecutor(recorder, 0, 1); + + // Build heterogeneous graph with two node types + AMSHeterogeneousGraph graph; + + // Node type "fluid": 5 nodes, 3 features + auto& fluid_store = graph.getOrCreateNodeStore("fluid"); + auto fluid_feat = makeTensor({5, 3}); + float* ff = fluid_feat.data(); + for (int i = 0; i < 15; i++) ff[i] = static_cast(i) * 0.1f; + insertTensor(fluid_store, "features", std::move(fluid_feat)); + + // Node type "solid": 3 nodes, 2 features + auto& solid_store = graph.getOrCreateNodeStore("solid"); + auto solid_feat = makeTensor({3, 2}); + float* sf = solid_feat.data(); + for (int i = 0; i < 6; i++) sf[i] = static_cast(i) * 0.2f; + insertTensor(solid_store, "features", std::move(solid_feat)); + + // Edge type: fluid->solid (4 edges) + EdgeType edge_type("fluid", "interacts", "solid"); + auto& edge_store = graph.getOrCreateEdgeStore(edge_type); + auto edge_idx = makeTensor({2, 4}); + int64_t* ei = edge_idx.data(); + ei[0] = 0; ei[1] = 1; ei[2] = 2; ei[3] = 3; // fluid nodes (sources) + ei[4] = 0; ei[5] = 1; ei[6] = 1; ei[7] = 2; // solid nodes (destinations) + insertTensor(edge_store, "edge_index", std::move(edge_idx)); + + // Global features + auto global_feat = makeTensor({1, 2}); + float* gf = global_feat.data(); + gf[0] = 0.01f; // timestep + gf[1] = 0.5f; // time + insertTensor(graph.global_store, "time", std::move(global_feat)); + + AMSHeterogeneousGraphFields outputs; + + int callback_count = 0; + HeterogeneousGraphDomainFn physics = + [&](const AMSHeterogeneousGraph& g, AMSHeterogeneousGraphFields& o) { + callback_count++; + + // Output for fluid nodes + auto& fluid_out = o.getOrCreateNodeStore("fluid"); + auto fluid_delta = makeTensor({5, 1}); + double* fd = fluid_delta.data(); + for (int i = 0; i < 5; i++) fd[i] = static_cast(i) + 10.0; + fluid_out.insert("delta_u", std::move(fluid_delta)); + + // Output for solid nodes + auto& solid_out = o.getOrCreateNodeStore("solid"); + auto solid_delta = makeTensor({3, 1}); + double* sd = solid_delta.data(); + for (int i = 0; i < 3; i++) sd[i] = static_cast(i) + 20.0; + solid_out.insert("delta_u", std::move(solid_delta)); + }; + + AMSExecute(executor, physics, graph, outputs); + + CATCH_REQUIRE(callback_count == 1); + CATCH_REQUIRE(outputs.node_stores.find("fluid") != nullptr); + CATCH_REQUIRE(outputs.node_stores.find("solid") != nullptr); + + // Note: Not destroying executor to avoid triggering AMSFinalize between tests + // The executor will be cleaned up at program exit + + // Verify heterogeneous storage + CATCH_REQUIRE(fs::exists(test_dir / "manifest.json")); + + std::ifstream manifest_file(test_dir / "manifest.json"); + nlohmann::json manifest; + manifest_file >> manifest; + + CATCH_REQUIRE(manifest["cases"].size() == 1); + auto case0 = manifest["cases"][0]; + + // Verify typed node stores stored + CATCH_REQUIRE(case0["tensors"].contains("node_fluid__features")); + CATCH_REQUIRE(case0["tensors"].contains("node_solid__features")); + + // Verify typed outputs stored with target_ prefix + CATCH_REQUIRE(case0["tensors"].contains("target_node_fluid__delta_u")); + CATCH_REQUIRE(case0["tensors"].contains("target_node_solid__delta_u")); + + // Verify dtypes + CATCH_REQUIRE(case0["tensors"]["target_node_fluid__delta_u"]["dtype"] == "float64"); + CATCH_REQUIRE(case0["tensors"]["target_node_solid__delta_u"]["dtype"] == "float64"); + + // Verify shapes + auto fluid_shape = case0["tensors"]["target_node_fluid__delta_u"]["shape"]; + CATCH_REQUIRE(fluid_shape[0] == 5); + CATCH_REQUIRE(fluid_shape[1] == 1); + + auto solid_shape = case0["tensors"]["target_node_solid__delta_u"]["shape"]; + CATCH_REQUIRE(solid_shape[0] == 3); + CATCH_REQUIRE(solid_shape[1] == 1); + + fs::remove_all(test_dir); +} + +CATCH_TEST_CASE("Surrogate success: zero callbacks and zero stored cases", + "[.][wf][graph][storage][surrogate][future]") +{ + // This test requires a working surrogate model that returns low uncertainty. + // For now, verify the no-model case (which always runs physics). + // TODO: Add actual surrogate test when model fixture is available. + + fs::path test_dir = + fs::temp_directory_path() / "ams_graph_surrogate_success_test"; + fs::remove_all(test_dir); + fs::create_directories(test_dir); + + AMSInit(); + AMSConfigureFSDatabase(AMSDBType::AMS_JSON, test_dir.string().c_str()); + + // No model (empty path) means surrogate cannot succeed + AMSCAbstrModel no_model = + AMSRegisterAbstractModel("surrogate_domain", 0.5, "", true); + AMSExecutor executor = AMSCreateExecutor(no_model, 0, 1); + + auto node_feat = makeTensor({3, 2}); + auto edge_idx = makeTensor({2, 2}); + auto edge_feat = makeTensor({2, 1}); + auto global_feat = makeTensor({1, 1}); + + // Fill all tensors + float* nf = node_feat.data(); + for (int i = 0; i < 6; i++) nf[i] = static_cast(i); + + int64_t* ei = edge_idx.data(); + ei[0] = 0; ei[1] = 1; // sources + ei[2] = 1; ei[3] = 2; // destinations + + float* ef = edge_feat.data(); + ef[0] = 0.5f; ef[1] = 1.0f; + + float* gf = global_feat.data(); + gf[0] = 1.0f; + + AMSHomogeneousGraph graph(std::move(node_feat), + std::move(edge_idx), + std::move(edge_feat), + std::move(global_feat)); + + AMSHomogeneousGraphFields outputs; + + int callback_count = 0; + HomogeneousGraphDomainFn physics = + [&](const AMSHomogeneousGraph& g, AMSHomogeneousGraphFields& o) { + callback_count++; + auto delta = makeTensor({3, 1}); + o.node_fields.insert("delta_u", std::move(delta)); + }; + + AMSExecute(executor, physics, graph, outputs); + + // Without model, physics runs + CATCH_REQUIRE(callback_count == 1); + CATCH_REQUIRE(outputs.node_fields.find("delta_u") != nullptr); + + // Note: Not destroying executor to avoid triggering AMSFinalize between tests + // The executor will be cleaned up at program exit + + // Note: This test will be updated when surrogate model fixture is available + // Expected behavior with working surrogate: + // - callback_count == 0 (surrogate used) + // - outputs populated by surrogate + // - manifest contains 0 cases (no fallback storage) + + fs::remove_all(test_dir); +} + +CATCH_TEST_CASE("No database configured: physics output returned without crash", + "[wf][graph][storage]") +{ + // No AMSConfigureFSDatabase call - DB should be null + + AMSInit(); + + // Register without configuring database (store_data=false → no DB needed) + AMSCAbstrModel no_db = + AMSRegisterAbstractModel("no_db_domain", -1.0, "", false); + AMSExecutor executor = AMSCreateExecutor(no_db, 0, 1); + + auto node_feat = makeTensor({4, 2}); + auto edge_idx = makeTensor({2, 3}); + auto edge_feat = makeTensor({3, 1}); + auto global_feat = makeTensor({1, 1}); + + // Fill all tensors + float* nf = node_feat.data(); + for (int i = 0; i < 8; i++) nf[i] = static_cast(i); + + int64_t* ei = edge_idx.data(); + ei[0] = 0; ei[1] = 1; ei[2] = 2; // sources + ei[3] = 1; ei[4] = 2; ei[5] = 3; // destinations + + float* ef = edge_feat.data(); + for (int i = 0; i < 3; i++) ef[i] = static_cast(i) * 0.3f; + + float* gf = global_feat.data(); + gf[0] = 1.0f; + + AMSHomogeneousGraph graph(std::move(node_feat), + std::move(edge_idx), + std::move(edge_feat), + std::move(global_feat)); + + AMSHomogeneousGraphFields outputs; + + int callback_count = 0; + HomogeneousGraphDomainFn physics = + [&](const AMSHomogeneousGraph& g, AMSHomogeneousGraphFields& o) { + callback_count++; + const int64_t N = g.node_features.shape()[0]; + auto delta = makeTensor({N, 1}); + double* data = delta.data(); + for (int64_t i = 0; i < N; i++) data[i] = static_cast(i) * 2.5; + o.node_fields.insert("delta_u", std::move(delta)); + }; + + // Should not crash even without DB + CATCH_REQUIRE_NOTHROW(AMSExecute(executor, physics, graph, outputs)); + + // Verify physics ran and output returned + CATCH_REQUIRE(callback_count == 1); + CATCH_REQUIRE(outputs.node_fields.find("delta_u") != nullptr); + + const auto& delta = outputs.node_fields.at("delta_u"); + CATCH_REQUIRE(delta.shape()[0] == 4); + + // Verify values + const double* delta_data = delta.data(); + for (int i = 0; i < 4; i++) { + CATCH_REQUIRE(delta_data[i] == static_cast(i) * 2.5); + } + + // Note: Not destroying executor to avoid triggering AMSFinalize between tests + // The executor will be cleaned up at program exit +} + +CATCH_TEST_CASE( + "Surrogate rejection: model runs but rejected, physics executes once, exact output stored", + "[wf][graph][storage][surrogate]") +{ + // This test validates the rejection path: when a surrogate model runs + // but returns high uncertainty (above threshold), physics should execute + // and store the fallback output exactly once. + // + // Current limitation: Requires a scripted model that returns high uncertainty. + // For now, test the no-model fallback path which exercises the same storage logic. + + fs::path test_dir = + fs::temp_directory_path() / "ams_graph_surrogate_reject_test"; + fs::remove_all(test_dir); + fs::create_directories(test_dir); + + AMSInit(); + AMSConfigureFSDatabase(AMSDBType::AMS_JSON, test_dir.string().c_str()); + + // No model path = immediate fallback to physics + // (Same storage path as uncertainty rejection would trigger) + AMSCAbstrModel fallback = + AMSRegisterAbstractModel("reject_domain", 0.1, "", true); + AMSExecutor executor = AMSCreateExecutor(fallback, 0, 1); + + auto node_feat = makeTensor({6, 2}); + auto edge_idx = makeTensor({2, 5}); + auto edge_feat = makeTensor({5, 1}); + auto global_feat = makeTensor({1, 1}); + + // Fill with specific values to verify exact storage + float* nf = node_feat.data(); + for (int i = 0; i < 12; i++) nf[i] = static_cast(i) * 1.5f; + + // Fill edge index (5 edges for 6 nodes) + int64_t* ei = edge_idx.data(); + ei[0] = 0; ei[1] = 1; ei[2] = 2; ei[3] = 3; ei[4] = 4; // sources + ei[5] = 1; ei[6] = 2; ei[7] = 3; ei[8] = 4; ei[9] = 5; // destinations + + // Fill edge features + float* ef = edge_feat.data(); + for (int i = 0; i < 5; i++) ef[i] = static_cast(i) * 0.25f; + + // Fill global features + float* gf = global_feat.data(); + gf[0] = 1.0f; + + AMSHomogeneousGraph graph(std::move(node_feat), + std::move(edge_idx), + std::move(edge_feat), + std::move(global_feat)); + + AMSHomogeneousGraphFields outputs; + + int callback_count = 0; + HomogeneousGraphDomainFn physics = + [&](const AMSHomogeneousGraph& g, AMSHomogeneousGraphFields& o) { + callback_count++; + const int64_t N = g.node_features.shape()[0]; + auto delta = makeTensor({N, 1}); + double* data = delta.data(); + for (int64_t i = 0; i < N; i++) data[i] = static_cast(i) * 3.7; + o.node_fields.insert("delta_u", std::move(delta)); + }; + + AMSExecute(executor, physics, graph, outputs); + + // Verify physics executed exactly once + CATCH_REQUIRE(callback_count == 1); + + // Verify output returned to caller + CATCH_REQUIRE(outputs.node_fields.find("delta_u") != nullptr); + const auto& delta = outputs.node_fields.at("delta_u"); + CATCH_REQUIRE(delta.shape()[0] == 6); + CATCH_REQUIRE(delta.dType() == ams::AMS_DOUBLE); + + // Destroy executor to flush manifest (but don't call AMSFinalize) + AMSDestroyExecutor(executor); + + // Verify exactly one case stored + CATCH_REQUIRE(fs::exists(test_dir / "manifest.json")); + + std::ifstream manifest_file(test_dir / "manifest.json"); + nlohmann::json manifest; + manifest_file >> manifest; + + CATCH_REQUIRE(manifest["cases"].size() == 1); + + auto case0 = manifest["cases"][0]; + + // Verify stored output matches exact physics output + std::string target_path = case0["tensors"]["target_delta_u"]["path"]; + fs::path full_path = test_dir / target_path; + CATCH_REQUIRE(fs::exists(full_path)); + + // Load stored values + std::ifstream f(full_path, std::ios::binary); + std::vector stored_values(6); + f.read(reinterpret_cast(stored_values.data()), 6 * sizeof(double)); + + // Verify exact match with physics output + for (int i = 0; i < 6; i++) { + double expected = static_cast(i) * 3.7; + CATCH_REQUIRE(std::abs(stored_values[i] - expected) < 1e-12); + } + + // Verify field name (runtime) vs storage name (target_ prefix) + // Runtime API: node_fields["delta_u"] + // Storage: "target_delta_u" in manifest + CATCH_REQUIRE(case0["tensors"].contains("target_delta_u")); + CATCH_REQUIRE(!case0["tensors"].contains("delta_u")); // No prefix in storage + + // Verify dtype preserved as float64 + CATCH_REQUIRE(case0["tensors"]["target_delta_u"]["dtype"] == "float64"); + + fs::remove_all(test_dir); +} From b250008882c282a8180a25ce7e2d85dcb5297a19 Mon Sep 17 00:00:00 2001 From: Yohann Dudouit Date: Fri, 28 Aug 2026 15:49:37 -0700 Subject: [PATCH 03/18] clang-format. --- src/AMSlib/wf/basedb.cpp | 5 +- src/AMSlib/wf/basedb.hpp | 3 +- src/AMSlib/wf/jsondb.cpp | 95 +++---- src/AMSlib/wf/jsondb.hpp | 3 +- src/AMSlib/wf/workflow.hpp | 6 +- .../test_graph_workflow_storage.cpp | 249 +++++++++++------- 6 files changed, 200 insertions(+), 161 deletions(-) diff --git a/src/AMSlib/wf/basedb.cpp b/src/AMSlib/wf/basedb.cpp index b5adaa1e..ac1a9cc3 100644 --- a/src/AMSlib/wf/basedb.cpp +++ b/src/AMSlib/wf/basedb.cpp @@ -1,8 +1,9 @@ -#include +#include "wf/basedb.hpp" + #include +#include #include "AMS.h" -#include "wf/basedb.hpp" #include "wf/jsondb.hpp" namespace ams diff --git a/src/AMSlib/wf/basedb.hpp b/src/AMSlib/wf/basedb.hpp index d54980ef..9bc64f59 100644 --- a/src/AMSlib/wf/basedb.hpp +++ b/src/AMSlib/wf/basedb.hpp @@ -156,7 +156,8 @@ class BaseDB const ams::AMSHeterogeneousGraphFields& outputs) { THROW(std::runtime_error, - (this->type() + " database does not support heterogeneous graph storage") + (this->type() + " database does not support heterogeneous graph " + "storage") .c_str()); } diff --git a/src/AMSlib/wf/jsondb.cpp b/src/AMSlib/wf/jsondb.cpp index 52edd159..71170571 100644 --- a/src/AMSlib/wf/jsondb.cpp +++ b/src/AMSlib/wf/jsondb.cpp @@ -54,13 +54,15 @@ std::string base64Encode(const uint8_t* data, size_t len) ((char_array_3[1] & 0x0f) << 2) + ((char_array_3[2] & 0xc0) >> 6); char_array_4[3] = char_array_3[2] & 0x3f; - for (i = 0; i < 4; i++) ret += base64_chars[char_array_4[i]]; + for (i = 0; i < 4; i++) + ret += base64_chars[char_array_4[i]]; i = 0; } } if (i) { - for (int j = i; j < 3; j++) char_array_3[j] = '\0'; + for (int j = i; j < 3; j++) + char_array_3[j] = '\0'; char_array_4[0] = (char_array_3[0] & 0xfc) >> 2; char_array_4[1] = @@ -68,9 +70,11 @@ std::string base64Encode(const uint8_t* data, size_t len) char_array_4[2] = ((char_array_3[1] & 0x0f) << 2) + ((char_array_3[2] & 0xc0) >> 6); - for (int j = 0; j < i + 1; j++) ret += base64_chars[char_array_4[j]]; + for (int j = 0; j < i + 1; j++) + ret += base64_chars[char_array_4[j]]; - while (i++ < 3) ret += '='; + while (i++ < 3) + ret += '='; } return ret; @@ -99,15 +103,11 @@ JSONDB::JSONDB(std::string path, if (json_mode_ != "binary" && json_mode_ != "json") { THROW(std::invalid_argument, - ("Invalid json_mode: " + json_mode_ + - ". Must be 'binary' or 'json'.") + ("Invalid json_mode: " + json_mode_ + ". Must be 'binary' or 'json'.") .c_str()); } - AMS_DBG(JSONDB, - "Created JSONDB at '{}' with mode '{}'", - fp, - json_mode_); + AMS_DBG(JSONDB, "Created JSONDB at '{}' with mode '{}'", fp, json_mode_); } JSONDB::~JSONDB() @@ -141,14 +141,10 @@ std::string JSONDB::dtypeToString(AMSDType dtype) const std::string JSONDB::torchDTypeToString(torch::Dtype dtype) const { - if (dtype == torch::kFloat32 || dtype == torch::kFloat) - return "float32"; - if (dtype == torch::kFloat64 || dtype == torch::kDouble) - return "float64"; - if (dtype == torch::kInt32) - return "int32"; - if (dtype == torch::kInt64 || dtype == torch::kLong) - return "int64"; + if (dtype == torch::kFloat32 || dtype == torch::kFloat) return "float32"; + if (dtype == torch::kFloat64 || dtype == torch::kDouble) return "float64"; + if (dtype == torch::kInt32) return "int32"; + if (dtype == torch::kInt64 || dtype == torch::kLong) return "int64"; return "unknown"; } @@ -169,7 +165,7 @@ size_t JSONDB::dtypeSize(AMSDType dtype) const } size_t JSONDB::writeBinaryTensor(const AMSTensor& tensor, - const std::string& path) + const std::string& path) { // Get tensor properties const void* data = tensor.raw_data(); @@ -179,12 +175,11 @@ size_t JSONDB::writeBinaryTensor(const AMSTensor& tensor, // Handle GPU tensors - copy to CPU first std::vector cpu_buffer; if (location != AMSResourceType::AMS_HOST) { - AMS_WARNING( - JSONDB, - "GPU tensor detected. Copying to CPU for serialization (not yet " - "implemented - will fail)."); - THROW(std::runtime_error, - "GPU tensor serialization not yet implemented"); + AMS_WARNING(JSONDB, + "GPU tensor detected. Copying to CPU for serialization (not " + "yet " + "implemented - will fail)."); + THROW(std::runtime_error, "GPU tensor serialization not yet implemented"); // TODO: Implement cudaMemcpy/hipMemcpy here } @@ -211,16 +206,13 @@ size_t JSONDB::writeBinaryTensor(const AMSTensor& tensor, file.write(static_cast(data), byte_size); file.close(); - AMS_DBG(JSONDB, - "Wrote binary tensor to '{}' ({} bytes)", - path, - byte_size); + AMS_DBG(JSONDB, "Wrote binary tensor to '{}' ({} bytes)", path, byte_size); return byte_size; } size_t JSONDB::writeBinaryTensor(const torch::Tensor& tensor, - const std::string& path) + const std::string& path) { // Ensure tensor is contiguous and on CPU torch::Tensor cpu_tensor = tensor.contiguous().cpu(); @@ -240,10 +232,7 @@ size_t JSONDB::writeBinaryTensor(const torch::Tensor& tensor, file.write(static_cast(cpu_tensor.data_ptr()), byte_size); file.close(); - AMS_DBG(JSONDB, - "Wrote PyTorch tensor to '{}' ({} bytes)", - path, - byte_size); + AMS_DBG(JSONDB, "Wrote PyTorch tensor to '{}' ({} bytes)", path, byte_size); return byte_size; } @@ -255,8 +244,7 @@ nlohmann::json JSONDB::encodeBase64Tensor(const AMSTensor& tensor) size_t byte_size = tensor.elements() * tensor.element_size(); // Handle GPU/non-contiguous tensors - if (tensor.location() != AMSResourceType::AMS_HOST || - !tensor.contiguous()) { + if (tensor.location() != AMSResourceType::AMS_HOST || !tensor.contiguous()) { THROW(std::runtime_error, "Base64 encoding only supports contiguous CPU tensors currently"); } @@ -289,8 +277,7 @@ void JSONDB::validateEdgeIndex(const AMSTensor& edge_index, int64_t num_nodes) // Check dtype is int64 if (edge_index.dType() != AMS_INT64) { - THROW(std::invalid_argument, - "edge_index must have dtype int64"); + THROW(std::invalid_argument, "edge_index must have dtype int64"); } // Validate indices are in range @@ -310,8 +297,8 @@ void JSONDB::validateEdgeIndex(const AMSTensor& edge_index, int64_t num_nodes) for (int64_t i = 0; i < num_edges; ++i) { if (indices[i] == indices[i + num_edges]) { std::ostringstream oss; - oss << "edge_index contains self-loop at edge " << i << ": " - << indices[i] << " -> " << indices[i + num_edges]; + oss << "edge_index contains self-loop at edge " << i << ": " << indices[i] + << " -> " << indices[i + num_edges]; THROW(std::invalid_argument, oss.str().c_str()); } } @@ -348,10 +335,11 @@ void JSONDB::store(ArrayRef Inputs, auto sizes = Inputs[i].sizes(); std::vector shape(sizes.begin(), sizes.end()); - tensors_json[name] = nlohmann::json{{"path", rel_path}, - {"dtype", torchDTypeToString(Inputs[i].scalar_type())}, - {"shape", shape}, - {"byte_size", byte_size}}; + tensors_json[name] = + nlohmann::json{{"path", rel_path}, + {"dtype", torchDTypeToString(Inputs[i].scalar_type())}, + {"shape", shape}, + {"byte_size", byte_size}}; } else { // Pure JSON mode - not yet fully implemented for torch tensors THROW(std::runtime_error, @@ -372,10 +360,12 @@ void JSONDB::store(ArrayRef Inputs, auto sizes = Outputs[i].sizes(); std::vector shape(sizes.begin(), sizes.end()); - tensors_json[name] = nlohmann::json{{"path", rel_path}, - {"dtype", torchDTypeToString(Outputs[i].scalar_type())}, - {"shape", shape}, - {"byte_size", byte_size}}; + tensors_json[name] = + nlohmann::json{{"path", rel_path}, + {"dtype", + torchDTypeToString(Outputs[i].scalar_type())}, + {"shape", shape}, + {"byte_size", byte_size}}; } } @@ -452,11 +442,10 @@ void JSONDB::store(const ams::AMSHomogeneousGraph& graph, std::string rel_path = case_dir + "/edge_index.bin"; size_t byte_size = writeBinaryTensor(graph.edge_index, rel_path); - tensors_json["edge_index"] = { - {"path", rel_path}, - {"dtype", "int64"}, - {"shape", std::vector{2, num_edges}}, - {"byte_size", byte_size}}; + tensors_json["edge_index"] = {{"path", rel_path}, + {"dtype", "int64"}, + {"shape", std::vector{2, num_edges}}, + {"byte_size", byte_size}}; } // Write edge_features diff --git a/src/AMSlib/wf/jsondb.hpp b/src/AMSlib/wf/jsondb.hpp index b68322f8..7505792f 100644 --- a/src/AMSlib/wf/jsondb.hpp +++ b/src/AMSlib/wf/jsondb.hpp @@ -8,10 +8,9 @@ #ifndef __AMS_JSON_DB__ #define __AMS_JSON_DB__ -#include - #include #include +#include #include #include diff --git a/src/AMSlib/wf/workflow.hpp b/src/AMSlib/wf/workflow.hpp index 34d448fe..b5e5ae5a 100644 --- a/src/AMSlib/wf/workflow.hpp +++ b/src/AMSlib/wf/workflow.hpp @@ -103,7 +103,8 @@ class AMSWorkflow const ams::AMSHomogeneousGraphFields& outputs) { if (!DB) { - AMS_WARNING(Workflow, "Cannot store graph data: database not initialized"); + AMS_WARNING(Workflow, + "Cannot store graph data: database not initialized"); return; } @@ -121,7 +122,8 @@ class AMSWorkflow const ams::AMSHeterogeneousGraphFields& outputs) { if (!DB) { - AMS_WARNING(Workflow, "Cannot store graph data: database not initialized"); + AMS_WARNING(Workflow, + "Cannot store graph data: database not initialized"); return; } diff --git a/tests/AMSlib/ams_interface/test_graph_workflow_storage.cpp b/tests/AMSlib/ams_interface/test_graph_workflow_storage.cpp index b6b5589d..96445286 100644 --- a/tests/AMSlib/ams_interface/test_graph_workflow_storage.cpp +++ b/tests/AMSlib/ams_interface/test_graph_workflow_storage.cpp @@ -6,7 +6,6 @@ */ #include - #include #include #include @@ -73,7 +72,7 @@ CATCH_TEST_CASE( auto edge_idx = makeTensor({2, E}); int64_t* ei = edge_idx.data(); for (int64_t e = 0; e < E; e++) { - ei[e] = e; // source in first half + ei[e] = e; // source in first half ei[E + e] = (e + 1) % N; // destination in second half } @@ -100,20 +99,20 @@ CATCH_TEST_CASE( // Physics callback that returns float64 delta_u int callback_count = 0; - HomogeneousGraphDomainFn physics = - [&](const AMSHomogeneousGraph& g, AMSHomogeneousGraphFields& o) { - callback_count++; - - // Return float64 delta_u [N, 1] (MFEM precision) - const int64_t num_nodes = g.node_features.shape()[0]; - auto delta = makeTensor({num_nodes, 1}); - double* data = delta.data(); - for (int64_t i = 0; i < num_nodes; i++) { - data[i] = static_cast(i) * 0.123; - } + HomogeneousGraphDomainFn physics = [&](const AMSHomogeneousGraph& g, + AMSHomogeneousGraphFields& o) { + callback_count++; + + // Return float64 delta_u [N, 1] (MFEM precision) + const int64_t num_nodes = g.node_features.shape()[0]; + auto delta = makeTensor({num_nodes, 1}); + double* data = delta.data(); + for (int64_t i = 0; i < num_nodes; i++) { + data[i] = static_cast(i) * 0.123; + } - o.node_fields.insert("delta_u", std::move(delta)); - }; + o.node_fields.insert("delta_u", std::move(delta)); + }; // Execute with forced physics + storage AMSExecute(executor, physics, graph, outputs); @@ -172,8 +171,9 @@ CATCH_TEST_CASE( "int64"); CATCH_REQUIRE(case0["tensors"]["edge_features"]["dtype"].get() == "float32"); - CATCH_REQUIRE(case0["tensors"]["global_features"]["dtype"].get() == - "float32"); + CATCH_REQUIRE( + case0["tensors"]["global_features"]["dtype"].get() == + "float32"); // Verify float64 delta_u (CRITICAL for MFEM precision) std::string delta_key = "delta_u"; @@ -204,11 +204,11 @@ CATCH_TEST_CASE( // A2 Tests: Complete Storage Test Coverage // ============================================================================ -CATCH_TEST_CASE("store_data=false returns physics output with zero recorded cases", - "[wf][graph][storage]") +CATCH_TEST_CASE( + "store_data=false returns physics output with zero recorded cases", + "[wf][graph][storage]") { - fs::path test_dir = - fs::temp_directory_path() / "ams_graph_no_storage_test"; + fs::path test_dir = fs::temp_directory_path() / "ams_graph_no_storage_test"; fs::remove_all(test_dir); fs::create_directories(test_dir); @@ -228,16 +228,24 @@ CATCH_TEST_CASE("store_data=false returns physics output with zero recorded case // Fill with identifiable values float* nf = node_feat.data(); - for (int i = 0; i < 10; i++) nf[i] = static_cast(i) + 100.0f; + for (int i = 0; i < 10; i++) + nf[i] = static_cast(i) + 100.0f; // Fill edge index (4 edges for 5 nodes) int64_t* ei = edge_idx.data(); - ei[0] = 0; ei[1] = 1; ei[2] = 2; ei[3] = 3; // sources - ei[4] = 1; ei[5] = 2; ei[6] = 3; ei[7] = 4; // destinations + ei[0] = 0; + ei[1] = 1; + ei[2] = 2; + ei[3] = 3; // sources + ei[4] = 1; + ei[5] = 2; + ei[6] = 3; + ei[7] = 4; // destinations // Fill edge features float* ef = edge_feat.data(); - for (int i = 0; i < 4; i++) ef[i] = static_cast(i) * 0.5f; + for (int i = 0; i < 4; i++) + ef[i] = static_cast(i) * 0.5f; // Fill global features float* gf = global_feat.data(); @@ -252,15 +260,16 @@ CATCH_TEST_CASE("store_data=false returns physics output with zero recorded case AMSHomogeneousGraphFields outputs; int callback_count = 0; - HomogeneousGraphDomainFn physics = - [&](const AMSHomogeneousGraph& g, AMSHomogeneousGraphFields& o) { - callback_count++; - const int64_t N = g.node_features.shape()[0]; - auto delta = makeTensor({N, 1}); - double* data = delta.data(); - for (int64_t i = 0; i < N; i++) data[i] = static_cast(i) + 200.0; - o.node_fields.insert("delta_u", std::move(delta)); - }; + HomogeneousGraphDomainFn physics = [&](const AMSHomogeneousGraph& g, + AMSHomogeneousGraphFields& o) { + callback_count++; + const int64_t N = g.node_features.shape()[0]; + auto delta = makeTensor({N, 1}); + double* data = delta.data(); + for (int64_t i = 0; i < N; i++) + data[i] = static_cast(i) + 200.0; + o.node_fields.insert("delta_u", std::move(delta)); + }; AMSExecute(executor, physics, graph, outputs); @@ -290,8 +299,7 @@ CATCH_TEST_CASE("store_data=false returns physics output with zero recorded case CATCH_TEST_CASE("Multiple calls accumulate cases with distinguishable values", "[wf][graph][storage]") { - fs::path test_dir = - fs::temp_directory_path() / "ams_graph_accumulation_test"; + fs::path test_dir = fs::temp_directory_path() / "ams_graph_accumulation_test"; fs::remove_all(test_dir); fs::create_directories(test_dir); @@ -320,8 +328,12 @@ CATCH_TEST_CASE("Multiple calls accumulate cases with distinguishable values", // Fill edge index with valid connectivity (3 edges for 4 nodes) int64_t* ei = edge_idx.data(); - ei[0] = 0; ei[1] = 1; ei[2] = 2; // sources - ei[3] = 1; ei[4] = 2; ei[5] = 3; // destinations + ei[0] = 0; + ei[1] = 1; + ei[2] = 2; // sources + ei[3] = 1; + ei[4] = 2; + ei[5] = 3; // destinations // Fill edge features float* ef = edge_feat.data(); @@ -395,8 +407,8 @@ CATCH_TEST_CASE("Multiple calls accumulate cases with distinguishable values", CATCH_TEST_CASE("Heterogeneous graph typed storage", "[.][wf][graph][storage][heterogeneous][future]") { - fs::path test_dir = - fs::temp_directory_path() / "ams_graph_heterogeneous_test"; + fs::path test_dir = fs::temp_directory_path() / "ams_graph_heterogeneous_" + "test"; fs::remove_all(test_dir); fs::create_directories(test_dir); @@ -414,14 +426,16 @@ CATCH_TEST_CASE("Heterogeneous graph typed storage", auto& fluid_store = graph.getOrCreateNodeStore("fluid"); auto fluid_feat = makeTensor({5, 3}); float* ff = fluid_feat.data(); - for (int i = 0; i < 15; i++) ff[i] = static_cast(i) * 0.1f; + for (int i = 0; i < 15; i++) + ff[i] = static_cast(i) * 0.1f; insertTensor(fluid_store, "features", std::move(fluid_feat)); // Node type "solid": 3 nodes, 2 features auto& solid_store = graph.getOrCreateNodeStore("solid"); auto solid_feat = makeTensor({3, 2}); float* sf = solid_feat.data(); - for (int i = 0; i < 6; i++) sf[i] = static_cast(i) * 0.2f; + for (int i = 0; i < 6; i++) + sf[i] = static_cast(i) * 0.2f; insertTensor(solid_store, "features", std::move(solid_feat)); // Edge type: fluid->solid (4 edges) @@ -429,8 +443,14 @@ CATCH_TEST_CASE("Heterogeneous graph typed storage", auto& edge_store = graph.getOrCreateEdgeStore(edge_type); auto edge_idx = makeTensor({2, 4}); int64_t* ei = edge_idx.data(); - ei[0] = 0; ei[1] = 1; ei[2] = 2; ei[3] = 3; // fluid nodes (sources) - ei[4] = 0; ei[5] = 1; ei[6] = 1; ei[7] = 2; // solid nodes (destinations) + ei[0] = 0; + ei[1] = 1; + ei[2] = 2; + ei[3] = 3; // fluid nodes (sources) + ei[4] = 0; + ei[5] = 1; + ei[6] = 1; + ei[7] = 2; // solid nodes (destinations) insertTensor(edge_store, "edge_index", std::move(edge_idx)); // Global features @@ -443,24 +463,26 @@ CATCH_TEST_CASE("Heterogeneous graph typed storage", AMSHeterogeneousGraphFields outputs; int callback_count = 0; - HeterogeneousGraphDomainFn physics = - [&](const AMSHeterogeneousGraph& g, AMSHeterogeneousGraphFields& o) { - callback_count++; - - // Output for fluid nodes - auto& fluid_out = o.getOrCreateNodeStore("fluid"); - auto fluid_delta = makeTensor({5, 1}); - double* fd = fluid_delta.data(); - for (int i = 0; i < 5; i++) fd[i] = static_cast(i) + 10.0; - fluid_out.insert("delta_u", std::move(fluid_delta)); - - // Output for solid nodes - auto& solid_out = o.getOrCreateNodeStore("solid"); - auto solid_delta = makeTensor({3, 1}); - double* sd = solid_delta.data(); - for (int i = 0; i < 3; i++) sd[i] = static_cast(i) + 20.0; - solid_out.insert("delta_u", std::move(solid_delta)); - }; + HeterogeneousGraphDomainFn physics = [&](const AMSHeterogeneousGraph& g, + AMSHeterogeneousGraphFields& o) { + callback_count++; + + // Output for fluid nodes + auto& fluid_out = o.getOrCreateNodeStore("fluid"); + auto fluid_delta = makeTensor({5, 1}); + double* fd = fluid_delta.data(); + for (int i = 0; i < 5; i++) + fd[i] = static_cast(i) + 10.0; + fluid_out.insert("delta_u", std::move(fluid_delta)); + + // Output for solid nodes + auto& solid_out = o.getOrCreateNodeStore("solid"); + auto solid_delta = makeTensor({3, 1}); + double* sd = solid_delta.data(); + for (int i = 0; i < 3; i++) + sd[i] = static_cast(i) + 20.0; + solid_out.insert("delta_u", std::move(solid_delta)); + }; AMSExecute(executor, physics, graph, outputs); @@ -490,8 +512,10 @@ CATCH_TEST_CASE("Heterogeneous graph typed storage", CATCH_REQUIRE(case0["tensors"].contains("target_node_solid__delta_u")); // Verify dtypes - CATCH_REQUIRE(case0["tensors"]["target_node_fluid__delta_u"]["dtype"] == "float64"); - CATCH_REQUIRE(case0["tensors"]["target_node_solid__delta_u"]["dtype"] == "float64"); + CATCH_REQUIRE(case0["tensors"]["target_node_fluid__delta_u"]["dtype"] == + "float64"); + CATCH_REQUIRE(case0["tensors"]["target_node_solid__delta_u"]["dtype"] == + "float64"); // Verify shapes auto fluid_shape = case0["tensors"]["target_node_fluid__delta_u"]["shape"]; @@ -532,14 +556,18 @@ CATCH_TEST_CASE("Surrogate success: zero callbacks and zero stored cases", // Fill all tensors float* nf = node_feat.data(); - for (int i = 0; i < 6; i++) nf[i] = static_cast(i); + for (int i = 0; i < 6; i++) + nf[i] = static_cast(i); int64_t* ei = edge_idx.data(); - ei[0] = 0; ei[1] = 1; // sources - ei[2] = 1; ei[3] = 2; // destinations + ei[0] = 0; + ei[1] = 1; // sources + ei[2] = 1; + ei[3] = 2; // destinations float* ef = edge_feat.data(); - ef[0] = 0.5f; ef[1] = 1.0f; + ef[0] = 0.5f; + ef[1] = 1.0f; float* gf = global_feat.data(); gf[0] = 1.0f; @@ -552,12 +580,12 @@ CATCH_TEST_CASE("Surrogate success: zero callbacks and zero stored cases", AMSHomogeneousGraphFields outputs; int callback_count = 0; - HomogeneousGraphDomainFn physics = - [&](const AMSHomogeneousGraph& g, AMSHomogeneousGraphFields& o) { - callback_count++; - auto delta = makeTensor({3, 1}); - o.node_fields.insert("delta_u", std::move(delta)); - }; + HomogeneousGraphDomainFn physics = [&](const AMSHomogeneousGraph& g, + AMSHomogeneousGraphFields& o) { + callback_count++; + auto delta = makeTensor({3, 1}); + o.node_fields.insert("delta_u", std::move(delta)); + }; AMSExecute(executor, physics, graph, outputs); @@ -596,14 +624,20 @@ CATCH_TEST_CASE("No database configured: physics output returned without crash", // Fill all tensors float* nf = node_feat.data(); - for (int i = 0; i < 8; i++) nf[i] = static_cast(i); + for (int i = 0; i < 8; i++) + nf[i] = static_cast(i); int64_t* ei = edge_idx.data(); - ei[0] = 0; ei[1] = 1; ei[2] = 2; // sources - ei[3] = 1; ei[4] = 2; ei[5] = 3; // destinations + ei[0] = 0; + ei[1] = 1; + ei[2] = 2; // sources + ei[3] = 1; + ei[4] = 2; + ei[5] = 3; // destinations float* ef = edge_feat.data(); - for (int i = 0; i < 3; i++) ef[i] = static_cast(i) * 0.3f; + for (int i = 0; i < 3; i++) + ef[i] = static_cast(i) * 0.3f; float* gf = global_feat.data(); gf[0] = 1.0f; @@ -616,15 +650,16 @@ CATCH_TEST_CASE("No database configured: physics output returned without crash", AMSHomogeneousGraphFields outputs; int callback_count = 0; - HomogeneousGraphDomainFn physics = - [&](const AMSHomogeneousGraph& g, AMSHomogeneousGraphFields& o) { - callback_count++; - const int64_t N = g.node_features.shape()[0]; - auto delta = makeTensor({N, 1}); - double* data = delta.data(); - for (int64_t i = 0; i < N; i++) data[i] = static_cast(i) * 2.5; - o.node_fields.insert("delta_u", std::move(delta)); - }; + HomogeneousGraphDomainFn physics = [&](const AMSHomogeneousGraph& g, + AMSHomogeneousGraphFields& o) { + callback_count++; + const int64_t N = g.node_features.shape()[0]; + auto delta = makeTensor({N, 1}); + double* data = delta.data(); + for (int64_t i = 0; i < N; i++) + data[i] = static_cast(i) * 2.5; + o.node_fields.insert("delta_u", std::move(delta)); + }; // Should not crash even without DB CATCH_REQUIRE_NOTHROW(AMSExecute(executor, physics, graph, outputs)); @@ -647,7 +682,8 @@ CATCH_TEST_CASE("No database configured: physics output returned without crash", } CATCH_TEST_CASE( - "Surrogate rejection: model runs but rejected, physics executes once, exact output stored", + "Surrogate rejection: model runs but rejected, physics executes once, " + "exact output stored", "[wf][graph][storage][surrogate]") { // This test validates the rejection path: when a surrogate model runs @@ -678,16 +714,26 @@ CATCH_TEST_CASE( // Fill with specific values to verify exact storage float* nf = node_feat.data(); - for (int i = 0; i < 12; i++) nf[i] = static_cast(i) * 1.5f; + for (int i = 0; i < 12; i++) + nf[i] = static_cast(i) * 1.5f; // Fill edge index (5 edges for 6 nodes) int64_t* ei = edge_idx.data(); - ei[0] = 0; ei[1] = 1; ei[2] = 2; ei[3] = 3; ei[4] = 4; // sources - ei[5] = 1; ei[6] = 2; ei[7] = 3; ei[8] = 4; ei[9] = 5; // destinations + ei[0] = 0; + ei[1] = 1; + ei[2] = 2; + ei[3] = 3; + ei[4] = 4; // sources + ei[5] = 1; + ei[6] = 2; + ei[7] = 3; + ei[8] = 4; + ei[9] = 5; // destinations // Fill edge features float* ef = edge_feat.data(); - for (int i = 0; i < 5; i++) ef[i] = static_cast(i) * 0.25f; + for (int i = 0; i < 5; i++) + ef[i] = static_cast(i) * 0.25f; // Fill global features float* gf = global_feat.data(); @@ -701,15 +747,16 @@ CATCH_TEST_CASE( AMSHomogeneousGraphFields outputs; int callback_count = 0; - HomogeneousGraphDomainFn physics = - [&](const AMSHomogeneousGraph& g, AMSHomogeneousGraphFields& o) { - callback_count++; - const int64_t N = g.node_features.shape()[0]; - auto delta = makeTensor({N, 1}); - double* data = delta.data(); - for (int64_t i = 0; i < N; i++) data[i] = static_cast(i) * 3.7; - o.node_fields.insert("delta_u", std::move(delta)); - }; + HomogeneousGraphDomainFn physics = [&](const AMSHomogeneousGraph& g, + AMSHomogeneousGraphFields& o) { + callback_count++; + const int64_t N = g.node_features.shape()[0]; + auto delta = makeTensor({N, 1}); + double* data = delta.data(); + for (int64_t i = 0; i < N; i++) + data[i] = static_cast(i) * 3.7; + o.node_fields.insert("delta_u", std::move(delta)); + }; AMSExecute(executor, physics, graph, outputs); From a583d2a3a836b2ad1acad79afcc760dc6395c2d5 Mon Sep 17 00:00:00 2001 From: Yohann Dudouit Date: Fri, 28 Aug 2026 17:46:27 -0700 Subject: [PATCH 04/18] Fix json db for updated graph format. --- src/AMSlib/wf/jsondb.cpp | 26 ++--- tests/AMSlib/ams_interface/CMakeLists.txt | 4 + .../test_graph_workflow_storage.cpp | 106 +++++++++++++++--- 3 files changed, 108 insertions(+), 28 deletions(-) diff --git a/src/AMSlib/wf/jsondb.cpp b/src/AMSlib/wf/jsondb.cpp index 71170571..3593e1f8 100644 --- a/src/AMSlib/wf/jsondb.cpp +++ b/src/AMSlib/wf/jsondb.cpp @@ -404,10 +404,9 @@ void JSONDB::store(const ams::AMSHomogeneousGraph& graph, edge_feature_dim = ef_shape[1]; } - if (graph.global_features.has_value() && - graph.global_features.value().elements() > 0) { - auto gf_shape = graph.global_features.value().shape(); - global_feature_dim = gf_shape[1]; + if (graph.global_features.elements() > 0) { + auto gf_shape = graph.global_features.shape(); + global_feature_dim = gf_shape[0]; } // Validate edge_index @@ -463,18 +462,15 @@ void JSONDB::store(const ams::AMSHomogeneousGraph& graph, } // Write global_features - if (global_feature_dim > 0 && graph.global_features.has_value()) { - if (json_mode_ == "binary") { - std::string rel_path = case_dir + "/global_features.bin"; - size_t byte_size = - writeBinaryTensor(graph.global_features.value(), rel_path); + if (global_feature_dim > 0 && json_mode_ == "binary") { + std::string rel_path = case_dir + "/global_features.bin"; + size_t byte_size = writeBinaryTensor(graph.global_features, rel_path); - tensors_json["global_features"] = { - {"path", rel_path}, - {"dtype", dtypeToString(graph.global_features.value().dType())}, - {"shape", std::vector{1, global_feature_dim}}, - {"byte_size", byte_size}}; - } + tensors_json["global_features"] = { + {"path", rel_path}, + {"dtype", dtypeToString(graph.global_features.dType())}, + {"shape", std::vector{global_feature_dim}}, + {"byte_size", byte_size}}; } // Write targets from outputs.node_fields diff --git a/tests/AMSlib/ams_interface/CMakeLists.txt b/tests/AMSlib/ams_interface/CMakeLists.txt index 10b5a25e..cde06056 100644 --- a/tests/AMSlib/ams_interface/CMakeLists.txt +++ b/tests/AMSlib/ams_interface/CMakeLists.txt @@ -44,6 +44,10 @@ ADD_AMS_UNIT_TEST(AMS_GRAPH_FALLBACK ams_graph_fallback) BUILD_UNIT_TEST(ams_graph_surrogate test_graph_surrogate.cpp Catch2::Catch2 ../ams_catch_main.cpp) ADD_AMS_UNIT_TEST(AMS_GRAPH_SURROGATE ams_graph_surrogate) +BUILD_UNIT_TEST(ams_graph_workflow_storage test_graph_workflow_storage.cpp Catch2::Catch2 ../ams_catch_main.cpp) +target_link_libraries(ams_graph_workflow_storage PRIVATE nlohmann_json::nlohmann_json) +ADD_AMS_UNIT_TEST(AMS_GRAPH_WORKFLOW_STORAGE ams_graph_workflow_storage) + BUILD_UNIT_TEST(ams_graph_mgn_surrogate test_graph_mgn_surrogate.cpp Catch2::Catch2 ../ams_catch_main.cpp) target_link_libraries(ams_graph_mgn_surrogate PRIVATE nlohmann_json::nlohmann_json) target_compile_definitions(ams_graph_mgn_surrogate diff --git a/tests/AMSlib/ams_interface/test_graph_workflow_storage.cpp b/tests/AMSlib/ams_interface/test_graph_workflow_storage.cpp index 96445286..60a9fdfd 100644 --- a/tests/AMSlib/ams_interface/test_graph_workflow_storage.cpp +++ b/tests/AMSlib/ams_interface/test_graph_workflow_storage.cpp @@ -6,9 +6,14 @@ */ #include +#include +#include #include #include #include +#include +#include +#include #include "AMS.h" #include "AMSGraph.hpp" @@ -83,8 +88,8 @@ CATCH_TEST_CASE( ef[e] = static_cast(e) * 0.01f; } - // Global features [1, 2] - float32 - auto global_feat = makeTensor({1, 2}); + // Global features [2] - float32 + auto global_feat = makeTensor({2}); float* gf = global_feat.data(); gf[0] = 0.01f; // dt gf[1] = 0.123f; // time_n @@ -156,6 +161,12 @@ CATCH_TEST_CASE( // CRITICAL: Verify global_features stored CATCH_REQUIRE(case0["tensors"].contains("global_features")); + CATCH_REQUIRE(case0["global_feature_dim"].get() == 2); + CATCH_REQUIRE(case0["tensors"]["global_features"]["shape"] == + nlohmann::json::array({2})); + CATCH_REQUIRE( + case0["tensors"]["global_features"]["byte_size"].get() == + 2 * sizeof(float)); // Verify output field stored // Note: Field name may be "delta_u" or "target_delta_u" - discover actual @@ -175,6 +186,18 @@ CATCH_TEST_CASE( case0["tensors"]["global_features"]["dtype"].get() == "float32"); + std::string global_path = + case0["tensors"]["global_features"]["path"].get(); + std::ifstream global_file(test_dir / global_path, std::ios::binary); + CATCH_REQUIRE(global_file.is_open()); + float stored_globals[2] = {}; + global_file.read(reinterpret_cast(stored_globals), + static_cast(sizeof(stored_globals))); + CATCH_REQUIRE(global_file.gcount() == + static_cast(sizeof(stored_globals))); + CATCH_REQUIRE(stored_globals[0] == 0.01f); + CATCH_REQUIRE(stored_globals[1] == 0.123f); + // Verify float64 delta_u (CRITICAL for MFEM precision) std::string delta_key = "delta_u"; if (case0["tensors"].contains("target_delta_u")) { @@ -200,6 +223,62 @@ CATCH_TEST_CASE( << std::endl; } +CATCH_TEST_CASE("Homogeneous graph without globals omits global storage", + "[wf][graph][storage]") +{ + fs::path test_dir = fs::temp_directory_path() / "ams_graph_empty_globals_" + "test"; + fs::remove_all(test_dir); + fs::create_directories(test_dir); + + AMSInit(); + AMSConfigureFSDatabase(AMSDBType::AMS_JSON, test_dir.string().c_str()); + + AMSCAbstrModel recorder = + AMSRegisterAbstractModel("empty_globals_domain", -1.0, "", true); + AMSExecutor executor = AMSCreateExecutor(recorder, 0, 1); + + auto node_feat = makeTensor({3, 2}); + auto edge_idx = makeTensor({2, 2}); + auto edge_feat = makeTensor({2, 1}); + + int64_t* edge_data = edge_idx.data(); + edge_data[0] = 0; + edge_data[1] = 1; + edge_data[2] = 1; + edge_data[3] = 2; + + AMSHomogeneousGraph graph(std::move(node_feat), + std::move(edge_idx), + std::move(edge_feat)); + CATCH_REQUIRE(graph.global_features.shape().size() == 1); + CATCH_REQUIRE(graph.global_features.shape()[0] == 0); + + AMSHomogeneousGraphFields outputs; + HomogeneousGraphDomainFn physics = [](const AMSHomogeneousGraph& g, + AMSHomogeneousGraphFields& o) { + auto delta = makeTensor({g.node_features.shape()[0], 1}); + o.node_fields.insert("delta_u", std::move(delta)); + }; + + AMSExecute(executor, physics, graph, outputs); + AMSDestroyExecutor(executor); + + std::ifstream manifest_file(test_dir / "manifest.json"); + CATCH_REQUIRE(manifest_file.is_open()); + nlohmann::json manifest; + manifest_file >> manifest; + + CATCH_REQUIRE(manifest["cases"].size() == 1); + const auto& stored_case = manifest["cases"][0]; + CATCH_REQUIRE(stored_case["global_feature_dim"].get() == 0); + CATCH_REQUIRE_FALSE(stored_case["tensors"].contains("global_features")); + CATCH_REQUIRE_FALSE( + fs::exists(test_dir / "step_000000" / "global_features.bin")); + + fs::remove_all(test_dir); +} + // ============================================================================ // A2 Tests: Complete Storage Test Coverage // ============================================================================ @@ -224,7 +303,7 @@ CATCH_TEST_CASE( auto node_feat = makeTensor({5, 2}); auto edge_idx = makeTensor({2, 4}); auto edge_feat = makeTensor({4, 1}); - auto global_feat = makeTensor({1, 2}); + auto global_feat = makeTensor({2}); // Fill with identifiable values float* nf = node_feat.data(); @@ -318,7 +397,7 @@ CATCH_TEST_CASE("Multiple calls accumulate cases with distinguishable values", auto node_feat = makeTensor({4, 2}); auto edge_idx = makeTensor({2, 3}); auto edge_feat = makeTensor({3, 1}); - auto global_feat = makeTensor({1, 1}); + auto global_feat = makeTensor({1}); // Fill node features with call-specific values float* nf = node_feat.data(); @@ -407,8 +486,9 @@ CATCH_TEST_CASE("Multiple calls accumulate cases with distinguishable values", CATCH_TEST_CASE("Heterogeneous graph typed storage", "[.][wf][graph][storage][heterogeneous][future]") { - fs::path test_dir = fs::temp_directory_path() / "ams_graph_heterogeneous_" - "test"; + fs::path test_dir = fs::temp_directory_path() / + "ams_graph_heterogeneous_" + "test"; fs::remove_all(test_dir); fs::create_directories(test_dir); @@ -473,7 +553,7 @@ CATCH_TEST_CASE("Heterogeneous graph typed storage", double* fd = fluid_delta.data(); for (int i = 0; i < 5; i++) fd[i] = static_cast(i) + 10.0; - fluid_out.insert("delta_u", std::move(fluid_delta)); + insertTensor(fluid_out, "delta_u", std::move(fluid_delta)); // Output for solid nodes auto& solid_out = o.getOrCreateNodeStore("solid"); @@ -481,14 +561,14 @@ CATCH_TEST_CASE("Heterogeneous graph typed storage", double* sd = solid_delta.data(); for (int i = 0; i < 3; i++) sd[i] = static_cast(i) + 20.0; - solid_out.insert("delta_u", std::move(solid_delta)); + insertTensor(solid_out, "delta_u", std::move(solid_delta)); }; AMSExecute(executor, physics, graph, outputs); CATCH_REQUIRE(callback_count == 1); - CATCH_REQUIRE(outputs.node_stores.find("fluid") != nullptr); - CATCH_REQUIRE(outputs.node_stores.find("solid") != nullptr); + CATCH_REQUIRE(outputs.node_stores.find("fluid") != outputs.node_stores.end()); + CATCH_REQUIRE(outputs.node_stores.find("solid") != outputs.node_stores.end()); // Note: Not destroying executor to avoid triggering AMSFinalize between tests // The executor will be cleaned up at program exit @@ -552,7 +632,7 @@ CATCH_TEST_CASE("Surrogate success: zero callbacks and zero stored cases", auto node_feat = makeTensor({3, 2}); auto edge_idx = makeTensor({2, 2}); auto edge_feat = makeTensor({2, 1}); - auto global_feat = makeTensor({1, 1}); + auto global_feat = makeTensor({1}); // Fill all tensors float* nf = node_feat.data(); @@ -620,7 +700,7 @@ CATCH_TEST_CASE("No database configured: physics output returned without crash", auto node_feat = makeTensor({4, 2}); auto edge_idx = makeTensor({2, 3}); auto edge_feat = makeTensor({3, 1}); - auto global_feat = makeTensor({1, 1}); + auto global_feat = makeTensor({1}); // Fill all tensors float* nf = node_feat.data(); @@ -710,7 +790,7 @@ CATCH_TEST_CASE( auto node_feat = makeTensor({6, 2}); auto edge_idx = makeTensor({2, 5}); auto edge_feat = makeTensor({5, 1}); - auto global_feat = makeTensor({1, 1}); + auto global_feat = makeTensor({1}); // Fill with specific values to verify exact storage float* nf = node_feat.data(); From bab0881003799bcdc42206900b7a37e495632147 Mon Sep 17 00:00:00 2001 From: Yohann Dudouit Date: Fri, 28 Aug 2026 17:59:31 -0700 Subject: [PATCH 05/18] clang format. --- .../test_graph_workflow_storage.cpp | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/tests/AMSlib/ams_interface/test_graph_workflow_storage.cpp b/tests/AMSlib/ams_interface/test_graph_workflow_storage.cpp index 60a9fdfd..098ffc33 100644 --- a/tests/AMSlib/ams_interface/test_graph_workflow_storage.cpp +++ b/tests/AMSlib/ams_interface/test_graph_workflow_storage.cpp @@ -226,8 +226,9 @@ CATCH_TEST_CASE( CATCH_TEST_CASE("Homogeneous graph without globals omits global storage", "[wf][graph][storage]") { - fs::path test_dir = fs::temp_directory_path() / "ams_graph_empty_globals_" - "test"; + fs::path test_dir = fs::temp_directory_path() / + "ams_graph_empty_globals_" + "test"; fs::remove_all(test_dir); fs::create_directories(test_dir); @@ -242,12 +243,21 @@ CATCH_TEST_CASE("Homogeneous graph without globals omits global storage", auto edge_idx = makeTensor({2, 2}); auto edge_feat = makeTensor({2, 1}); + float* node_data = node_feat.data(); + for (int i = 0; i < 6; ++i) { + node_data[i] = static_cast(i); + } + int64_t* edge_data = edge_idx.data(); edge_data[0] = 0; edge_data[1] = 1; edge_data[2] = 1; edge_data[3] = 2; + float* edge_feature_data = edge_feat.data(); + edge_feature_data[0] = 0.5f; + edge_feature_data[1] = 1.0f; + AMSHomogeneousGraph graph(std::move(node_feat), std::move(edge_idx), std::move(edge_feat)); @@ -258,6 +268,10 @@ CATCH_TEST_CASE("Homogeneous graph without globals omits global storage", HomogeneousGraphDomainFn physics = [](const AMSHomogeneousGraph& g, AMSHomogeneousGraphFields& o) { auto delta = makeTensor({g.node_features.shape()[0], 1}); + double* delta_data = delta.data(); + for (int64_t i = 0; i < g.node_features.shape()[0]; ++i) { + delta_data[i] = static_cast(i); + } o.node_fields.insert("delta_u", std::move(delta)); }; From d911e34dc7472a2b75671d00214e4e120724b045 Mon Sep 17 00:00:00 2001 From: Yohann Dudouit Date: Mon, 31 Aug 2026 10:28:21 -0700 Subject: [PATCH 06/18] Fix for clang tidy. --- src/AMSlib/wf/basedb.hpp | 13 ++++++++----- src/AMSlib/wf/jsondb.cpp | 4 ++-- .../ams_interface/test_graph_workflow_storage.cpp | 8 ++++---- 3 files changed, 14 insertions(+), 11 deletions(-) diff --git a/src/AMSlib/wf/basedb.hpp b/src/AMSlib/wf/basedb.hpp index 9bc64f59..1cb8aac0 100644 --- a/src/AMSlib/wf/basedb.hpp +++ b/src/AMSlib/wf/basedb.hpp @@ -139,8 +139,9 @@ class BaseDB * @param[in] graph The homogeneous graph containing input features * @param[in] outputs The graph fields containing output/target data */ - virtual void store(const ams::AMSHomogeneousGraph& graph, - const ams::AMSHomogeneousGraphFields& outputs) + 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()); @@ -152,8 +153,9 @@ class BaseDB * @param[in] graph The heterogeneous graph containing input features * @param[in] outputs The graph fields containing output/target data */ - virtual void store(const ams::AMSHeterogeneousGraph& graph, - const ams::AMSHeterogeneousGraphFields& outputs) + 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 " @@ -1595,8 +1597,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() diff --git a/src/AMSlib/wf/jsondb.cpp b/src/AMSlib/wf/jsondb.cpp index 3593e1f8..3e3ee9e3 100644 --- a/src/AMSlib/wf/jsondb.cpp +++ b/src/AMSlib/wf/jsondb.cpp @@ -511,8 +511,8 @@ void JSONDB::store(const ams::AMSHomogeneousGraph& graph, num_edges); } -void JSONDB::store(const ams::AMSHeterogeneousGraph& graph, - const ams::AMSHeterogeneousGraphFields& outputs) +void JSONDB::store(const ams::AMSHeterogeneousGraph&, + const ams::AMSHeterogeneousGraphFields&) { // Heterogeneous graph storage not yet implemented THROW(std::runtime_error, diff --git a/tests/AMSlib/ams_interface/test_graph_workflow_storage.cpp b/tests/AMSlib/ams_interface/test_graph_workflow_storage.cpp index 098ffc33..dc387929 100644 --- a/tests/AMSlib/ams_interface/test_graph_workflow_storage.cpp +++ b/tests/AMSlib/ams_interface/test_graph_workflow_storage.cpp @@ -557,7 +557,7 @@ CATCH_TEST_CASE("Heterogeneous graph typed storage", AMSHeterogeneousGraphFields outputs; int callback_count = 0; - HeterogeneousGraphDomainFn physics = [&](const AMSHeterogeneousGraph& g, + HeterogeneousGraphDomainFn physics = [&](const AMSHeterogeneousGraph&, AMSHeterogeneousGraphFields& o) { callback_count++; @@ -567,7 +567,7 @@ CATCH_TEST_CASE("Heterogeneous graph typed storage", double* fd = fluid_delta.data(); for (int i = 0; i < 5; i++) fd[i] = static_cast(i) + 10.0; - insertTensor(fluid_out, "delta_u", std::move(fluid_delta)); + fluid_out.insert("delta_u", std::move(fluid_delta)); // Output for solid nodes auto& solid_out = o.getOrCreateNodeStore("solid"); @@ -575,7 +575,7 @@ CATCH_TEST_CASE("Heterogeneous graph typed storage", double* sd = solid_delta.data(); for (int i = 0; i < 3; i++) sd[i] = static_cast(i) + 20.0; - insertTensor(solid_out, "delta_u", std::move(solid_delta)); + solid_out.insert("delta_u", std::move(solid_delta)); }; AMSExecute(executor, physics, graph, outputs); @@ -677,7 +677,7 @@ CATCH_TEST_CASE("Surrogate success: zero callbacks and zero stored cases", HomogeneousGraphDomainFn physics = [&](const AMSHomogeneousGraph& g, AMSHomogeneousGraphFields& o) { callback_count++; - auto delta = makeTensor({3, 1}); + auto delta = makeTensor({g.node_features.shape()[0], 1}); o.node_fields.insert("delta_u", std::move(delta)); }; From 713f1c2c674ecbe45859c386e2de8c2229255cdc Mon Sep 17 00:00:00 2001 From: Yohann Dudouit Date: Tue, 1 Sep 2026 14:50:07 -0700 Subject: [PATCH 07/18] Remove unnecessary if(DB). --- src/AMSlib/wf/workflow.hpp | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/src/AMSlib/wf/workflow.hpp b/src/AMSlib/wf/workflow.hpp index b5e5ae5a..730f459f 100644 --- a/src/AMSlib/wf/workflow.hpp +++ b/src/AMSlib/wf/workflow.hpp @@ -482,9 +482,7 @@ class AMSWorkflow CALIPER(CALI_MARK_END("PHYSICS MODULE");) // Store data after physics computation - if (DB) { - storeGraphData(graph_input, outputs); - } + storeGraphData(graph_input, outputs); CALIPER(CALI_MARK_END("AMSEvaluateGraph");) } @@ -508,9 +506,7 @@ class AMSWorkflow CALIPER(CALI_MARK_END("PHYSICS MODULE");) // Store data after physics computation - if (DB) { - storeGraphData(graph_input, outputs); - } + storeGraphData(graph_input, outputs); CALIPER(CALI_MARK_END("AMSEvaluateGraph");) } From fda4614f6a554e75eb32e3bed2c615420bdad3ef Mon Sep 17 00:00:00 2001 From: Yohann Dudouit Date: Tue, 1 Sep 2026 14:55:43 -0700 Subject: [PATCH 08/18] Support pure Json export. --- src/AMSlib/wf/jsondb.cpp | 88 ++++++++++++++++++++++------------------ src/AMSlib/wf/jsondb.hpp | 32 ++++++++------- 2 files changed, 66 insertions(+), 54 deletions(-) diff --git a/src/AMSlib/wf/jsondb.cpp b/src/AMSlib/wf/jsondb.cpp index 3e3ee9e3..bdcc53d6 100644 --- a/src/AMSlib/wf/jsondb.cpp +++ b/src/AMSlib/wf/jsondb.cpp @@ -114,10 +114,10 @@ JSONDB::~JSONDB() { if (!finalized_) { try { - finalize(); + close(); } catch (const std::exception& e) { AMS_WARNING(JSONDB, - "Exception during automatic finalization: {}", + "Exception while automatically closing JSONDB: {}", e.what()); } } @@ -148,22 +148,6 @@ std::string JSONDB::torchDTypeToString(torch::Dtype dtype) const return "unknown"; } -size_t JSONDB::dtypeSize(AMSDType dtype) const -{ - switch (dtype) { - case AMS_SINGLE: - return 4; - case AMS_DOUBLE: - return 8; - case AMS_INT32: - return 4; - case AMS_INT64: - return 8; - default: - return 0; - } -} - size_t JSONDB::writeBinaryTensor(const AMSTensor& tensor, const std::string& path) { @@ -172,13 +156,9 @@ size_t JSONDB::writeBinaryTensor(const AMSTensor& tensor, size_t byte_size = tensor.elements() * tensor.element_size(); AMSResourceType location = tensor.location(); - // Handle GPU tensors - copy to CPU first - std::vector cpu_buffer; + // AMSTensor device transfers are not implemented here yet. if (location != AMSResourceType::AMS_HOST) { - AMS_WARNING(JSONDB, - "GPU tensor detected. Copying to CPU for serialization (not " - "yet " - "implemented - will fail)."); + AMS_WARNING(JSONDB, "GPU tensor serialization is not implemented"); THROW(std::runtime_error, "GPU tensor serialization not yet implemented"); // TODO: Implement cudaMemcpy/hipMemcpy here } @@ -253,6 +233,7 @@ nlohmann::json JSONDB::encodeBase64Tensor(const AMSTensor& tensor) result["encoding"] = "base64"; result["data"] = base64Encode(data, byte_size); result["dtype"] = dtypeToString(tensor.dType()); + result["byte_size"] = byte_size; // Add shape auto shape_ref = tensor.shape(); @@ -261,6 +242,22 @@ nlohmann::json JSONDB::encodeBase64Tensor(const AMSTensor& tensor) return result; } +nlohmann::json JSONDB::encodeBase64Tensor(const torch::Tensor& tensor) +{ + torch::Tensor cpu_tensor = tensor.contiguous().cpu(); + const uint8_t* data = static_cast(cpu_tensor.data_ptr()); + size_t byte_size = cpu_tensor.nbytes(); + + auto sizes = cpu_tensor.sizes(); + std::vector shape(sizes.begin(), sizes.end()); + + return nlohmann::json{{"encoding", "base64"}, + {"data", base64Encode(data, byte_size)}, + {"dtype", torchDTypeToString(cpu_tensor.scalar_type())}, + {"shape", shape}, + {"byte_size", byte_size}}; +} + void JSONDB::validateEdgeIndex(const AMSTensor& edge_index, int64_t num_nodes) { auto shape_ref = edge_index.shape(); @@ -340,10 +337,8 @@ void JSONDB::store(ArrayRef Inputs, {"dtype", torchDTypeToString(Inputs[i].scalar_type())}, {"shape", shape}, {"byte_size", byte_size}}; - } else { - // Pure JSON mode - not yet fully implemented for torch tensors - THROW(std::runtime_error, - "Pure JSON mode not yet implemented for tensor storage"); + } else { // Pure json mode + tensors_json[name] = encodeBase64Tensor(Inputs[i]); } } @@ -366,6 +361,8 @@ void JSONDB::store(ArrayRef Inputs, torchDTypeToString(Outputs[i].scalar_type())}, {"shape", shape}, {"byte_size", byte_size}}; + } else { // Pure json mode + tensors_json[name] = encodeBase64Tensor(Outputs[i]); } } @@ -434,6 +431,8 @@ void JSONDB::store(const ams::AMSHomogeneousGraph& graph, {"dtype", dtypeToString(graph.node_features.dType())}, {"shape", std::vector{num_nodes, node_feature_dim}}, {"byte_size", byte_size}}; + } else { // Pure json mode + tensors_json["node_features"] = encodeBase64Tensor(graph.node_features); } // Write edge_index @@ -445,6 +444,8 @@ void JSONDB::store(const ams::AMSHomogeneousGraph& graph, {"dtype", "int64"}, {"shape", std::vector{2, num_edges}}, {"byte_size", byte_size}}; + } else { // Pure json mode + tensors_json["edge_index"] = encodeBase64Tensor(graph.edge_index); } // Write edge_features @@ -458,19 +459,26 @@ void JSONDB::store(const ams::AMSHomogeneousGraph& graph, {"dtype", dtypeToString(graph.edge_features.dType())}, {"shape", std::vector{num_edges, edge_feature_dim}}, {"byte_size", byte_size}}; + } else { // Pure json mode + tensors_json["edge_features"] = encodeBase64Tensor(graph.edge_features); } } // Write global_features - if (global_feature_dim > 0 && json_mode_ == "binary") { - std::string rel_path = case_dir + "/global_features.bin"; - size_t byte_size = writeBinaryTensor(graph.global_features, rel_path); + if (global_feature_dim > 0) { + if (json_mode_ == "binary") { + std::string rel_path = case_dir + "/global_features.bin"; + size_t byte_size = writeBinaryTensor(graph.global_features, rel_path); - tensors_json["global_features"] = { - {"path", rel_path}, - {"dtype", dtypeToString(graph.global_features.dType())}, - {"shape", std::vector{global_feature_dim}}, - {"byte_size", byte_size}}; + tensors_json["global_features"] = { + {"path", rel_path}, + {"dtype", dtypeToString(graph.global_features.dType())}, + {"shape", std::vector{global_feature_dim}}, + {"byte_size", byte_size}}; + } else { // Pure json mode + tensors_json["global_features"] = + encodeBase64Tensor(graph.global_features); + } } // Write targets from outputs.node_fields @@ -493,6 +501,8 @@ void JSONDB::store(const ams::AMSHomogeneousGraph& graph, {"dtype", dtypeToString(delta_u->dType())}, {"shape", std::vector{num_nodes, target_dim}}, {"byte_size", byte_size}}; + } else { // Pure json mode + tensors_json["target_delta_u"] = encodeBase64Tensor(*delta_u); } } @@ -519,7 +529,7 @@ void JSONDB::store(const ams::AMSHeterogeneousGraph&, "Heterogeneous graph storage not yet implemented in JSONDB"); } -void JSONDB::finalize() +void JSONDB::close() { if (finalized_) { AMS_DBG(JSONDB, "Manifest already finalized, skipping"); @@ -544,8 +554,8 @@ void JSONDB::finalize() // Add all cases manifest["cases"] = cases_; - // Write manifest.json - fs::path manifest_path = fs::path(fp) / "manifest.json"; + // FileDB provides a domain- and rank-specific manifest filename. + fs::path manifest_path = fn; std::ofstream manifest_file(manifest_path.string()); if (!manifest_file.is_open()) { THROW(std::runtime_error, diff --git a/src/AMSlib/wf/jsondb.hpp b/src/AMSlib/wf/jsondb.hpp index 7505792f..713317f3 100644 --- a/src/AMSlib/wf/jsondb.hpp +++ b/src/AMSlib/wf/jsondb.hpp @@ -28,16 +28,18 @@ namespace db /** * @brief JSON database backend for storing both tensor and graph data. * - * This class provides a JSON-based storage backend that can serialize: - * - Flat tensors (via store(ArrayRef, ArrayRef)) - * - Homogeneous graphs (via store(AMSHomogeneousGraph, AMSHomogeneousGraphFields)) - * - Heterogeneous graphs (via store(AMSHeterogeneousGraph, AMSHeterogeneousGraphFields)) + * This class provides a JSON-based storage backend that can serialize flat + * tensors and homogeneous graphs. Heterogeneous graph storage is not yet + * supported. * * Supports two modes: * - "binary": Binary tensor files + JSON manifest (default, efficient) - * - "json": Pure JSON with base64-encoded binary data (human-readable) + * - "json": Self-contained JSON with base64-encoded tensor data * * Output format is compatible with PyTorch Geometric data loaders. + * + * @note Case and step directories are not rank-scoped. Concurrent MPI ranks + * writing to the same output directory are not currently supported. */ class JSONDB final : public FileDB { @@ -79,6 +81,13 @@ class JSONDB final : public FileDB */ nlohmann::json encodeBase64Tensor(const AMSTensor& tensor); + /** + * @brief Encode a PyTorch tensor as base64 JSON (for pure JSON mode) + * @param[in] tensor The PyTorch tensor to encode + * @return JSON object with encoded data + */ + nlohmann::json encodeBase64Tensor(const torch::Tensor& tensor); + /** * @brief Validate edge_index tensor format * @param[in] edge_index Edge connectivity tensor [2, E] @@ -100,13 +109,6 @@ class JSONDB final : public FileDB */ std::string torchDTypeToString(torch::Dtype dtype) const; - /** - * @brief Get byte size for a data type - * @param[in] dtype The data type enum - * @return Size in bytes - */ - size_t dtypeSize(AMSDType dtype) const; - public: /** * @brief Construct a JSON database @@ -123,7 +125,7 @@ class JSONDB final : public FileDB /** * @brief Destructor - finalizes manifest if not already done */ - ~JSONDB(); + ~JSONDB() override; // Delete copy/move constructors JSONDB(const JSONDB&) = delete; @@ -154,9 +156,9 @@ class JSONDB final : public FileDB const ams::AMSHeterogeneousGraphFields& outputs) override; /** - * @brief Finalize and write manifest.json + * @brief Finalize and write the JSON manifest */ - void finalize(); + void close() override; /** * @brief Set application metadata From e2e01d0d4f3d783428d8efcf475ae598d753b247 Mon Sep 17 00:00:00 2001 From: Yohann Dudouit Date: Tue, 1 Sep 2026 14:58:57 -0700 Subject: [PATCH 09/18] Remove callApplication for graphs. --- src/AMSlib/wf/interface.cpp | 20 -------------------- src/AMSlib/wf/interface.hpp | 9 --------- 2 files changed, 29 deletions(-) diff --git a/src/AMSlib/wf/interface.cpp b/src/AMSlib/wf/interface.cpp index 2cb7fa61..cc575362 100644 --- a/src/AMSlib/wf/interface.cpp +++ b/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) // ============================================================================ diff --git a/src/AMSlib/wf/interface.hpp b/src/AMSlib/wf/interface.hpp index 463c121e..b8233d0d 100644 --- a/src/AMSlib/wf/interface.hpp +++ b/src/AMSlib/wf/interface.hpp @@ -14,15 +14,6 @@ void callApplication(ams::DomainLambda CallBack, ams::MutableArrayRef InOuts, ams::MutableArrayRef Outs); -void callApplication(ams::HomogeneousGraphDomainFn CallBack, - const ams::AMSHomogeneousGraph& graph, - ams::AMSHomogeneousGraphFields& outputs); - -void callApplication(ams::HeterogeneousGraphDomainFn CallBack, - const ams::AMSHeterogeneousGraph& graph, - ams::AMSHeterogeneousGraphFields& outputs); - - void callAMS(ams::AMSWorkflow* executor, ams::DomainLambda Physics, const ams::SmallVector& ins, From c2a257cc036f0fbe8666db25f35491ca0c19abc9 Mon Sep 17 00:00:00 2001 From: Yohann Dudouit Date: Tue, 1 Sep 2026 15:12:02 -0700 Subject: [PATCH 10/18] Remove exposed `fs` namespace. --- src/AMSlib/AMS.cpp | 3 +++ src/AMSlib/wf/basedb.hpp | 22 ++++++++++------------ src/AMSlib/wf/hdf5db.cpp | 6 ++++++ src/AMSlib/wf/jsondb.cpp | 17 ++++++++++------- src/AMSlib/wf/jsondb.hpp | 3 --- 5 files changed, 29 insertions(+), 22 deletions(-) diff --git a/src/AMSlib/AMS.cpp b/src/AMSlib/AMS.cpp index 9c480815..ed7e7e17 100644 --- a/src/AMSlib/AMS.cpp +++ b/src/AMSlib/AMS.cpp @@ -12,6 +12,7 @@ #endif #include +#include #include #include #include @@ -34,6 +35,8 @@ using namespace ams; namespace { +namespace fs = std::experimental::filesystem; + struct AMSAbstractModel { public: std::string SPath; diff --git a/src/AMSlib/wf/basedb.hpp b/src/AMSlib/wf/basedb.hpp index 1cb8aac0..5ac32e6f 100644 --- a/src/AMSlib/wf/basedb.hpp +++ b/src/AMSlib/wf/basedb.hpp @@ -31,8 +31,6 @@ #include "wf/resource_manager.hpp" #include "wf/utils.hpp" -namespace fs = std::experimental::filesystem; - // Forward declarations for graph types namespace ams { @@ -213,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) } @@ -1707,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); @@ -1916,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.", diff --git a/src/AMSlib/wf/hdf5db.cpp b/src/AMSlib/wf/hdf5db.cpp index 1cb064c9..d76c16f3 100644 --- a/src/AMSlib/wf/hdf5db.cpp +++ b/src/AMSlib/wf/hdf5db.cpp @@ -13,6 +13,7 @@ #include #include +#include #include #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 shape) { std::ostringstream oss; diff --git a/src/AMSlib/wf/jsondb.cpp b/src/AMSlib/wf/jsondb.cpp index bdcc53d6..cb4f0b89 100644 --- a/src/AMSlib/wf/jsondb.cpp +++ b/src/AMSlib/wf/jsondb.cpp @@ -10,6 +10,7 @@ #include #include +#include #include #include #include @@ -26,6 +27,8 @@ using namespace ams; namespace { +namespace fs = std::experimental::filesystem; + // Check system endianness bool isLittleEndian() { @@ -337,7 +340,7 @@ void JSONDB::store(ArrayRef Inputs, {"dtype", torchDTypeToString(Inputs[i].scalar_type())}, {"shape", shape}, {"byte_size", byte_size}}; - } else { // Pure json mode + } else { // Pure json mode tensors_json[name] = encodeBase64Tensor(Inputs[i]); } } @@ -361,7 +364,7 @@ void JSONDB::store(ArrayRef Inputs, torchDTypeToString(Outputs[i].scalar_type())}, {"shape", shape}, {"byte_size", byte_size}}; - } else { // Pure json mode + } else { // Pure json mode tensors_json[name] = encodeBase64Tensor(Outputs[i]); } } @@ -431,7 +434,7 @@ void JSONDB::store(const ams::AMSHomogeneousGraph& graph, {"dtype", dtypeToString(graph.node_features.dType())}, {"shape", std::vector{num_nodes, node_feature_dim}}, {"byte_size", byte_size}}; - } else { // Pure json mode + } else { // Pure json mode tensors_json["node_features"] = encodeBase64Tensor(graph.node_features); } @@ -444,7 +447,7 @@ void JSONDB::store(const ams::AMSHomogeneousGraph& graph, {"dtype", "int64"}, {"shape", std::vector{2, num_edges}}, {"byte_size", byte_size}}; - } else { // Pure json mode + } else { // Pure json mode tensors_json["edge_index"] = encodeBase64Tensor(graph.edge_index); } @@ -459,7 +462,7 @@ void JSONDB::store(const ams::AMSHomogeneousGraph& graph, {"dtype", dtypeToString(graph.edge_features.dType())}, {"shape", std::vector{num_edges, edge_feature_dim}}, {"byte_size", byte_size}}; - } else { // Pure json mode + } else { // Pure json mode tensors_json["edge_features"] = encodeBase64Tensor(graph.edge_features); } } @@ -475,7 +478,7 @@ void JSONDB::store(const ams::AMSHomogeneousGraph& graph, {"dtype", dtypeToString(graph.global_features.dType())}, {"shape", std::vector{global_feature_dim}}, {"byte_size", byte_size}}; - } else { // Pure json mode + } else { // Pure json mode tensors_json["global_features"] = encodeBase64Tensor(graph.global_features); } @@ -501,7 +504,7 @@ void JSONDB::store(const ams::AMSHomogeneousGraph& graph, {"dtype", dtypeToString(delta_u->dType())}, {"shape", std::vector{num_nodes, target_dim}}, {"byte_size", byte_size}}; - } else { // Pure json mode + } else { // Pure json mode tensors_json["target_delta_u"] = encodeBase64Tensor(*delta_u); } } diff --git a/src/AMSlib/wf/jsondb.hpp b/src/AMSlib/wf/jsondb.hpp index 713317f3..0a0c179e 100644 --- a/src/AMSlib/wf/jsondb.hpp +++ b/src/AMSlib/wf/jsondb.hpp @@ -8,7 +8,6 @@ #ifndef __AMS_JSON_DB__ #define __AMS_JSON_DB__ -#include #include #include #include @@ -18,8 +17,6 @@ #include "AMSTensor.hpp" #include "wf/basedb.hpp" -namespace fs = std::experimental::filesystem; - namespace ams { namespace db From f12edb4242d9aee496b6f483822634381065fdfe Mon Sep 17 00:00:00 2001 From: Yohann Dudouit Date: Tue, 1 Sep 2026 15:16:12 -0700 Subject: [PATCH 11/18] Test inline pure json mode, and use AMS defined fn name for DB. --- .../test_graph_workflow_storage.cpp | 269 +++++++++++++++++- 1 file changed, 258 insertions(+), 11 deletions(-) diff --git a/tests/AMSlib/ams_interface/test_graph_workflow_storage.cpp b/tests/AMSlib/ams_interface/test_graph_workflow_storage.cpp index dc387929..d745e79c 100644 --- a/tests/AMSlib/ams_interface/test_graph_workflow_storage.cpp +++ b/tests/AMSlib/ams_interface/test_graph_workflow_storage.cpp @@ -5,9 +5,13 @@ * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception */ +#include + #include #include #include +#include +#include #include #include #include @@ -19,6 +23,7 @@ #include "AMSGraph.hpp" #include "AMSTensor.hpp" #include "nlohmann/json.hpp" +#include "wf/jsondb.hpp" using namespace ams; namespace fs = std::filesystem; @@ -43,6 +48,94 @@ static AMSTensor makeTensor(std::vector shape) return AMSTensor::create(shape, strides, AMSResourceType::AMS_HOST); } +static std::vector decodeBase64(const std::string& encoded) +{ + std::vector decoded; + uint32_t accumulator = 0; + int bits = 0; + + for (unsigned char c : encoded) { + if (c == '=') break; + + int value = -1; + if (c >= 'A' && c <= 'Z') + value = c - 'A'; + else if (c >= 'a' && c <= 'z') + value = c - 'a' + 26; + else if (c >= '0' && c <= '9') + value = c - '0' + 52; + else if (c == '+') + value = 62; + else if (c == '/') + value = 63; + + CATCH_REQUIRE(value >= 0); + accumulator = (accumulator << 6) | static_cast(value); + bits += 6; + if (bits >= 8) { + bits -= 8; + decoded.push_back(static_cast((accumulator >> bits) & 0xffU)); + } + } + + return decoded; +} + +static void requireInlineTensor(const nlohmann::json& tensor, + const void* expected_data, + size_t expected_byte_size, + const std::string& expected_dtype, + const nlohmann::json& expected_shape) +{ + CATCH_REQUIRE(tensor["encoding"] == "base64"); + CATCH_REQUIRE(tensor["dtype"] == expected_dtype); + CATCH_REQUIRE(tensor["shape"] == expected_shape); + CATCH_REQUIRE(tensor["byte_size"] == expected_byte_size); + CATCH_REQUIRE_FALSE(tensor.contains("path")); + + auto decoded = decodeBase64(tensor["data"].get()); + CATCH_REQUIRE(decoded.size() == expected_byte_size); + CATCH_REQUIRE( + std::memcmp(decoded.data(), expected_data, expected_byte_size) == 0); +} + +static bool containsBinaryFile(const fs::path& directory) +{ + for (const auto& entry : fs::recursive_directory_iterator(directory)) { + if (entry.is_regular_file() && entry.path().extension() == ".bin") { + return true; + } + } + return false; +} + +class ScopedEnvironmentVariable +{ + std::string name_; + std::string previous_value_; + bool had_previous_value_; + +public: + ScopedEnvironmentVariable(std::string name, const std::string& value) + : name_(std::move(name)), had_previous_value_(false) + { + const char* previous_value = std::getenv(name_.c_str()); + if (previous_value != nullptr) { + previous_value_ = previous_value; + had_previous_value_ = true; + } + setenv(name_.c_str(), value.c_str(), 1); + } + + ~ScopedEnvironmentVariable() + { + if (had_previous_value_) + setenv(name_.c_str(), previous_value_.c_str(), 1); + else + unsetenv(name_.c_str()); + } +}; + CATCH_TEST_CASE( "Homogeneous graph forced physics stores data and reveals native schema", "[wf][graph][storage]") @@ -61,6 +154,8 @@ CATCH_TEST_CASE( AMSCAbstrModel recorder = AMSRegisterAbstractModel("heat_graph_schema", -1.0, "", true); AMSExecutor executor = AMSCreateExecutor(recorder, 0, 1); + const fs::path manifest_path = AMSGetDatabaseName(executor); + CATCH_REQUIRE(manifest_path.filename() == "heat_graph_schema_0_jsondb.json"); // Build small test graph const int64_t N = 10; // nodes @@ -133,9 +228,9 @@ CATCH_TEST_CASE( // INSPECT ACTUAL JSONDB OUTPUT - source of truth for schema // ======================================================================== - CATCH_REQUIRE(fs::exists(test_dir / "manifest.json")); + CATCH_REQUIRE(fs::exists(manifest_path)); - std::ifstream manifest_file(test_dir / "manifest.json"); + std::ifstream manifest_file(manifest_path); nlohmann::json manifest; manifest_file >> manifest; @@ -238,6 +333,7 @@ CATCH_TEST_CASE("Homogeneous graph without globals omits global storage", AMSCAbstrModel recorder = AMSRegisterAbstractModel("empty_globals_domain", -1.0, "", true); AMSExecutor executor = AMSCreateExecutor(recorder, 0, 1); + const fs::path manifest_path = AMSGetDatabaseName(executor); auto node_feat = makeTensor({3, 2}); auto edge_idx = makeTensor({2, 2}); @@ -278,7 +374,7 @@ CATCH_TEST_CASE("Homogeneous graph without globals omits global storage", AMSExecute(executor, physics, graph, outputs); AMSDestroyExecutor(executor); - std::ifstream manifest_file(test_dir / "manifest.json"); + std::ifstream manifest_file(manifest_path); CATCH_REQUIRE(manifest_file.is_open()); nlohmann::json manifest; manifest_file >> manifest; @@ -293,6 +389,154 @@ CATCH_TEST_CASE("Homogeneous graph without globals omits global storage", fs::remove_all(test_dir); } +CATCH_TEST_CASE("JSONDB pure JSON mode stores flat tensors inline", + "[wf][graph][storage][json]") +{ + fs::path test_dir = fs::temp_directory_path() / "ams_json_inline_tensor_test"; + fs::remove_all(test_dir); + fs::create_directories(test_dir); + + torch::Tensor input = + torch::tensor({1.25f, -2.5f}, + torch::TensorOptions().dtype(torch::kFloat32)) + .reshape({1, 2}); + torch::Tensor output = + torch::tensor({3, 4}, torch::TensorOptions().dtype(torch::kInt64)); + std::vector inputs{input}; + std::vector outputs{output}; + + fs::path manifest_path; + { + ams::db::JSONDB db(test_dir.string(), "inline_tensor", 7, "json"); + manifest_path = db.getFilename(); + CATCH_REQUIRE(manifest_path.filename() == "inline_tensor_7_jsondb.json"); + + db.store(inputs, outputs); + db.close(); + CATCH_REQUIRE_NOTHROW(db.close()); + } + + CATCH_REQUIRE(fs::exists(manifest_path)); + std::ifstream manifest_file(manifest_path); + nlohmann::json manifest; + manifest_file >> manifest; + + CATCH_REQUIRE(manifest["cases"].size() == 1); + const auto& tensors = manifest["cases"][0]["tensors"]; + requireInlineTensor(tensors["input_0"], + input.data_ptr(), + input.nbytes(), + "float32", + nlohmann::json::array({1, 2})); + requireInlineTensor(tensors["output_0"], + output.data_ptr(), + output.nbytes(), + "int64", + nlohmann::json::array({2})); + CATCH_REQUIRE_FALSE(containsBinaryFile(test_dir)); + + fs::remove_all(test_dir); +} + +CATCH_TEST_CASE("AMS pure JSON mode stores homogeneous graphs inline", + "[wf][graph][storage][json]") +{ + ScopedEnvironmentVariable json_mode("AMS_JSON_MODE", "json"); + fs::path test_dir = fs::temp_directory_path() / "ams_json_inline_graph_test"; + fs::remove_all(test_dir); + fs::create_directories(test_dir); + + AMSInit(); + AMSConfigureFSDatabase(AMSDBType::AMS_JSON, test_dir.string().c_str()); + + AMSCAbstrModel recorder = + AMSRegisterAbstractModel("inline_graph", -1.0, "", true); + AMSExecutor executor = AMSCreateExecutor(recorder, 3, 4); + const fs::path manifest_path = AMSGetDatabaseName(executor); + CATCH_REQUIRE(manifest_path.filename() == "inline_graph_3_jsondb.json"); + + auto node_features = makeTensor({3, 2}); + float* node_data = node_features.data(); + for (int i = 0; i < 6; ++i) + node_data[i] = static_cast(i) * 0.5f; + + auto edge_index = makeTensor({2, 2}); + int64_t* edge_data = edge_index.data(); + edge_data[0] = 0; + edge_data[1] = 1; + edge_data[2] = 1; + edge_data[3] = 2; + + auto edge_features = makeTensor({2, 1}); + float* edge_feature_data = edge_features.data(); + edge_feature_data[0] = 1.5f; + edge_feature_data[1] = 2.5f; + + auto global_features = makeTensor({2}); + double* global_data = global_features.data(); + global_data[0] = 0.25; + global_data[1] = 1.25; + + AMSHomogeneousGraph graph(std::move(node_features), + std::move(edge_index), + std::move(edge_features), + std::move(global_features)); + AMSHomogeneousGraphFields outputs; + HomogeneousGraphDomainFn physics = [](const AMSHomogeneousGraph& g, + AMSHomogeneousGraphFields& o) { + auto delta = makeTensor({g.node_features.shape()[0], 1}); + double* delta_data = delta.data(); + for (int64_t i = 0; i < g.node_features.shape()[0]; ++i) + delta_data[i] = static_cast(i) + 10.0; + o.node_fields.insert("delta_u", std::move(delta)); + }; + + AMSExecute(executor, physics, graph, outputs); + AMSDestroyExecutor(executor); + + CATCH_REQUIRE(fs::exists(manifest_path)); + std::ifstream manifest_file(manifest_path); + nlohmann::json manifest; + manifest_file >> manifest; + + CATCH_REQUIRE(manifest["cases"].size() == 1); + const auto& tensors = manifest["cases"][0]["tensors"]; + requireInlineTensor(tensors["node_features"], + graph.node_features.raw_data(), + graph.node_features.elements() * + graph.node_features.element_size(), + "float32", + nlohmann::json::array({3, 2})); + requireInlineTensor(tensors["edge_index"], + graph.edge_index.raw_data(), + graph.edge_index.elements() * + graph.edge_index.element_size(), + "int64", + nlohmann::json::array({2, 2})); + requireInlineTensor(tensors["edge_features"], + graph.edge_features.raw_data(), + graph.edge_features.elements() * + graph.edge_features.element_size(), + "float32", + nlohmann::json::array({2, 1})); + requireInlineTensor(tensors["global_features"], + graph.global_features.raw_data(), + graph.global_features.elements() * + graph.global_features.element_size(), + "float64", + nlohmann::json::array({2})); + + const auto& delta = outputs.node_fields.at("delta_u"); + requireInlineTensor(tensors["target_delta_u"], + delta.raw_data(), + delta.elements() * delta.element_size(), + "float64", + nlohmann::json::array({3, 1})); + CATCH_REQUIRE_FALSE(containsBinaryFile(test_dir)); + + fs::remove_all(test_dir); +} + // ============================================================================ // A2 Tests: Complete Storage Test Coverage // ============================================================================ @@ -383,8 +627,8 @@ CATCH_TEST_CASE( // Note: Not destroying executor to avoid triggering AMSFinalize between tests // The executor will be cleaned up at program exit - // Verify NO manifest created (store_data=false) - CATCH_REQUIRE(!fs::exists(test_dir / "manifest.json")); + // Verify no database artifacts were created (store_data=false). + CATCH_REQUIRE(fs::is_empty(test_dir)); fs::remove_all(test_dir); } @@ -402,6 +646,7 @@ CATCH_TEST_CASE("Multiple calls accumulate cases with distinguishable values", AMSCAbstrModel recorder = AMSRegisterAbstractModel("accumulation_domain", -1.0, "", true); AMSExecutor executor = AMSCreateExecutor(recorder, 0, 1); + const fs::path manifest_path = AMSGetDatabaseName(executor); const int num_calls = 3; int total_callbacks = 0; @@ -467,9 +712,9 @@ CATCH_TEST_CASE("Multiple calls accumulate cases with distinguishable values", AMSDestroyExecutor(executor); // Verify manifest exists with correct case count - CATCH_REQUIRE(fs::exists(test_dir / "manifest.json")); + CATCH_REQUIRE(fs::exists(manifest_path)); - std::ifstream manifest_file(test_dir / "manifest.json"); + std::ifstream manifest_file(manifest_path); nlohmann::json manifest; manifest_file >> manifest; @@ -512,6 +757,7 @@ CATCH_TEST_CASE("Heterogeneous graph typed storage", AMSCAbstrModel recorder = AMSRegisterAbstractModel("hetero_domain", -1.0, "", true); AMSExecutor executor = AMSCreateExecutor(recorder, 0, 1); + const fs::path manifest_path = AMSGetDatabaseName(executor); // Build heterogeneous graph with two node types AMSHeterogeneousGraph graph; @@ -588,9 +834,9 @@ CATCH_TEST_CASE("Heterogeneous graph typed storage", // The executor will be cleaned up at program exit // Verify heterogeneous storage - CATCH_REQUIRE(fs::exists(test_dir / "manifest.json")); + CATCH_REQUIRE(fs::exists(manifest_path)); - std::ifstream manifest_file(test_dir / "manifest.json"); + std::ifstream manifest_file(manifest_path); nlohmann::json manifest; manifest_file >> manifest; @@ -800,6 +1046,7 @@ CATCH_TEST_CASE( AMSCAbstrModel fallback = AMSRegisterAbstractModel("reject_domain", 0.1, "", true); AMSExecutor executor = AMSCreateExecutor(fallback, 0, 1); + const fs::path manifest_path = AMSGetDatabaseName(executor); auto node_feat = makeTensor({6, 2}); auto edge_idx = makeTensor({2, 5}); @@ -867,9 +1114,9 @@ CATCH_TEST_CASE( AMSDestroyExecutor(executor); // Verify exactly one case stored - CATCH_REQUIRE(fs::exists(test_dir / "manifest.json")); + CATCH_REQUIRE(fs::exists(manifest_path)); - std::ifstream manifest_file(test_dir / "manifest.json"); + std::ifstream manifest_file(manifest_path); nlohmann::json manifest; manifest_file >> manifest; From a23a74661f2e66f4a7b25b4e57182ad621ef700f Mon Sep 17 00:00:00 2001 From: Yohann Dudouit Date: Tue, 1 Sep 2026 15:17:11 -0700 Subject: [PATCH 12/18] Update CHANGELOG.md --- CHANGELOG.md | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index d9763d28..7da89df8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,11 @@ ## [Unreleased] +### Added + +- JSON-backed storage can emit binary tensor files or self-contained base64 + manifests named for each domain and rank. + ### Changed - Workflow environments can now use active system Flux Python bindings instead From 00aa3a044ad27054c840c9081c58366ea86257b8 Mon Sep 17 00:00:00 2001 From: Yohann Dudouit Date: Tue, 1 Sep 2026 15:52:23 -0700 Subject: [PATCH 13/18] Remove hardcoded `delta_u` --- src/AMSlib/include/AMSGraph.hpp | 7 +++ src/AMSlib/wf/jsondb.cpp | 78 +++++++++++++++++++++------------ src/AMSlib/wf/jsondb.hpp | 24 ++++++++++ 3 files changed, 81 insertions(+), 28 deletions(-) diff --git a/src/AMSlib/include/AMSGraph.hpp b/src/AMSlib/include/AMSGraph.hpp index 6278541d..55408da4 100644 --- a/src/AMSlib/include/AMSGraph.hpp +++ b/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(); } diff --git a/src/AMSlib/wf/jsondb.cpp b/src/AMSlib/wf/jsondb.cpp index cb4f0b89..d5d1f9a5 100644 --- a/src/AMSlib/wf/jsondb.cpp +++ b/src/AMSlib/wf/jsondb.cpp @@ -9,11 +9,13 @@ #include +#include #include #include #include #include #include +#include #include "wf/debug.h" @@ -261,6 +263,49 @@ nlohmann::json JSONDB::encodeBase64Tensor(const torch::Tensor& tensor) {"byte_size", byte_size}}; } +nlohmann::json JSONDB::serializeTensor(const AMSTensor& tensor, + const std::string& binary_path) +{ + if (json_mode_ == "json") { + return encodeBase64Tensor(tensor); + } + + auto shape_ref = tensor.shape(); + std::vector shape(shape_ref.begin(), shape_ref.end()); + size_t byte_size = writeBinaryTensor(tensor, binary_path); + return nlohmann::json{{"path", binary_path}, + {"dtype", dtypeToString(tensor.dType())}, + {"shape", shape}, + {"byte_size", byte_size}}; +} + +nlohmann::json JSONDB::serializeOutputFields(const AMSTensorFieldMap& fields, + const std::string& case_dir, + const std::string& association) +{ + std::vector> sorted_fields; + sorted_fields.reserve(fields.size()); + for (const auto& [name, tensor] : fields) { + sorted_fields.emplace_back(name, &tensor); + } + std::sort(sorted_fields.begin(), + sorted_fields.end(), + [](const auto& lhs, const auto& rhs) { + return lhs.first < rhs.first; + }); + + nlohmann::json output_json = nlohmann::json::object(); + for (size_t i = 0; i < sorted_fields.size(); ++i) { + std::ostringstream filename; + filename << "field_" << std::setw(6) << std::setfill('0') << i << ".bin"; + std::string binary_path = + case_dir + "/outputs/" + association + "/" + filename.str(); + output_json[sorted_fields[i].first] = + serializeTensor(*sorted_fields[i].second, binary_path); + } + return output_json; +} + void JSONDB::validateEdgeIndex(const AMSTensor& edge_index, int64_t num_nodes) { auto shape_ref = edge_index.shape(); @@ -484,35 +529,12 @@ void JSONDB::store(const ams::AMSHomogeneousGraph& graph, } } - // Write targets from outputs.node_fields - // For now, we look for specific known target names - // TODO: Make this more generic with iterator support in AMSTensorFieldMap - int target_dim = 0; - - // Check for "delta_u" target (heat_equation convention) - const AMSTensor* delta_u = outputs.node_fields.find("delta_u"); - if (delta_u != nullptr) { - auto t_shape = delta_u->shape(); - target_dim = (t_shape.size() > 1) ? t_shape[1] : 1; - - if (json_mode_ == "binary") { - std::string rel_path = case_dir + "/target_delta_u.bin"; - size_t byte_size = writeBinaryTensor(*delta_u, rel_path); - - tensors_json["target_delta_u"] = { - {"path", rel_path}, - {"dtype", dtypeToString(delta_u->dType())}, - {"shape", std::vector{num_nodes, target_dim}}, - {"byte_size", byte_size}}; - } else { // Pure json mode - tensors_json["target_delta_u"] = encodeBase64Tensor(*delta_u); - } - } - - // Add target_dim to case metadata - case_json["target_dim"] = target_dim; - case_json["tensors"] = tensors_json; + case_json["outputs"] = { + {"node", serializeOutputFields(outputs.node_fields, case_dir, "node")}, + {"edge", serializeOutputFields(outputs.edge_fields, case_dir, "edge")}, + {"global", + serializeOutputFields(outputs.global_fields, case_dir, "global")}}; cases_.push_back(case_json); case_counter_++; diff --git a/src/AMSlib/wf/jsondb.hpp b/src/AMSlib/wf/jsondb.hpp index 0a0c179e..e27e2bd4 100644 --- a/src/AMSlib/wf/jsondb.hpp +++ b/src/AMSlib/wf/jsondb.hpp @@ -33,6 +33,10 @@ namespace db * - "binary": Binary tensor files + JSON manifest (default, efficient) * - "json": Self-contained JSON with base64-encoded tensor data * + * Homogeneous graph inputs are stored under each case's "tensors" object. + * Named graph outputs are stored under "outputs", grouped into "node", "edge", + * and "global" objects using the application-provided field names. + * * Output format is compatible with PyTorch Geometric data loaders. * * @note Case and step directories are not rank-scoped. Concurrent MPI ranks @@ -85,6 +89,26 @@ class JSONDB final : public FileDB */ nlohmann::json encodeBase64Tensor(const torch::Tensor& tensor); + /** + * @brief Serialize an AMSTensor using the configured JSON mode + * @param[in] tensor The tensor to serialize + * @param[in] binary_path Relative binary path used in binary mode + * @return JSON tensor descriptor + */ + nlohmann::json serializeTensor(const AMSTensor& tensor, + const std::string& binary_path); + + /** + * @brief Serialize all named output fields for one graph association + * @param[in] fields Named output fields + * @param[in] case_dir Relative case directory + * @param[in] association Output association: node, edge, or global + * @return JSON object keyed by application field name + */ + nlohmann::json serializeOutputFields(const AMSTensorFieldMap& fields, + const std::string& case_dir, + const std::string& association); + /** * @brief Validate edge_index tensor format * @param[in] edge_index Edge connectivity tensor [2, E] From 96732e2699af7b59f06e9fd824602943391c3d52 Mon Sep 17 00:00:00 2001 From: Yohann Dudouit Date: Tue, 1 Sep 2026 15:52:40 -0700 Subject: [PATCH 14/18] Strengthen tests. --- .../ams_interface/test_graph_fallback.cpp | 19 ++ .../test_graph_workflow_storage.cpp | 238 ++++++++++++------ 2 files changed, 186 insertions(+), 71 deletions(-) diff --git a/tests/AMSlib/ams_interface/test_graph_fallback.cpp b/tests/AMSlib/ams_interface/test_graph_fallback.cpp index 16e6e75f..d6a055c1 100644 --- a/tests/AMSlib/ams_interface/test_graph_fallback.cpp +++ b/tests/AMSlib/ams_interface/test_graph_fallback.cpp @@ -116,6 +116,25 @@ CATCH_TEST_CASE("AMSTensorFieldMap explicit field API", "[wf][graph]") CATCH_REQUIRE_THROWS_AS(fields.insert("flux", makeTensor({2, 1})), std::runtime_error); CATCH_REQUIRE_THROWS_AS(fields.at("absent"), std::out_of_range); + + const AMSTensorFieldMap& const_fields = fields; + CATCH_REQUIRE(const_fields.cbegin() != const_fields.cend()); + std::size_t visited = 0; + bool saw_prediction = false; + bool saw_flux = false; + for (const auto& [name, tensor] : const_fields) { + ++visited; + if (name == "prediction") { + saw_prediction = true; + CATCH_REQUIRE(tensor.shape()[0] == 3); + } else if (name == "flux") { + saw_flux = true; + CATCH_REQUIRE(tensor.shape()[0] == 2); + } + } + CATCH_REQUIRE(visited == 2); + CATCH_REQUIRE(saw_prediction); + CATCH_REQUIRE(saw_flux); } CATCH_TEST_CASE("AMSHomogeneousGraph validates construction", diff --git a/tests/AMSlib/ams_interface/test_graph_workflow_storage.cpp b/tests/AMSlib/ams_interface/test_graph_workflow_storage.cpp index d745e79c..523571b7 100644 --- a/tests/AMSlib/ams_interface/test_graph_workflow_storage.cpp +++ b/tests/AMSlib/ams_interface/test_graph_workflow_storage.cpp @@ -14,7 +14,6 @@ #include #include #include -#include #include #include #include @@ -99,6 +98,35 @@ static void requireInlineTensor(const nlohmann::json& tensor, std::memcmp(decoded.data(), expected_data, expected_byte_size) == 0); } +static void requireBinaryTensor(const fs::path& root, + const nlohmann::json& tensor, + const void* expected_data, + size_t expected_byte_size, + const std::string& expected_dtype, + const nlohmann::json& expected_shape, + const fs::path& expected_relative_path) +{ + CATCH_REQUIRE(tensor["dtype"] == expected_dtype); + CATCH_REQUIRE(tensor["shape"] == expected_shape); + CATCH_REQUIRE(tensor["byte_size"] == expected_byte_size); + CATCH_REQUIRE_FALSE(tensor.contains("encoding")); + CATCH_REQUIRE_FALSE(tensor.contains("data")); + + const fs::path relative_path = tensor["path"].get(); + CATCH_REQUIRE(relative_path == expected_relative_path); + CATCH_REQUIRE_FALSE(relative_path.is_absolute()); + + std::ifstream file(root / relative_path, std::ios::binary); + CATCH_REQUIRE(file.is_open()); + std::vector actual(expected_byte_size); + file.read(reinterpret_cast(actual.data()), + static_cast(actual.size())); + CATCH_REQUIRE(file.gcount() == + static_cast(expected_byte_size)); + CATCH_REQUIRE(std::memcmp(actual.data(), expected_data, expected_byte_size) == + 0); +} + static bool containsBinaryFile(const fs::path& directory) { for (const auto& entry : fs::recursive_directory_iterator(directory)) { @@ -136,13 +164,12 @@ class ScopedEnvironmentVariable } }; -CATCH_TEST_CASE( - "Homogeneous graph forced physics stores data and reveals native schema", - "[wf][graph][storage]") +CATCH_TEST_CASE("Homogeneous graph stores every named output in binary mode", + "[wf][graph][storage]") { // Use unique temp directory fs::path test_dir = - fs::temp_directory_path() / "ams_graph_schema_discovery_test"; + fs::temp_directory_path() / "ams_graph_named_output_storage_test"; fs::remove_all(test_dir); fs::create_directories(test_dir); @@ -152,10 +179,11 @@ CATCH_TEST_CASE( // Recorder configuration: threshold < 0 forces physics, store_data=true enables DB AMSCAbstrModel recorder = - AMSRegisterAbstractModel("heat_graph_schema", -1.0, "", true); + AMSRegisterAbstractModel("generic_graph_schema", -1.0, "", true); AMSExecutor executor = AMSCreateExecutor(recorder, 0, 1); const fs::path manifest_path = AMSGetDatabaseName(executor); - CATCH_REQUIRE(manifest_path.filename() == "heat_graph_schema_0_jsondb.json"); + CATCH_REQUIRE(manifest_path.filename() == + "generic_graph_schema_0_jsondb.json"); // Build small test graph const int64_t N = 10; // nodes @@ -197,21 +225,41 @@ CATCH_TEST_CASE( AMSHomogeneousGraphFields outputs; - // Physics callback that returns float64 delta_u + // Physics callback with multiple application-defined output fields. int callback_count = 0; HomogeneousGraphDomainFn physics = [&](const AMSHomogeneousGraph& g, AMSHomogeneousGraphFields& o) { callback_count++; - // Return float64 delta_u [N, 1] (MFEM precision) const int64_t num_nodes = g.node_features.shape()[0]; - auto delta = makeTensor({num_nodes, 1}); - double* data = delta.data(); + const int64_t num_edges = g.edge_index.shape()[1]; + + auto pressure = makeTensor({num_nodes, 2}); + float* pressure_data = pressure.data(); + for (int64_t i = 0; i < pressure.elements(); ++i) { + pressure_data[i] = static_cast(i) + 0.5f; + } + o.node_fields.insert("pressure/drop", std::move(pressure)); + + auto temperature = makeTensor({num_nodes, 1}); + double* temperature_data = temperature.data(); for (int64_t i = 0; i < num_nodes; i++) { - data[i] = static_cast(i) * 0.123; + temperature_data[i] = static_cast(i) * 0.123; } + o.node_fields.insert("temperature", std::move(temperature)); - o.node_fields.insert("delta_u", std::move(delta)); + auto flux = makeTensor({num_edges, 1}); + float* flux_data = flux.data(); + for (int64_t i = 0; i < num_edges; ++i) { + flux_data[i] = static_cast(i) * 1.5f; + } + o.edge_fields.insert("flux", std::move(flux)); + + auto loss = makeTensor({1, 2}); + double* loss_data = loss.data(); + loss_data[0] = 0.25; + loss_data[1] = 0.75; + o.global_fields.insert("loss", std::move(loss)); }; // Execute with forced physics + storage @@ -219,26 +267,20 @@ CATCH_TEST_CASE( // Verify physics ran CATCH_REQUIRE(callback_count == 1); - CATCH_REQUIRE(outputs.node_fields.find("delta_u") != nullptr); + CATCH_REQUIRE(outputs.node_fields.contains("pressure/drop")); + CATCH_REQUIRE(outputs.node_fields.contains("temperature")); + CATCH_REQUIRE(outputs.edge_fields.contains("flux")); + CATCH_REQUIRE(outputs.global_fields.contains("loss")); // Destroy executor to flush manifest (but don't call AMSFinalize) AMSDestroyExecutor(executor); - // ======================================================================== - // INSPECT ACTUAL JSONDB OUTPUT - source of truth for schema - // ======================================================================== - CATCH_REQUIRE(fs::exists(manifest_path)); std::ifstream manifest_file(manifest_path); nlohmann::json manifest; manifest_file >> manifest; - // Document native schema structure - std::cout << "\n=== NATIVE JSONDB SCHEMA (A1 Discovery) ===\n"; - std::cout << manifest.dump(2) << std::endl; - std::cout << "==========================================\n" << std::endl; - // Verify essential structure exists CATCH_REQUIRE(manifest.contains("format_version")); CATCH_REQUIRE(manifest.contains("endianness")); @@ -263,12 +305,12 @@ CATCH_TEST_CASE( case0["tensors"]["global_features"]["byte_size"].get() == 2 * sizeof(float)); - // Verify output field stored - // Note: Field name may be "delta_u" or "target_delta_u" - discover actual - bool has_delta_u = case0["tensors"].contains("delta_u") || - case0["tensors"].contains("target_delta_u") || - case0["tensors"].contains("node:delta_u"); - CATCH_REQUIRE(has_delta_u); + CATCH_REQUIRE(case0.contains("outputs")); + CATCH_REQUIRE(case0["outputs"]["node"].size() == 2); + CATCH_REQUIRE(case0["outputs"]["edge"].size() == 1); + CATCH_REQUIRE(case0["outputs"]["global"].size() == 1); + CATCH_REQUIRE_FALSE(case0.contains("target_dim")); + CATCH_REQUIRE_FALSE(case0["tensors"].contains("target_delta_u")); // Verify dtypes preserved CATCH_REQUIRE(case0["tensors"]["node_features"]["dtype"].get() == @@ -293,15 +335,41 @@ CATCH_TEST_CASE( CATCH_REQUIRE(stored_globals[0] == 0.01f); CATCH_REQUIRE(stored_globals[1] == 0.123f); - // Verify float64 delta_u (CRITICAL for MFEM precision) - std::string delta_key = "delta_u"; - if (case0["tensors"].contains("target_delta_u")) { - delta_key = "target_delta_u"; - } else if (case0["tensors"].contains("node:delta_u")) { - delta_key = "node:delta_u"; - } - CATCH_REQUIRE(case0["tensors"][delta_key]["dtype"].get() == - "float64"); + const auto& pressure = outputs.node_fields.at("pressure/drop"); + requireBinaryTensor(test_dir, + case0["outputs"]["node"]["pressure/drop"], + pressure.raw_data(), + pressure.elements() * pressure.element_size(), + "float32", + nlohmann::json::array({N, 2}), + fs::path("step_000000/outputs/node/field_000000.bin")); + + const auto& temperature = outputs.node_fields.at("temperature"); + requireBinaryTensor(test_dir, + case0["outputs"]["node"]["temperature"], + temperature.raw_data(), + temperature.elements() * temperature.element_size(), + "float64", + nlohmann::json::array({N, 1}), + fs::path("step_000000/outputs/node/field_000001.bin")); + + const auto& flux = outputs.edge_fields.at("flux"); + requireBinaryTensor(test_dir, + case0["outputs"]["edge"]["flux"], + flux.raw_data(), + flux.elements() * flux.element_size(), + "float32", + nlohmann::json::array({E, 1}), + fs::path("step_000000/outputs/edge/field_000000.bin")); + + const auto& loss = outputs.global_fields.at("loss"); + requireBinaryTensor(test_dir, + case0["outputs"]["global"]["loss"], + loss.raw_data(), + loss.elements() * loss.element_size(), + "float64", + nlohmann::json::array({1, 2}), + fs::path("step_000000/outputs/global/field_000000.bin")); // Verify paths are relative to dataset root std::string node_path = @@ -309,13 +377,7 @@ CATCH_TEST_CASE( CATCH_REQUIRE(!fs::path(node_path).is_absolute()); CATCH_REQUIRE(fs::exists(test_dir / node_path)); - // Note: Not destroying executor or calling AMSFinalize to avoid lifecycle - // issues between tests. Cleanup happens at process exit. - // fs::remove_all(test_dir); // Keep for manual inspection - - std::cout << "A1 schema discovery test PASSED. Review manifest output above " - "before proceeding to A2." - << std::endl; + fs::remove_all(test_dir); } CATCH_TEST_CASE("Homogeneous graph without globals omits global storage", @@ -361,15 +423,8 @@ CATCH_TEST_CASE("Homogeneous graph without globals omits global storage", CATCH_REQUIRE(graph.global_features.shape()[0] == 0); AMSHomogeneousGraphFields outputs; - HomogeneousGraphDomainFn physics = [](const AMSHomogeneousGraph& g, - AMSHomogeneousGraphFields& o) { - auto delta = makeTensor({g.node_features.shape()[0], 1}); - double* delta_data = delta.data(); - for (int64_t i = 0; i < g.node_features.shape()[0]; ++i) { - delta_data[i] = static_cast(i); - } - o.node_fields.insert("delta_u", std::move(delta)); - }; + HomogeneousGraphDomainFn physics = [](const AMSHomogeneousGraph&, + AMSHomogeneousGraphFields&) {}; AMSExecute(executor, physics, graph, outputs); AMSDestroyExecutor(executor); @@ -385,6 +440,10 @@ CATCH_TEST_CASE("Homogeneous graph without globals omits global storage", CATCH_REQUIRE_FALSE(stored_case["tensors"].contains("global_features")); CATCH_REQUIRE_FALSE( fs::exists(test_dir / "step_000000" / "global_features.bin")); + CATCH_REQUIRE(stored_case["outputs"]["node"].empty()); + CATCH_REQUIRE(stored_case["outputs"]["edge"].empty()); + CATCH_REQUIRE(stored_case["outputs"]["global"].empty()); + CATCH_REQUIRE_FALSE(stored_case.contains("target_dim")); fs::remove_all(test_dir); } @@ -484,11 +543,29 @@ CATCH_TEST_CASE("AMS pure JSON mode stores homogeneous graphs inline", AMSHomogeneousGraphFields outputs; HomogeneousGraphDomainFn physics = [](const AMSHomogeneousGraph& g, AMSHomogeneousGraphFields& o) { - auto delta = makeTensor({g.node_features.shape()[0], 1}); - double* delta_data = delta.data(); + auto pressure = makeTensor({g.node_features.shape()[0], 2}); + float* pressure_data = pressure.data(); + for (int64_t i = 0; i < pressure.elements(); ++i) + pressure_data[i] = static_cast(i) + 0.25f; + o.node_fields.insert("pressure/drop", std::move(pressure)); + + auto temperature = makeTensor({g.node_features.shape()[0], 1}); + double* temperature_data = temperature.data(); for (int64_t i = 0; i < g.node_features.shape()[0]; ++i) - delta_data[i] = static_cast(i) + 10.0; - o.node_fields.insert("delta_u", std::move(delta)); + temperature_data[i] = static_cast(i) + 10.0; + o.node_fields.insert("temperature", std::move(temperature)); + + auto flux = makeTensor({g.edge_index.shape()[1], 1}); + float* flux_data = flux.data(); + for (int64_t i = 0; i < g.edge_index.shape()[1]; ++i) + flux_data[i] = static_cast(i) + 20.0f; + o.edge_fields.insert("flux", std::move(flux)); + + auto loss = makeTensor({1, 2}); + double* loss_data = loss.data(); + loss_data[0] = 30.0; + loss_data[1] = 31.0; + o.global_fields.insert("loss", std::move(loss)); }; AMSExecute(executor, physics, graph, outputs); @@ -526,12 +603,33 @@ CATCH_TEST_CASE("AMS pure JSON mode stores homogeneous graphs inline", "float64", nlohmann::json::array({2})); - const auto& delta = outputs.node_fields.at("delta_u"); - requireInlineTensor(tensors["target_delta_u"], - delta.raw_data(), - delta.elements() * delta.element_size(), + const auto& stored_outputs = manifest["cases"][0]["outputs"]; + const auto& pressure = outputs.node_fields.at("pressure/drop"); + requireInlineTensor(stored_outputs["node"]["pressure/drop"], + pressure.raw_data(), + pressure.elements() * pressure.element_size(), + "float32", + nlohmann::json::array({3, 2})); + const auto& temperature = outputs.node_fields.at("temperature"); + requireInlineTensor(stored_outputs["node"]["temperature"], + temperature.raw_data(), + temperature.elements() * temperature.element_size(), "float64", nlohmann::json::array({3, 1})); + const auto& flux = outputs.edge_fields.at("flux"); + requireInlineTensor(stored_outputs["edge"]["flux"], + flux.raw_data(), + flux.elements() * flux.element_size(), + "float32", + nlohmann::json::array({2, 1})); + const auto& loss = outputs.global_fields.at("loss"); + requireInlineTensor(stored_outputs["global"]["loss"], + loss.raw_data(), + loss.elements() * loss.element_size(), + "float64", + nlohmann::json::array({1, 2})); + CATCH_REQUIRE_FALSE(manifest["cases"][0].contains("target_dim")); + CATCH_REQUIRE_FALSE(tensors.contains("target_delta_u")); CATCH_REQUIRE_FALSE(containsBinaryFile(test_dir)); fs::remove_all(test_dir); @@ -735,8 +833,8 @@ CATCH_TEST_CASE("Multiple calls accumulate cases with distinguishable values", f.read(reinterpret_cast(&global_val), sizeof(float)); CATCH_REQUIRE(global_val == static_cast(call)); - // Verify target output exists - CATCH_REQUIRE(case_entry["tensors"].contains("target_delta_u")); + // Verify output exists + CATCH_REQUIRE(case_entry["outputs"]["node"].contains("delta_u")); } fs::remove_all(test_dir); @@ -1125,7 +1223,7 @@ CATCH_TEST_CASE( auto case0 = manifest["cases"][0]; // Verify stored output matches exact physics output - std::string target_path = case0["tensors"]["target_delta_u"]["path"]; + std::string target_path = case0["outputs"]["node"]["delta_u"]["path"]; fs::path full_path = test_dir / target_path; CATCH_REQUIRE(fs::exists(full_path)); @@ -1140,14 +1238,12 @@ CATCH_TEST_CASE( CATCH_REQUIRE(std::abs(stored_values[i] - expected) < 1e-12); } - // Verify field name (runtime) vs storage name (target_ prefix) - // Runtime API: node_fields["delta_u"] - // Storage: "target_delta_u" in manifest - CATCH_REQUIRE(case0["tensors"].contains("target_delta_u")); - CATCH_REQUIRE(!case0["tensors"].contains("delta_u")); // No prefix in storage + // Verify the application field name is preserved in the node output group. + CATCH_REQUIRE(case0["outputs"]["node"].contains("delta_u")); + CATCH_REQUIRE_FALSE(case0["tensors"].contains("target_delta_u")); // Verify dtype preserved as float64 - CATCH_REQUIRE(case0["tensors"]["target_delta_u"]["dtype"] == "float64"); + CATCH_REQUIRE(case0["outputs"]["node"]["delta_u"]["dtype"] == "float64"); fs::remove_all(test_dir); } From 3039f5fbcef42b7faa94f8860340a794a3ad3483 Mon Sep 17 00:00:00 2001 From: Yohann Dudouit Date: Wed, 2 Sep 2026 09:40:23 -0700 Subject: [PATCH 15/18] Fix CMakeLists for ams_interface, privately link with MPI and caliper. --- tests/AMSlib/ams_interface/CMakeLists.txt | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/tests/AMSlib/ams_interface/CMakeLists.txt b/tests/AMSlib/ams_interface/CMakeLists.txt index cde06056..1f070c8f 100644 --- a/tests/AMSlib/ams_interface/CMakeLists.txt +++ b/tests/AMSlib/ams_interface/CMakeLists.txt @@ -13,7 +13,16 @@ function(BUILD_UNIT_TEST exe source) target_link_libraries(${exe} PRIVATE hip::host) target_link_libraries(${exe} PRIVATE Threads::Threads) endif() - + + if (ENABLE_MPI) + target_link_libraries(${exe} PRIVATE MPI::MPI_CXX) + endif() + + if (ENABLE_CALIPER) + target_link_libraries(${exe} PRIVATE caliper) + target_include_directories(${exe} PRIVATE ${caliper_INCLUDE_DIR}) + endif() + target_include_directories(${exe} PRIVATE ${CMAKE_SOURCE_DIR}/src/AMSlib/) target_include_directories(${exe} PRIVATE ${CMAKE_BINARY_DIR}/include/) target_include_directories(${exe} PRIVATE ${AMS_TEST_ROOT}) From c44809abfb622e5af80160627fb14dd25f6aaf8d Mon Sep 17 00:00:00 2001 From: Yohann Dudouit Date: Tue, 8 Sep 2026 13:40:53 -0700 Subject: [PATCH 16/18] Use rank ID in filenames. --- src/AMSlib/wf/jsondb.cpp | 6 ++++-- src/AMSlib/wf/jsondb.hpp | 5 +++-- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/src/AMSlib/wf/jsondb.cpp b/src/AMSlib/wf/jsondb.cpp index d5d1f9a5..b1c456d5 100644 --- a/src/AMSlib/wf/jsondb.cpp +++ b/src/AMSlib/wf/jsondb.cpp @@ -358,7 +358,8 @@ void JSONDB::store(ArrayRef Inputs, { // Create case directory std::ostringstream case_name; - case_name << "case_" << std::setw(6) << std::setfill('0') << case_counter_; + case_name << "case_" << getId() << "_" << std::setw(6) << std::setfill('0') + << case_counter_; std::string case_dir = case_name.str(); nlohmann::json case_json; @@ -431,7 +432,8 @@ void JSONDB::store(const ams::AMSHomogeneousGraph& graph, { // Create case directory std::ostringstream case_name; - case_name << "step_" << std::setw(6) << std::setfill('0') << case_counter_; + case_name << "step_" << getId() << "_" << std::setw(6) << std::setfill('0') + << case_counter_; std::string case_dir = case_name.str(); // Extract graph dimensions diff --git a/src/AMSlib/wf/jsondb.hpp b/src/AMSlib/wf/jsondb.hpp index e27e2bd4..7b54bc83 100644 --- a/src/AMSlib/wf/jsondb.hpp +++ b/src/AMSlib/wf/jsondb.hpp @@ -39,8 +39,9 @@ namespace db * * Output format is compatible with PyTorch Geometric data loaders. * - * @note Case and step directories are not rank-scoped. Concurrent MPI ranks - * writing to the same output directory are not currently supported. + * @note Case and step directory names include the database rank ID so + * concurrent MPI ranks can write to the same output directory without tensor + * file collisions. */ class JSONDB final : public FileDB { From df31d8dc30084646e35742eb540edf409840fdd4 Mon Sep 17 00:00:00 2001 From: Yohann Dudouit Date: Tue, 8 Sep 2026 13:41:11 -0700 Subject: [PATCH 17/18] Use isLittleEndian in manifest. --- src/AMSlib/wf/jsondb.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/AMSlib/wf/jsondb.cpp b/src/AMSlib/wf/jsondb.cpp index b1c456d5..0c003bd6 100644 --- a/src/AMSlib/wf/jsondb.cpp +++ b/src/AMSlib/wf/jsondb.cpp @@ -566,7 +566,7 @@ void JSONDB::close() // Build complete manifest nlohmann::json manifest; manifest["format_version"] = 1; - manifest["endianness"] = "little"; + manifest["endianness"] = isLittleEndian() ? "little" : "big"; // Add metadata if set if (!metadata_.is_null()) { From 07272f4a6e93bff2f8604e7d588aa8af4dc88f43 Mon Sep 17 00:00:00 2001 From: Yohann Dudouit Date: Tue, 8 Sep 2026 13:42:14 -0700 Subject: [PATCH 18/18] Update graph storage tests. --- CHANGELOG.md | 5 +- .../test_graph_workflow_storage.cpp | 100 ++++++++++++++++-- 2 files changed, 97 insertions(+), 8 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7da89df8..8531da4b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,8 +4,9 @@ ### Added -- JSON-backed storage can emit binary tensor files or self-contained base64 - manifests named for each domain and rank. +- 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 diff --git a/tests/AMSlib/ams_interface/test_graph_workflow_storage.cpp b/tests/AMSlib/ams_interface/test_graph_workflow_storage.cpp index 523571b7..1a626584 100644 --- a/tests/AMSlib/ams_interface/test_graph_workflow_storage.cpp +++ b/tests/AMSlib/ams_interface/test_graph_workflow_storage.cpp @@ -80,6 +80,13 @@ static std::vector decodeBase64(const std::string& encoded) return decoded; } +static std::string nativeEndianness() +{ + const uint16_t test = 0x0102; + const auto* bytes = reinterpret_cast(&test); + return bytes[0] == 0x02 ? "little" : "big"; +} + static void requireInlineTensor(const nlohmann::json& tensor, const void* expected_data, size_t expected_byte_size, @@ -283,12 +290,13 @@ CATCH_TEST_CASE("Homogeneous graph stores every named output in binary mode", // Verify essential structure exists CATCH_REQUIRE(manifest.contains("format_version")); - CATCH_REQUIRE(manifest.contains("endianness")); + CATCH_REQUIRE(manifest["endianness"] == nativeEndianness()); CATCH_REQUIRE(manifest.contains("cases")); CATCH_REQUIRE(manifest["cases"].is_array()); CATCH_REQUIRE(manifest["cases"].size() == 1); auto case0 = manifest["cases"][0]; + CATCH_REQUIRE(case0["name"] == "step_0_000000"); CATCH_REQUIRE(case0.contains("tensors")); // Verify all graph components stored @@ -342,7 +350,7 @@ CATCH_TEST_CASE("Homogeneous graph stores every named output in binary mode", pressure.elements() * pressure.element_size(), "float32", nlohmann::json::array({N, 2}), - fs::path("step_000000/outputs/node/field_000000.bin")); + fs::path("step_0_000000/outputs/node/field_000000.bin")); const auto& temperature = outputs.node_fields.at("temperature"); requireBinaryTensor(test_dir, @@ -351,7 +359,7 @@ CATCH_TEST_CASE("Homogeneous graph stores every named output in binary mode", temperature.elements() * temperature.element_size(), "float64", nlohmann::json::array({N, 1}), - fs::path("step_000000/outputs/node/field_000001.bin")); + fs::path("step_0_000000/outputs/node/field_000001.bin")); const auto& flux = outputs.edge_fields.at("flux"); requireBinaryTensor(test_dir, @@ -360,7 +368,7 @@ CATCH_TEST_CASE("Homogeneous graph stores every named output in binary mode", flux.elements() * flux.element_size(), "float32", nlohmann::json::array({E, 1}), - fs::path("step_000000/outputs/edge/field_000000.bin")); + fs::path("step_0_000000/outputs/edge/field_000000.bin")); const auto& loss = outputs.global_fields.at("loss"); requireBinaryTensor(test_dir, @@ -369,7 +377,8 @@ CATCH_TEST_CASE("Homogeneous graph stores every named output in binary mode", loss.elements() * loss.element_size(), "float64", nlohmann::json::array({1, 2}), - fs::path("step_000000/outputs/global/field_000000.bin")); + fs::path("step_0_000000/outputs/global/" + "field_000000.bin")); // Verify paths are relative to dataset root std::string node_path = @@ -439,7 +448,7 @@ CATCH_TEST_CASE("Homogeneous graph without globals omits global storage", CATCH_REQUIRE(stored_case["global_feature_dim"].get() == 0); CATCH_REQUIRE_FALSE(stored_case["tensors"].contains("global_features")); CATCH_REQUIRE_FALSE( - fs::exists(test_dir / "step_000000" / "global_features.bin")); + fs::exists(test_dir / "step_0_000000" / "global_features.bin")); CATCH_REQUIRE(stored_case["outputs"]["node"].empty()); CATCH_REQUIRE(stored_case["outputs"]["edge"].empty()); CATCH_REQUIRE(stored_case["outputs"]["global"].empty()); @@ -480,7 +489,9 @@ CATCH_TEST_CASE("JSONDB pure JSON mode stores flat tensors inline", nlohmann::json manifest; manifest_file >> manifest; + CATCH_REQUIRE(manifest["endianness"] == nativeEndianness()); CATCH_REQUIRE(manifest["cases"].size() == 1); + CATCH_REQUIRE(manifest["cases"][0]["name"] == "case_7_000000"); const auto& tensors = manifest["cases"][0]["tensors"]; requireInlineTensor(tensors["input_0"], input.data_ptr(), @@ -497,6 +508,82 @@ CATCH_TEST_CASE("JSONDB pure JSON mode stores flat tensors inline", fs::remove_all(test_dir); } +CATCH_TEST_CASE("JSONDB binary tensor paths distinguish rank IDs", + "[wf][graph][storage]") +{ + fs::path test_dir = + fs::temp_directory_path() / "ams_json_ranked_tensor_storage_test"; + fs::remove_all(test_dir); + fs::create_directories(test_dir); + + torch::Tensor rank_three_input = + torch::tensor({3.25f}, torch::TensorOptions().dtype(torch::kFloat32)); + torch::Tensor rank_seven_input = + torch::tensor({7.5f}, torch::TensorOptions().dtype(torch::kFloat32)); + std::vector no_outputs; + + fs::path rank_three_manifest; + { + ams::db::JSONDB db(test_dir.string(), "shared_domain", 3, "binary"); + rank_three_manifest = db.getFilename(); + std::vector inputs{rank_three_input}; + db.store(inputs, no_outputs); + db.close(); + } + + fs::path rank_seven_manifest; + { + ams::db::JSONDB db(test_dir.string(), "shared_domain", 7, "binary"); + rank_seven_manifest = db.getFilename(); + std::vector inputs{rank_seven_input}; + db.store(inputs, no_outputs); + db.close(); + } + + const fs::path rank_three_filename = "shared_domain_3_jsondb.json"; + const fs::path rank_seven_filename = "shared_domain_7_jsondb.json"; + CATCH_REQUIRE(rank_three_manifest.filename() == rank_three_filename); + CATCH_REQUIRE(rank_seven_manifest.filename() == rank_seven_filename); + + std::ifstream rank_three_file(rank_three_manifest); + std::ifstream rank_seven_file(rank_seven_manifest); + CATCH_REQUIRE(rank_three_file.is_open()); + CATCH_REQUIRE(rank_seven_file.is_open()); + + nlohmann::json rank_three_json; + nlohmann::json rank_seven_json; + rank_three_file >> rank_three_json; + rank_seven_file >> rank_seven_json; + + CATCH_REQUIRE(rank_three_json["endianness"] == nativeEndianness()); + CATCH_REQUIRE(rank_seven_json["endianness"] == nativeEndianness()); + CATCH_REQUIRE(rank_three_json["cases"].size() == 1); + CATCH_REQUIRE(rank_seven_json["cases"].size() == 1); + const auto& rank_three_case = rank_three_json["cases"][0]; + const auto& rank_seven_case = rank_seven_json["cases"][0]; + CATCH_REQUIRE(rank_three_case["name"] == "case_3_000000"); + CATCH_REQUIRE(rank_seven_case["name"] == "case_7_000000"); + + requireBinaryTensor(test_dir, + rank_three_case["tensors"]["input_0"], + rank_three_input.data_ptr(), + rank_three_input.nbytes(), + "float32", + nlohmann::json::array({1}), + fs::path("case_3_000000/input_0.bin")); + requireBinaryTensor(test_dir, + rank_seven_case["tensors"]["input_0"], + rank_seven_input.data_ptr(), + rank_seven_input.nbytes(), + "float32", + nlohmann::json::array({1}), + fs::path("case_7_000000/input_0.bin")); + CATCH_REQUIRE(rank_three_case["tensors"]["input_0"]["path"] != + rank_seven_case["tensors"]["input_0"]["path"]); + + fs::remove_all(test_dir); +} + CATCH_TEST_CASE("AMS pure JSON mode stores homogeneous graphs inline", "[wf][graph][storage][json]") { @@ -577,6 +664,7 @@ CATCH_TEST_CASE("AMS pure JSON mode stores homogeneous graphs inline", manifest_file >> manifest; CATCH_REQUIRE(manifest["cases"].size() == 1); + CATCH_REQUIRE(manifest["cases"][0]["name"] == "step_3_000000"); const auto& tensors = manifest["cases"][0]["tensors"]; requireInlineTensor(tensors["node_features"], graph.node_features.raw_data(),