Skip to content

Trace ID to link before and after events#589

Merged
fonsp merged 5 commits into
release-6from
trace-ID-v1
Mar 26, 2026
Merged

Trace ID to link before and after events#589
fonsp merged 5 commits into
release-6from
trace-ID-v1

Conversation

@fonsp
Copy link
Copy Markdown
Contributor

@fonsp fonsp commented Mar 24, 2026

I think it is useful to link before and after events. Here is one implementation, using a field trace_id added to events. This also needs to be done to the events defined in RxInfer.jl.

It works, but maybe we can do better?

Alternative 1: before field in AfterSomething events

We can add a field .before that stores the counterpart of an AfterSomething event. Instead of trace_id. (Then the 'before' objectID functions the trace ID.)

Aternative 2: Before{...} and After{...} types

We could change the types to eg:

Before{ProductOfMessagesEvent{V, C, Ms, R}} <: After{Event{:product_of_messages}}

The After type can have a .result field.

Then we can replace this:

# Before
trace_id = uuid4()
invoke_callback(
    context.callbacks,
    BeforeProductOfMessagesEvent(variable, context, messages, trace_id),
)

result = as_message(...)

invoke_callback(
    context.callbacks,
    AfterProductOfMessagesEvent(variable, context, messages, trace_id),
)

with:

# After
result = with_span(
	BeforeProductOfMessagesEvent(variable, context, messages, trace_id)
) do
	as_message(...)
end

@github-actions
Copy link
Copy Markdown
Contributor

🤖 Code Formatting

Your PR still has some code formatting issues. I've updated PR #590 with the necessary formatting changes.

You can merge that PR into this branch to fix the code style check.

Alternatively, you can run make format locally and push the changes yourself.

@codecov
Copy link
Copy Markdown

codecov Bot commented Mar 24, 2026

Codecov Report

✅ All modified and coverable lines are covered by tests.
⚠️ Please upload report for BASE (release-6@a40573e). Learn more about missing BASE report.

Additional details and impacted files
@@             Coverage Diff              @@
##             release-6     #589   +/-   ##
============================================
  Coverage             ?   80.74%           
============================================
  Files                ?      212           
  Lines                ?     6462           
  Branches             ?        0           
============================================
  Hits                 ?     5218           
  Misses               ?     1244           
  Partials             ?        0           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@bvdmitri
Copy link
Copy Markdown
Member

I think storing the .before field might be interesting indeed, the invoke_callback returns the event, so the code should be a simple as

before_event = invoke_callback(callbacks, BeforeEvent(...))

...

after_event = invoke_callback(callbacks, AfterEvent(..., before = before_event)))

@bvdmitri
Copy link
Copy Markdown
Member

bvdmitri commented Mar 25, 2026

I would strongly not suggest doing the with_span idea. I experimented with that a bit, the problem is with Julia's lambda implementation and that it needs to capture variables outside. That hits performance really bad when those variables are not fully inferred (and they are not fully inferred in our case, because the result of the product/rule call/form constrain application is often type unstable).

Also we discussed that a bit in the corresponding RxInferBoard discussion and reached the conclusion that the tracing system should be dead-simple linear and should not have a notion of a span by itself. Spans are OpenTelemetry related and do not exist in TensorBoard for example. So ReactiveMP should not inherently create those spans but it should make it easy to recover those spans (e.g. with extra fields like .trace_id or .before)

@bvdmitri
Copy link
Copy Markdown
Member

The current solution is not bad either btw, so I don't have a strong preference of this implementation over the .before field

@bvdmitri
Copy link
Copy Markdown
Member

I only want to make sure that this extra field does not have performance implications in a sense that if there are no callbacks, Julia compile should be able to remove both callsites. Can you write a test for that to make sure that is what is happening (with the @allocated macro)

@github-actions
Copy link
Copy Markdown
Contributor

🤖 Code Formatting

Your PR still has some code formatting issues. I've updated PR #590 with the necessary formatting changes.

You can merge that PR into this branch to fix the code style check.

Alternatively, you can run make format locally and push the changes yourself.

@fonsp fonsp marked this pull request as ready for review March 26, 2026 10:29
Copy link
Copy Markdown
Member

@wouterwln wouterwln left a comment

Choose a reason for hiding this comment

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

Super nice, I like this a lot

@bvdmitri
Copy link
Copy Markdown
Member

It seems that Julia v1.10 isn't that smart :( I'm ok with just guarding the test with @static VERSION check and only check on v1.11+

@bvdmitri
Copy link
Copy Markdown
Member

OR - it might also happening because sqrt can throw, an maybe Julia v1.10 isn't good enough to understand that and shows allocations for potential error path. Its better to call functions that truly do not allocate, e.g. +

@github-actions
Copy link
Copy Markdown
Contributor

🤖 Code Formatting

Your PR still has some code formatting issues. I've updated PR #590 with the necessary formatting changes.

You can merge that PR into this branch to fix the code style check.

Alternatively, you can run make format locally and push the changes yourself.

@fonsp fonsp marked this pull request as draft March 26, 2026 11:54
@github-actions
Copy link
Copy Markdown
Contributor

🤖 Code Formatting

Your PR still has some code formatting issues. I've updated PR #590 with the necessary formatting changes.

You can merge that PR into this branch to fix the code style check.

Alternatively, you can run make format locally and push the changes yourself.

Copy link
Copy Markdown
Member

@bvdmitri bvdmitri left a comment

