Skip to content

[Solvergraph] Add callback API for node tracking #1890

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
tdavidcl wants to merge 2 commits into
base: main
Choose a base branch
from
Draft

[Solvergraph] Add callback API for node tracking #1890

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
3ac17ab
basic node tracking
tdavidcl Jun 23, 2026
cc8e495
basic node tracking
tdavidcl Jun 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
146 changes: 146 additions & 0 deletions src/shammodels/ramses/src/Solver.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -269,9 +269,155 @@ class PatchDataLayerToVtk : public shamrock::solvergraph::INode {
std::string _impl_get_tex() { return "TODO"; }
};

class GraphTracer {
public:
std::string filename;

std::queue<std::string> events;

public:
inline void add_event(std::string event) { events.push(event); }

inline void maybe_flush() {

// open the file for writing
std::ofstream file(filename, std::ios::app);

// lock the file with flock

// flush the events
while (!events.empty()) {
std::string event = events.front();
events.pop();
file << event << std::endl;
}

// close the file
file.close();
}
Comment on lines +281 to +297

@gemini-code-assist gemini-code-assist Bot Jun 23, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Opening and closing the file stream on every single event flush (std::ofstream file(filename, std::ios::app)) introduces a significant performance bottleneck, especially in high-performance simulation codes. Consider keeping the file stream open as a member of GraphTracer (opening it once during initialization and closing it in the destructor), and calling file.flush() or letting the buffer flush naturally to minimize disk I/O overhead.

};

GraphTracer graph_tracer("graph_trace.txt");
Comment on lines +272 to +300

@gemini-code-assist gemini-code-assist Bot Jun 23, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Issue: Missing Constructor, MPI Race Condition, and Static Initialization Order Fiasco

  1. Missing Constructor: GraphTracer does not define a constructor taking a std::string, which will cause a compilation error when initializing graph_tracer("graph_trace.txt") on compilers that do not support parenthesized aggregate initialization (pre-C++20).
  2. MPI Race Condition: In an MPI-enabled run, every rank will attempt to write to the same file (graph_trace.txt) concurrently, leading to file corruption or interleaved output.
  3. Static Initialization Order: graph_tracer is initialized as a global variable. Since it may call MPI-related functions or be accessed during static destruction, this can lead to undefined behavior if accessed before MPI is initialized or after global destruction.

Solution

  • Add an explicit constructor and destructor to GraphTracer to ensure safe initialization and automatic flushing of remaining events on destruction.
  • Use a function-local static get_graph_tracer() to lazily initialize the tracer after MPI has been initialized, appending the MPI rank to the filename to avoid conflicts.
  • Update all other occurrences of graph_tracer in the callback functions to use get_graph_tracer().
class GraphTracer {
public:
    std::string filename;
    std::queue<std::string> events;

    explicit GraphTracer(std::string filename) : filename(std::move(filename)) {}

    ~GraphTracer() {
        maybe_flush();
    }

    inline void add_event(std::string event) { events.push(event); }

    inline void maybe_flush() {
        std::ofstream file(filename, std::ios::app);
        if (!file) {
            return;
        }

        while (!events.empty()) {
            std::string event = events.front();
            events.pop();
            file << event << std::endl;
        }
        file.close();
    }
};

inline GraphTracer& get_graph_tracer() {
    static GraphTracer tracer("graph_trace_" + std::to_string(shamcomm::world_rank()) + ".txt");
    return tracer;
}


void on_node_create(u64 uuid) {
graph_tracer.add_event(
nlohmann::json{
{"event", "node_create"},
{"uuid", uuid},
{"wtime", shambase::details::get_wtime()},
}
.dump());
graph_tracer.maybe_flush();
}

void on_node_destroy(u64 uuid) {
graph_tracer.add_event(
nlohmann::json{
{"event", "node_destroy"},
{"uuid", uuid},
{"wtime", shambase::details::get_wtime()},
}
.dump());
graph_tracer.maybe_flush();
}

void on_edge_create(u64 uuid) {
graph_tracer.add_event(
nlohmann::json{
{"event", "edge_create"},
{"uuid", uuid},
{"wtime", shambase::details::get_wtime()},
}
.dump());
}

void on_edge_destroy(u64 uuid) {
graph_tracer.add_event(
nlohmann::json{
{"event", "edge_destroy"},
{"uuid", uuid},
{"wtime", shambase::details::get_wtime()},
}
.dump());
}
Comment on lines +324 to +342

@gemini-code-assist gemini-code-assist Bot Jun 23, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The callbacks on_edge_create and on_edge_destroy share duplicated logic for constructing, dumping, and queuing JSON events, and both lack a call to maybe_flush(). To improve readability and maintainability, refactor this duplicated logic into a helper function or lambda that also ensures maybe_flush() is called.

auto log_event = [&](const std::string& event, u64 uuid) {
    graph_tracer.add_event(
        nlohmann::json{
            {"event", event},
            {"uuid", uuid},
            {"wtime", shambase::details::get_wtime()},
        }
            .dump());
    graph_tracer.maybe_flush();
};

void on_edge_create(u64 uuid) {
    log_event("edge_create", uuid);
}

void on_edge_destroy(u64 uuid) {
    log_event("edge_destroy", uuid);
}
References
  1. Refactor duplicated logic into a helper function or lambda to improve readability and maintainability.


void on_node_update(shamrock::solvergraph::INode &node) {

nlohmann::json edge_list = nlohmann::json::array();

node.on_edge_ro_edges([&](shamrock::solvergraph::IEdge &e) {
edge_list.push_back(
nlohmann::json{{"type", "read"}, {"label", e.get_label()}, {"uuid", e.get_uuid()}});
});

node.on_edge_rw_edges([&](shamrock::solvergraph::IEdge &e) {
edge_list.push_back(
nlohmann::json{
{"type", "read_write"}, {"label", e.get_label()}, {"uuid", e.get_uuid()}});
});

graph_tracer.add_event(
nlohmann::json{
{"event", "node_update"},
{"uuid", node.get_uuid()},
{"type", typeid(node).name()},
{"label", node.get_label()},
{"wtime", shambase::details::get_wtime()},
{"edges", edge_list}}
.dump());
graph_tracer.maybe_flush();
}

void on_node_op(u64 uuid, u64 op_id) {
if (op_id == 0) {
graph_tracer.add_event(
nlohmann::json{
{"event", "node_evaluate_begin"},
{"uuid", uuid},
{"wtime", shambase::details::get_wtime()},
}
.dump());
} else if (op_id == 1) {
graph_tracer.add_event(
nlohmann::json{
{"event", "node_evaluate_end"},
{"uuid", uuid},
{"wtime", shambase::details::get_wtime()},
}
.dump());
} else {
shambase::throw_with_loc<std::runtime_error>(shambase::format("Invalid op_id: {}", op_id));
}
graph_tracer.maybe_flush();
}

void on_node_evaluate_end(u64 uuid) {
graph_tracer.add_event(
nlohmann::json{
{"event", "node_evaluate_end"},
{"uuid", uuid},
{"wtime", shambase::details::get_wtime()},
}
.dump());
graph_tracer.maybe_flush();
}
Comment on lines +394 to +403

@gemini-code-assist gemini-code-assist Bot Jun 23, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Redundant Function

on_node_evaluate_end is redundant because on_node_op already handles op_id == 1 as "node_evaluate_end". Additionally, on_node_evaluate_end is never registered as a callback in init_solver_graph(). Removing it keeps the codebase clean and maintainable.

Comment on lines +394 to +403

@gemini-code-assist gemini-code-assist Bot Jun 23, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The function on_node_evaluate_end is defined but never registered as a callback in init_solver_graph(). Instead, on_node_op handles the evaluation end event (when op_id == 1). This function is dead code and should be removed to keep the codebase clean.


template<class Tvec, class TgridVec>
void shammodels::basegodunov::Solver<Tvec, TgridVec>::init_solver_graph() {

Comment on lines 406 to 407

@gemini-code-assist gemini-code-assist Bot Jun 23, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

In a parallel MPI execution, multiple ranks will run concurrently. If all ranks write to the same hardcoded file "graph_trace.txt" without synchronization, the file writes will interleave, leading to corrupted and unparseable JSON data. To prevent this, the filename should be rank-specific, for example by appending shamcomm::world_rank(). Since shamcomm::world_rank() might not be initialized during global static construction, you should initialize or update the filename dynamically inside init_solver_graph().

template<class Tvec, class TgridVec>
void shammodels::basegodunov::Solver<Tvec, TgridVec>::init_solver_graph() {

    graph_tracer.filename = "graph_trace_" + std::to_string(shamcomm::world_rank()) + ".txt";

shamrock::solvergraph::LifetimeTracker<shamrock::solvergraph::INode>::on_create
= on_node_create;
shamrock::solvergraph::LifetimeTracker<shamrock::solvergraph::INode>::on_destroy
= on_node_destroy;
shamrock::solvergraph::LifetimeTracker<shamrock::solvergraph::INode>::on_state_update
= on_node_update;
shamrock::solvergraph::LifetimeTracker<shamrock::solvergraph::INode>::on_op = on_node_op;

shamrock::solvergraph::LifetimeTracker<shamrock::solvergraph::IEdge>::on_create
= on_edge_create;
shamrock::solvergraph::LifetimeTracker<shamrock::solvergraph::IEdge>::on_destroy
= on_edge_destroy;

bool enable_mem_free = false;

auto get_optional_free_mem = [&](auto &bind_to, auto &add_to) {
Expand Down
14 changes: 13 additions & 1 deletion src/shamrock/include/shamrock/solvergraph/IEdge.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,21 +18,33 @@

#include "shambase/WithUUID.hpp"
#include "shambase/aliases_int.hpp"
#include "shambase/memory.hpp"
#include "shamrock/solvergraph/IFreeable.hpp"
#include "shamrock/solvergraph/LifetimeTracker.hpp"
#include <string>

namespace shamrock::solvergraph {

class INode;

class IEdge : public shambase::WithUUID<IEdge, u64>, public IFreeable {
class IEdge : public IFreeable {

/// track the UUID of the node to notify changes.
/// Sadly we can not use a WithUUID directly since it is hard to log destruction in the
/// destructor if the object was moved which could lead to double free notifications.
std::shared_ptr<LifetimeTracker<IEdge>> tracker
= std::make_shared<LifetimeTracker<IEdge>>();

public:
inline std::string get_label() const { return _impl_get_dot_label(); }
inline std::string get_tex_symbol() const { return _impl_get_tex_symbol(); }

virtual std::string _impl_get_dot_label() const = 0;
virtual std::string _impl_get_tex_symbol() const = 0;

/// Get the UUID of the node

@gemini-code-assist gemini-code-assist Bot Jun 23, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The comment says "Get the UUID of the node", but this is the IEdge class. It should say "Get the UUID of the edge".

Suggested change
/// Get the UUID of the node
/// Get the UUID of the edge

inline u64 get_uuid() const { return shambase::get_check_ref(tracker).get_uuid(); }

inline virtual ~IEdge() {}
};

Expand Down
37 changes: 32 additions & 5 deletions src/shamrock/include/shamrock/solvergraph/INode.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,19 +20,25 @@
#include "shambase/memory.hpp"
#include "shambase/stacktrace.hpp"
#include "shamrock/solvergraph/IEdge.hpp"
#include "shamrock/solvergraph/LifetimeTracker.hpp"
#include <memory>
#include <vector>

namespace shamrock::solvergraph {

/// Inode is node between data edges, takes multiple inputs, multiple outputs
class INode : public std::enable_shared_from_this<INode>,
public shambase::WithUUID<INode, u64> {
class INode : public std::enable_shared_from_this<INode> {

/// Read only edges
std::vector<std::shared_ptr<IEdge>> ro_edges;
std::vector<std::shared_ptr<IEdge>> ro_edges = {};
/// Read write edges
std::vector<std::shared_ptr<IEdge>> rw_edges;
std::vector<std::shared_ptr<IEdge>> rw_edges = {};

/// track the UUID of the node to notify changes.
/// Sadly we can not use a WithUUID directly since it is hard to log destruction in the
/// destructor if the object was moved which could lead to double free notifications.
std::shared_ptr<LifetimeTracker<INode>> tracker
= std::make_shared<LifetimeTracker<INode>>();

public:
INode() = default;
Expand All @@ -46,6 +52,9 @@ namespace shamrock::solvergraph {
/// Move assignment - automatically delegates to base classes and members
INode &operator=(INode &&) noexcept = default;

/// Get the UUID of the node
inline u64 get_uuid() const { return shambase::get_check_ref(tracker).get_uuid(); }

/// Get a shared pointer to this node
inline std::shared_ptr<INode> getptr_shared() { return shared_from_this(); }
/// Get a weak pointer to this node
Expand All @@ -71,6 +80,7 @@ namespace shamrock::solvergraph {

/// Destructor (virtual) & reset the edges
virtual ~INode() {
tracker.reset();
__internal_set_ro_edges({});
__internal_set_rw_edges({});
}
Expand Down Expand Up @@ -106,7 +116,15 @@ namespace shamrock::solvergraph {
}

/// Evaluate the node
inline void evaluate() { _impl_evaluate_internal(); }
inline void evaluate() {
if (tracker) {
tracker->notify_op(0);
}
_impl_evaluate_internal();
if (tracker) {
tracker->notify_op(1);
}
}

/// Get the dot graph of the node (Currently only an alias to get_dot_graph_partial)
inline std::string get_dot_graph() { return get_dot_graph_partial(); };
Expand All @@ -124,6 +142,9 @@ namespace shamrock::solvergraph {
/// Get the TeX of the node partial
inline std::string get_tex_partial() { return _impl_get_tex(); };

/// Get the label of the node
inline std::string get_label() { return _impl_get_label(); };

/// print the node info
inline virtual std::string print_node_info() const {
std::string node_info = shambase::format("Node info :\n");
Expand Down Expand Up @@ -176,6 +197,9 @@ namespace shamrock::solvergraph {
for (auto e : ro_edges) {
// shambase::get_check_ref(e).parent = getptr_weak();
}
if (tracker) {
tracker->notify_update(*this);
}
}

inline void INode::__internal_set_rw_edges(std::vector<std::shared_ptr<IEdge>> new_rw_edges) {
Expand All @@ -186,6 +210,9 @@ namespace shamrock::solvergraph {
for (auto e : rw_edges) {
// shambase::get_check_ref(e).child = getptr_weak();
}
if (tracker) {
tracker->notify_update(*this);
}
}

template<class Func>
Expand Down
63 changes: 63 additions & 0 deletions src/shamrock/include/shamrock/solvergraph/LifetimeTracker.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,63 @@
// -------------------------------------------------------//
//
// SHAMROCK code for hydrodynamics
// Copyright (c) 2021-2026 Timothée David--Cléris <tim.shamrock@proton.me>
// SPDX-License-Identifier: CeCILL Free Software License Agreement v2.1
// Shamrock is licensed under the CeCILL 2.1 License, see LICENSE for more information
//
// -------------------------------------------------------//

#pragma once

/**
* @file LifetimeTracker.hpp
* @author Timothée David--Cléris (tim.shamrock@proton.me)
* @brief
*
*/

#include "shambase/WithUUID.hpp"
#include "shambase/aliases_int.hpp"
namespace shamrock::solvergraph {

template<typename T>
class LifetimeTracker : public shambase::WithUUID<LifetimeTracker<T>, u64> {
public:
inline static void (*on_create)(u64 uuid) = nullptr;
inline static void (*on_destroy)(u64 uuid) = nullptr;

inline static void (*on_state_update)(T &node) = nullptr;
inline static void (*on_op)(u64 uuid, u64 op_id) = nullptr;

LifetimeTracker() : shambase::WithUUID<LifetimeTracker, u64>() {
if (on_create != nullptr) {
on_create(this->get_uuid());
}
};

LifetimeTracker(const LifetimeTracker &) = delete;
LifetimeTracker &operator=(const LifetimeTracker &) = delete;

LifetimeTracker(LifetimeTracker &&) noexcept = default;
LifetimeTracker &operator=(LifetimeTracker &&) noexcept = default;
Comment on lines +41 to +42

@gemini-code-assist gemini-code-assist Bot Jun 23, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Potential Double-Destroy Bug on Move

If a LifetimeTracker is moved, the default move constructor/assignment will copy the uuid member. When the temporary/source LifetimeTracker is destroyed, its destructor will trigger on_destroy with the copied uuid. When the moved-to object is eventually destroyed, it will trigger on_destroy again with the same uuid, leading to a double-destroy notification.

Since LifetimeTracker is managed via std::shared_ptr and should never be moved directly, the move constructor and move assignment operator should be explicitly deleted to prevent accidental moves.

Suggested change
LifetimeTracker(LifetimeTracker &&) noexcept = default;
LifetimeTracker &operator=(LifetimeTracker &&) noexcept = default;
LifetimeTracker(LifetimeTracker &&) = delete;
LifetimeTracker &operator=(LifetimeTracker &&) = delete;

Comment on lines +41 to +42

@gemini-code-assist gemini-code-assist Bot Jun 23, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Since LifetimeTracker triggers on_destroy in its destructor using its uuid, allowing it to be moved can lead to double-destroy notifications (as both the moved-from temporary and the moved-to destination will share the same uuid and trigger on_destroy upon destruction). Since LifetimeTracker is managed via std::shared_ptr and should never be moved directly, you should explicitly delete its move constructor and move assignment operator to prevent accidental moves.

Suggested change
LifetimeTracker(LifetimeTracker &&) noexcept = default;
LifetimeTracker &operator=(LifetimeTracker &&) noexcept = default;
LifetimeTracker(LifetimeTracker &&) = delete;
LifetimeTracker &operator=(LifetimeTracker &&) = delete;


inline void notify_update(T &node) {
if (on_state_update != nullptr) {
on_state_update(node);
}
}

inline void notify_op(u64 op_id) {
if (on_op != nullptr) {
on_op(this->get_uuid(), op_id);
}
}

~LifetimeTracker() {
if (on_destroy != nullptr) {
on_destroy(this->get_uuid());
}
};
};

} // namespace shamrock::solvergraph
Loading