diff --git a/CHANGELOG.md b/CHANGELOG.md index d9763d28..8531da4b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,12 @@ ## [Unreleased] +### Added + +- JSON-backed storage can emit binary tensor files in rank-qualified case + directories or self-contained base64 manifests, with a separate manifest for + each domain and rank (#205). + ### Changed - Workflow environments can now use active system Flux Python bindings instead 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/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/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/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..ac1a9cc3 100644 --- a/src/AMSlib/wf/basedb.cpp +++ b/src/AMSlib/wf/basedb.cpp @@ -1,6 +1,10 @@ +#include "wf/basedb.hpp" + +#include #include #include "AMS.h" +#include "wf/jsondb.hpp" namespace ams { @@ -17,6 +21,8 @@ AMSDBType getDBType(std::string type) return AMSDBType::AMS_HDF5; } else if (type.compare("rmq") == 0) { return AMSDBType::AMS_RMQ; + } else if (type.compare("json") == 0) { + return AMSDBType::AMS_JSON; } return AMSDBType::AMS_NONE; } @@ -30,10 +36,52 @@ std::string getDBTypeAsStr(AMSDBType type) return "hdf5"; case AMSDBType::AMS_RMQ: return "rmq"; + case AMSDBType::AMS_JSON: + return "json"; } return "Unknown"; } +std::shared_ptr 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..5ac32e6f 100644 --- a/src/AMSlib/wf/basedb.hpp +++ b/src/AMSlib/wf/basedb.hpp @@ -31,7 +31,14 @@ #include "wf/resource_manager.hpp" #include "wf/utils.hpp" -namespace fs = std::experimental::filesystem; +// Forward declarations for graph types +namespace ams +{ +struct AMSHomogeneousGraph; +struct AMSHomogeneousGraphFields; +struct AMSHeterogeneousGraph; +struct AMSHeterogeneousGraphFields; +} // namespace ams #ifdef __AMS_ENABLE_HDF5__ #include @@ -124,6 +131,35 @@ 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( + [[maybe_unused]] const ams::AMSHomogeneousGraph& graph, + [[maybe_unused]] const ams::AMSHomogeneousGraphFields& outputs) + { + THROW(std::runtime_error, + (this->type() + " database does not support graph storage").c_str()); + } + + /** + * @brief Store heterogeneous graph data with outputs/targets for training. + * Default implementation throws - only backends that support graphs override this. + * @param[in] graph The heterogeneous graph containing input features + * @param[in] outputs The graph fields containing output/target data + */ + virtual void store( + [[maybe_unused]] const ams::AMSHeterogeneousGraph& graph, + [[maybe_unused]] const ams::AMSHeterogeneousGraphFields& outputs) + { + THROW(std::runtime_error, + (this->type() + " database does not support heterogeneous graph " + "storage") + .c_str()); + } uint64_t getId() const { return id; } @@ -175,29 +211,29 @@ class FileDB : public BaseDB uint64_t rId) : BaseDB(rId) { - fs::path Path(path); + std::experimental::filesystem::path Path(path); std::error_code ec; - if (!fs::exists(Path, ec)) { + if (!std::experimental::filesystem::exists(Path, ec)) { std::cerr << "[ERROR]: Path:'" << path << "' does not exist\n"; exit(-1); } checkError(ec); - if (!fs::is_directory(Path, ec)) { + if (!std::experimental::filesystem::is_directory(Path, ec)) { std::cerr << "[ERROR]: Path:'" << path << "' is a file NOT a directory\n"; exit(-1); } - Path = fs::absolute(Path); + Path = std::experimental::filesystem::absolute(Path); fp = Path.string(); // We can now create the filename std::string dbfn(fn + "_"); dbfn += std::to_string(rId) + suffix; - Path /= fs::path(dbfn); - this->fn = fs::absolute(Path).string(); + Path /= std::experimental::filesystem::path(dbfn); + this->fn = std::experimental::filesystem::absolute(Path).string(); AMS_DBG(DB, "File System DB writes to file {}", this->fn) } @@ -1559,8 +1595,9 @@ class RMQInterface flush(100, 100); _publishingManager->stop(); auto size = MessagesBuffer::getInstance().size(); - if (size != 0) + if (size != 0) { AMS_DBG(RMQInterface, "Rank {} did not ack {} messages", _rId, size) + } } ~RMQInterface() @@ -1668,10 +1705,10 @@ class FilesystemInterface bool connect(std::string& path) { connected = true; - fs::path Path(path); + std::experimental::filesystem::path Path(path); std::error_code ec; - if (!fs::exists(Path, ec)) { + if (!std::experimental::filesystem::exists(Path, ec)) { THROW(std::runtime_error, ("Path: :'" + path + "' does not exist").c_str()); exit(-1); @@ -1760,37 +1797,10 @@ class DBManager * @param[in] rId a unique Id for each process taking part in a distributed * execution (rank-id) */ + // Declared here, implemented in basedb.cpp to avoid including jsondb.hpp in header std::shared_ptr 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. @@ -1904,10 +1914,10 @@ class DBManager std::string& routing_key, bool update_surrogate) { - fs::path Path(rmq_cert); + std::experimental::filesystem::path Path(rmq_cert); std::error_code ec; AMS_CWARNING(AMS, - !fs::exists(Path, ec), + !std::experimental::filesystem::exists(Path, ec), "Certificate file '{}' for RMQ server does not exist. AMS " "will " "try to connect without it.", 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/interface.cpp b/src/AMSlib/wf/interface.cpp index 71c474ea..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) // ============================================================================ @@ -675,16 +655,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 +664,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/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, diff --git a/src/AMSlib/wf/jsondb.cpp b/src/AMSlib/wf/jsondb.cpp new file mode 100644 index 00000000..0c003bd6 --- /dev/null +++ b/src/AMSlib/wf/jsondb.cpp @@ -0,0 +1,603 @@ +/* + * 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 +#include +#include + +#include "wf/debug.h" + +using namespace ams::db; +using namespace ams; + +// ---------------------------------------------------------------------- +// Helper functions +// ---------------------------------------------------------------------- + +namespace +{ + +namespace fs = std::experimental::filesystem; + +// 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 { + close(); + } catch (const std::exception& e) { + AMS_WARNING(JSONDB, + "Exception while automatically closing JSONDB: {}", + 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::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(); + + // AMSTensor device transfers are not implemented here yet. + if (location != AMSResourceType::AMS_HOST) { + AMS_WARNING(JSONDB, "GPU tensor serialization is not implemented"); + 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()); + result["byte_size"] = byte_size; + + // Add shape + auto shape_ref = tensor.shape(); + result["shape"] = std::vector(shape_ref.begin(), shape_ref.end()); + + 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}}; +} + +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(); + 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_" << getId() << "_" << 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 + tensors_json[name] = encodeBase64Tensor(Inputs[i]); + } + } + + // 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}}; + } else { // Pure json mode + tensors_json[name] = encodeBase64Tensor(Outputs[i]); + } + } + + 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_" << getId() << "_" << 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.elements() > 0) { + auto gf_shape = graph.global_features.shape(); + global_feature_dim = gf_shape[0]; + } + + // 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}}; + } else { // Pure json mode + tensors_json["node_features"] = encodeBase64Tensor(graph.node_features); + } + + // 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}}; + } else { // Pure json mode + tensors_json["edge_index"] = encodeBase64Tensor(graph.edge_index); + } + + // 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}}; + } else { // Pure json mode + tensors_json["edge_features"] = encodeBase64Tensor(graph.edge_features); + } + } + + // Write global_features + 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}}; + } else { // Pure json mode + tensors_json["global_features"] = + encodeBase64Tensor(graph.global_features); + } + } + + 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_++; + + AMS_DBG(JSONDB, + "Stored graph data for step {} ({} nodes, {} edges)", + case_dir, + num_nodes, + num_edges); +} + +void JSONDB::store(const ams::AMSHeterogeneousGraph&, + const ams::AMSHeterogeneousGraphFields&) +{ + // Heterogeneous graph storage not yet implemented + THROW(std::runtime_error, + "Heterogeneous graph storage not yet implemented in JSONDB"); +} + +void JSONDB::close() +{ + if (finalized_) { + AMS_DBG(JSONDB, "Manifest already finalized, skipping"); + return; + } + + // Build complete manifest + nlohmann::json manifest; + manifest["format_version"] = 1; + manifest["endianness"] = isLittleEndian() ? "little" : "big"; + + // 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_; + + // 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, + ("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..7b54bc83 --- /dev/null +++ b/src/AMSlib/wf/jsondb.hpp @@ -0,0 +1,214 @@ +/* + * 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 "AMSGraph.hpp" +#include "AMSTensor.hpp" +#include "wf/basedb.hpp" + +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 and homogeneous graphs. Heterogeneous graph storage is not yet + * supported. + * + * Supports two modes: + * - "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 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 +{ +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 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 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] + * @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; + +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() override; + + // 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 the JSON manifest + */ + void close() override; + + /** + * @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..730f459f 100644 --- a/src/AMSlib/wf/workflow.hpp +++ b/src/AMSlib/wf/workflow.hpp @@ -102,21 +102,39 @@ 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 +462,55 @@ 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 + 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 + storeGraphData(graph_input, outputs); + + CALIPER(CALI_MARK_END("AMSEvaluateGraph");) + } + std::string getDBName() { if (!DB) return ""; diff --git a/tests/AMSlib/ams_interface/CMakeLists.txt b/tests/AMSlib/ams_interface/CMakeLists.txt index 10b5a25e..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}) @@ -44,6 +53,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_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 new file mode 100644 index 00000000..1a626584 --- /dev/null +++ b/tests/AMSlib/ams_interface/test_graph_workflow_storage.cpp @@ -0,0 +1,1337 @@ +/* + * 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 +#include +#include +#include +#include +#include +#include + +#include "AMS.h" +#include "AMSGraph.hpp" +#include "AMSTensor.hpp" +#include "nlohmann/json.hpp" +#include "wf/jsondb.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); +} + +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 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, + 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 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)) { + 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 stores every named output in binary mode", + "[wf][graph][storage]") +{ + // Use unique temp directory + fs::path test_dir = + fs::temp_directory_path() / "ams_graph_named_output_storage_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("generic_graph_schema", -1.0, "", true); + AMSExecutor executor = AMSCreateExecutor(recorder, 0, 1); + const fs::path manifest_path = AMSGetDatabaseName(executor); + CATCH_REQUIRE(manifest_path.filename() == + "generic_graph_schema_0_jsondb.json"); + + // 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 [2] - float32 + auto global_feat = makeTensor({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 with multiple application-defined output fields. + int callback_count = 0; + HomogeneousGraphDomainFn physics = [&](const AMSHomogeneousGraph& g, + AMSHomogeneousGraphFields& o) { + callback_count++; + + const int64_t num_nodes = g.node_features.shape()[0]; + 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++) { + temperature_data[i] = static_cast(i) * 0.123; + } + o.node_fields.insert("temperature", std::move(temperature)); + + 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 + AMSExecute(executor, physics, graph, outputs); + + // Verify physics ran + CATCH_REQUIRE(callback_count == 1); + 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); + + CATCH_REQUIRE(fs::exists(manifest_path)); + + std::ifstream manifest_file(manifest_path); + nlohmann::json manifest; + manifest_file >> manifest; + + // Verify essential structure exists + CATCH_REQUIRE(manifest.contains("format_version")); + 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 + 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")); + 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)); + + 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() == + "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"); + + 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); + + 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_0_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_0_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_0_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_0_000000/outputs/global/" + "field_000000.bin")); + + // 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)); + + fs::remove_all(test_dir); +} + +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); + const fs::path manifest_path = AMSGetDatabaseName(executor); + + auto node_feat = makeTensor({3, 2}); + 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)); + CATCH_REQUIRE(graph.global_features.shape().size() == 1); + CATCH_REQUIRE(graph.global_features.shape()[0] == 0); + + AMSHomogeneousGraphFields outputs; + HomogeneousGraphDomainFn physics = [](const AMSHomogeneousGraph&, + AMSHomogeneousGraphFields&) {}; + + AMSExecute(executor, physics, graph, outputs); + AMSDestroyExecutor(executor); + + std::ifstream manifest_file(manifest_path); + 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_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()); + CATCH_REQUIRE_FALSE(stored_case.contains("target_dim")); + + 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["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(), + 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("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]") +{ + 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 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) + 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); + 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); + 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(), + 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& 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); +} + +// ============================================================================ +// 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({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 database artifacts were created (store_data=false). + CATCH_REQUIRE(fs::is_empty(test_dir)); + + 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 fs::path manifest_path = AMSGetDatabaseName(executor); + + 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}); + + // 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(manifest_path)); + + std::ifstream manifest_file(manifest_path); + 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 output exists + CATCH_REQUIRE(case_entry["outputs"]["node"].contains("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); + const fs::path manifest_path = AMSGetDatabaseName(executor); + + // 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&, + 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") != 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 + + // Verify heterogeneous storage + CATCH_REQUIRE(fs::exists(manifest_path)); + + std::ifstream manifest_file(manifest_path); + 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}); + + // 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({g.node_features.shape()[0], 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}); + + // 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); + const fs::path manifest_path = AMSGetDatabaseName(executor); + + auto node_feat = makeTensor({6, 2}); + auto edge_idx = makeTensor({2, 5}); + auto edge_feat = makeTensor({5, 1}); + auto global_feat = makeTensor({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(manifest_path)); + + std::ifstream manifest_file(manifest_path); + 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["outputs"]["node"]["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 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["outputs"]["node"]["delta_u"]["dtype"] == "float64"); + + fs::remove_all(test_dir); +} 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