Choose a reason for hiding this comment

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

lets gooo!

@fonsp
Copy link
Copy Markdown
Contributor Author

fonsp commented Mar 26, 2026

I'll work through the CI Julia 1.10 and then I merge 👍

@fonsp fonsp marked this pull request as ready for review March 26, 2026 12:24
@fonsp fonsp merged commit d8459de into release-6 Mar 26, 2026
3 of 5 checks passed
@fonsp fonsp deleted the trace-ID-v1 branch March 26, 2026 12:24
@fonsp
Copy link
Copy Markdown
Contributor Author

fonsp commented Mar 26, 2026

For future reference, here is a function that matches traces based on their ID:

function find_spans(trace_events::Vector{TracedEvent})
	# Found spans
	spans = Tuple{TracedEvent,TracedEvent}[]
	# Currently waiting to be matched
	lonely = OrderedDict{Base.UUID, Any}()
	# No match possible
	not_matched = TracedEvent[]

	for te in trace_events
		event = te.event
		if hasfield(typeof(event), :trace_id)
			tid = event.trace_id
			friend = get(lonely, tid, nothing)
			if !isnothing(friend)
				delete!(lonely, tid)
				push!(spans, (friend, te))
			else
				lonely[tid] = te
			end
		else
			push!(not_matched, te)
		end
	end

	push!(not_matched, values(lonely)...)

	return (; spans, not_matched )
end

bvdmitri added a commit that referenced this pull request Apr 17, 2026
* add Changelog

* Implement event handler system for message passing and callbacks (#587)

* feat: add event handler system for message passing procedure

Introduce a new event handler mechanism that allows users to hook into
the message passing procedure via `Event`, `handle_event`, and
`broadcast_event` functions. Events are broadcast before and after
message rule calls, enabling debugging and monitoring capabilities.

Key changes:
- Add `event_handler.jl` with event types and handler interfaces
- Extend `MessageMapping` to include an event handler parameter
- Support both custom handler types and NamedTuple-based handlers
- Add documentation for the events system
- Include comprehensive tests for event handler functionality

* type in docstring

* inject the event handler in the factor node activation options

* make format

* document the existing events better

* rename event handler to callbacks to better match RxInfer

* update docs

* use make format

* fix typo

* allow to merge callbacks

* allow reduce the result of the callbacks

* add per-event callback reducer

* make format

* remove test method

* start reimplement of the product of two messages

this is required to inject the callbacks properly, plus we do the breaking
release already. Could break (and improve) more things as well then

* fix initial integration with RxInfer.jl

* fix Aqua tests

* caught a small bug in RxInfer tests

* caught another bug from RxInfer.jl tests

* merge stricter formatting

* Refactor the variables, add docstrings, add labels

- `compute_product_of_messages` now accepts the `AbstractVariable`,
  makes it easier to identify the variable inside the callback

* fix documentation build

* 2prev

* add before/after product of two messages callbacks

* fix failing tests

* add before/after product of messages callbacks. add before/after form constraints

* add before/after marginal compute callbacks

* add error hint in case of wrong passed callbacks

* fix test

* temporary fix for the logscale switch rule

* support Dict as callback handler

* fix warnings in the tests

* make format

* refactor to use an abstract Event{E} structure

* add new method for event_name

* use mutable fields instead

* update CHANGELOG

* Trace ID to link before and after events (#589)

* Trace ID to link before and after events

* Add allocations test

* disable some allocation tests on old julia versions

* Change sqrt to + and remove version check

* only check on julia 1.12 again

* Better trace ids (#595)

* create IDs only if callbacks are set

* fix one failing test

* make format

* Refactor addon system to annotations (#596)

* save plan

* update

* refine the plan

* refine plan 2

* implement post_product/post_rule annotations infrastructure

* reimplement the memory addon, remove old implementation

* commit new file

* make format

* update plan

* refactor the message.jl and message_tests.jl

* update marginals.jl and marginal_tests.jl

* remove old bridge functions

* refactor the rule macro and the rule generating expression

* refactor the nodes and tests

* fix remaining places where addons were used

* fix a lot of tests by hand

* fix more tests by hand

* update the plan

* finish the refactor, update the documentation with migration guide,
update CHANGELOG

* fix from RxInfer

* fix bugs caught in RxInfer

* fix documentation build

* add pretty output for input argument records

* 🤖 Auto-format Julia code (#597)

Co-authored-by: bvdmitri <6557701+bvdmitri@users.noreply.github.com>

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: bvdmitri <6557701+bvdmitri@users.noreply.github.com>

* checkpoint

* checkpoint

* checkpoint

* make format

* fix some tests by hand

* fix more tests by hand

* local update

* fix more tests

* rename vague messagein/messageout

* make format

* major documentation upgrade

* update the CHANGELOG and migration guide

* structural documentation update, add concepts

* improve documentation for existing pages about specific nodes

* update helpers documentation

* add concepts pages

* finish the documentation overhaul

* update the docs

* update the version to 6.0.0

* rename pipeline to stream postprocessors

* update the scheduler to use stream postprocessors instead

* add postprocess marginal stream

* simplify API even further

* update docs

* make format

* fix typo in mixture nodes

* update

* add pre_rule_annotations

* add logo, update README.md

* make logo half the size in README

---------

Co-authored-by: Fons van der Plas <fonsvdplas@gmail.com>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: bvdmitri <6557701+bvdmitri@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants