KEMBAR78
Add function tracking decorator and update IntermediateStep by dnandakumar-nv · Pull Request #98 · NVIDIA/NeMo-Agent-Toolkit · GitHub
Skip to content

Conversation

@dnandakumar-nv
Copy link
Contributor

@dnandakumar-nv dnandakumar-nv commented Apr 8, 2025

Introduce a track_function decorator to wrap sync, async, and generator functions for tracking execution spans. Added new SPAN_START, SPAN_CHUNK, and SPAN_END types to IntermediateStep. Updated tests to ensure tracking functionality is properly validated.

Description

Feature: Function Tracking Decorator

This PR adds a straightforward decorator, @track_function, that makes it easy to log details about function calls. It works on synchronous and asynchronous functions, generators (sync and async), and class methods—all with a consistent interface.

Key Points:

  • Versatile Support: Works with regular functions, async functions, generators, and even class methods.
  • Optional Metadata: You can add extra context by passing a dictionary to the decorator. Metadata keys are validated to ensure they’re strings.
  • Data Serialization: Automatically serializes input arguments and outputs (or each yield) into a JSON-friendly format. It even handles Pydantic models by calling model_dump() if needed.
  • Intermediate Event Logging: The decorator sends tracking events to a reactive stream in three stages:
    -- SPAN_START: When the function starts.
    -- SPAN_CHUNK: For each yielded item (for generators).
    -- SPAN_END: When the function returns (or ends in generators).

Usage Example:

from aiq.profiler.decorators.function_tracking import track_function

@track_function(metadata={"action": "compute", "source": "api"})
def add_numbers(a, b):
    return a + b

Just apply @track_function to any function (or method) where you need a bit of extra insight into the inputs and outputs. Subscribe to the reactive stream to see the published events and get a better sense of the function’s flow.

By Submitting this PR I confirm:

  • I am familiar with the Contributing Guidelines.
  • We require that all contributors "sign-off" on their commits. This certifies that the contribution is your original work, or you have rights to submit it under the same license, or a compatible license.
    • Any contribution which contains commits that are not Signed-Off will not be accepted.
  • When the PR is ready for review, new or existing tests cover these changes.
  • When the PR is ready for review, the documentation is up to date with these changes.

Introduce a `track_function` decorator to wrap sync, async, and generator functions for tracking execution spans. Added new SPAN_START, SPAN_CHUNK, and SPAN_END types to IntermediateStep. Updated tests to ensure tracking functionality is properly validated.

Signed-off-by: dnandakumar-nv <168006707+dnandakumar-nv@users.noreply.github.com>
@dnandakumar-nv dnandakumar-nv requested a review from a team as a code owner April 8, 2025 00:49
@copy-pr-bot
Copy link

copy-pr-bot bot commented Apr 8, 2025

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@dnandakumar-nv dnandakumar-nv added enhancement non-breaking Non-breaking change improvement Improvement to existing functionality and removed enhancement labels Apr 8, 2025
@dnandakumar-nv
Copy link
Contributor Author

/ok to test

This update adds the SPDX-based copyright and license headers to comply with legal and attribution requirements. The changes ensure proper licensing information is included for the test file.

Signed-off-by: dnandakumar-nv <168006707+dnandakumar-nv@users.noreply.github.com>
Document how to use the `@track_function` decorator in the profiler, outlining its stages, compatibility, and key benefits. This update helps users profile custom functions with ease, providing metadata, real-time event streaming, and automatic data serialization.

Signed-off-by: dnandakumar-nv <168006707+dnandakumar-nv@users.noreply.github.com>
Replaced the SpanPayload class with TraceMetadata within IntermediateStepPayload to standardize metadata handling. Updated the function tracking logic and corresponding tests to reflect this change, improving both clarity and maintainability.

Signed-off-by: dnandakumar-nv <168006707+dnandakumar-nv@users.noreply.github.com>
This commit introduces the `function_name` parameter to intermediate step tracking, ensuring the function's name is included in logs and metadata. This enhances traceability and debugging by providing more context about the tracked steps.

Signed-off-by: dnandakumar-nv <168006707+dnandakumar-nv@users.noreply.github.com>
@dnandakumar-nv
Copy link
Contributor Author

/ok to test

Simplified and condensed docstrings to enhance readability and maintain focus on key functionality. Removed verbose examples and redundant details while retaining essential information about function behavior.

Signed-off-by: dnandakumar-nv <168006707+dnandakumar-nv@users.noreply.github.com>
@dnandakumar-nv dnandakumar-nv self-assigned this Apr 8, 2025
@dnandakumar-nv
Copy link
Contributor Author

/ok to test 27890f2

Signed-off-by: dnandakumar-nv <168006707+dnandakumar-nv@users.noreply.github.com>
@dnandakumar-nv
Copy link
Contributor Author

/ok to test 34a179a

1 similar comment
@dnandakumar-nv
Copy link
Contributor Author

/ok to test 34a179a

@dnandakumar-nv
Copy link
Contributor Author

/merge

@rapids-bot rapids-bot bot merged commit 2c512ce into NVIDIA:develop Apr 15, 2025
10 of 11 checks passed
yczhang-nv pushed a commit to yczhang-nv/NeMo-Agent-Toolkit that referenced this pull request Apr 21, 2025
Feature: Function Tracking Decorator

This PR adds a straightforward decorator, `@track_function`, that makes it easy to log details about function calls. It works on synchronous and asynchronous functions, generators (sync and async), and class methods—all with a consistent interface.

Key Points:
- Versatile Support: Works with regular functions, async functions, generators, and even class methods.
- Optional Metadata: You can add extra context by passing a dictionary to the decorator. Metadata keys are validated to ensure they’re strings.
- Data Serialization: Automatically serializes input arguments and outputs (or each yield) into a JSON-friendly format. It even handles Pydantic models by calling model_dump() if needed.
- Intermediate Event Logging: The decorator sends tracking events to a reactive stream in three stages:
-- SPAN_START: When the function starts.
-- SPAN_CHUNK: For each yielded item (for generators).
-- SPAN_END: When the function returns (or ends in generators).

Usage Example:

```python
from aiq.profiler.decorators.function_tracking import track_function

@track_function(metadata={"action": "compute", "source": "api"})
def add_numbers(a, b):
    return a + b
```
Just apply `@track_function` to any function (or method) where you need a bit of extra insight into the inputs and outputs. Subscribe to the reactive stream to see the published events and get a better sense of the function’s flow.

## By Submitting this PR I confirm:
- I am familiar with the [Contributing Guidelines](https://github.com/NVIDIA/AgentIQ/blob/develop/docs/source/advanced/contributing.md).
- We require that all contributors "sign-off" on their commits. This certifies that the contribution is your original work, or you have rights to submit it under the same license, or a compatible license.
  - Any contribution which contains commits that are not Signed-Off will not be accepted.
- When the PR is ready for review, new or existing tests cover these changes.
- When the PR is ready for review, the documentation is up to date with these changes.

Authors:
  - Dhruv Nandakumar (https://github.com/dnandakumar-nv)

Approvers:
  - Michael Demoret (https://github.com/mdemoret-nv)

URL: NVIDIA#98
Signed-off-by: Yuchen Zhang <134643420+yczhang-nv@users.noreply.github.com>
ericevans-nv pushed a commit to ericevans-nv/agent-iq that referenced this pull request Apr 23, 2025
Feature: Function Tracking Decorator

This PR adds a straightforward decorator, `@track_function`, that makes it easy to log details about function calls. It works on synchronous and asynchronous functions, generators (sync and async), and class methods—all with a consistent interface.

Key Points:
- Versatile Support: Works with regular functions, async functions, generators, and even class methods.
- Optional Metadata: You can add extra context by passing a dictionary to the decorator. Metadata keys are validated to ensure they’re strings.
- Data Serialization: Automatically serializes input arguments and outputs (or each yield) into a JSON-friendly format. It even handles Pydantic models by calling model_dump() if needed.
- Intermediate Event Logging: The decorator sends tracking events to a reactive stream in three stages:
-- SPAN_START: When the function starts.
-- SPAN_CHUNK: For each yielded item (for generators).
-- SPAN_END: When the function returns (or ends in generators).

Usage Example:

```python
from aiq.profiler.decorators.function_tracking import track_function

@track_function(metadata={"action": "compute", "source": "api"})
def add_numbers(a, b):
    return a + b
```
Just apply `@track_function` to any function (or method) where you need a bit of extra insight into the inputs and outputs. Subscribe to the reactive stream to see the published events and get a better sense of the function’s flow.

## By Submitting this PR I confirm:
- I am familiar with the [Contributing Guidelines](https://github.com/NVIDIA/AgentIQ/blob/develop/docs/source/advanced/contributing.md).
- We require that all contributors "sign-off" on their commits. This certifies that the contribution is your original work, or you have rights to submit it under the same license, or a compatible license.
  - Any contribution which contains commits that are not Signed-Off will not be accepted.
- When the PR is ready for review, new or existing tests cover these changes.
- When the PR is ready for review, the documentation is up to date with these changes.

Authors:
  - Dhruv Nandakumar (https://github.com/dnandakumar-nv)

Approvers:
  - Michael Demoret (https://github.com/mdemoret-nv)

URL: NVIDIA#98
Signed-off-by: Eric Evans <194135482+ericevans-nv@users.noreply.github.com>
ericevans-nv pushed a commit to ericevans-nv/agent-iq that referenced this pull request Apr 23, 2025
Feature: Function Tracking Decorator

This PR adds a straightforward decorator, `@track_function`, that makes it easy to log details about function calls. It works on synchronous and asynchronous functions, generators (sync and async), and class methods—all with a consistent interface.

Key Points:
- Versatile Support: Works with regular functions, async functions, generators, and even class methods.
- Optional Metadata: You can add extra context by passing a dictionary to the decorator. Metadata keys are validated to ensure they’re strings.
- Data Serialization: Automatically serializes input arguments and outputs (or each yield) into a JSON-friendly format. It even handles Pydantic models by calling model_dump() if needed.
- Intermediate Event Logging: The decorator sends tracking events to a reactive stream in three stages:
-- SPAN_START: When the function starts.
-- SPAN_CHUNK: For each yielded item (for generators).
-- SPAN_END: When the function returns (or ends in generators).

Usage Example:

```python
from aiq.profiler.decorators.function_tracking import track_function

@track_function(metadata={"action": "compute", "source": "api"})
def add_numbers(a, b):
    return a + b
```
Just apply `@track_function` to any function (or method) where you need a bit of extra insight into the inputs and outputs. Subscribe to the reactive stream to see the published events and get a better sense of the function’s flow.

## By Submitting this PR I confirm:
- I am familiar with the [Contributing Guidelines](https://github.com/NVIDIA/AgentIQ/blob/develop/docs/source/advanced/contributing.md).
- We require that all contributors "sign-off" on their commits. This certifies that the contribution is your original work, or you have rights to submit it under the same license, or a compatible license.
  - Any contribution which contains commits that are not Signed-Off will not be accepted.
- When the PR is ready for review, new or existing tests cover these changes.
- When the PR is ready for review, the documentation is up to date with these changes.

Authors:
  - Dhruv Nandakumar (https://github.com/dnandakumar-nv)

Approvers:
  - Michael Demoret (https://github.com/mdemoret-nv)

URL: NVIDIA#98
Signed-off-by: Eric Evans <194135482+ericevans-nv@users.noreply.github.com>
yczhang-nv pushed a commit to yczhang-nv/NeMo-Agent-Toolkit that referenced this pull request May 8, 2025
Feature: Function Tracking Decorator

This PR adds a straightforward decorator, `@track_function`, that makes it easy to log details about function calls. It works on synchronous and asynchronous functions, generators (sync and async), and class methods—all with a consistent interface.

Key Points:
- Versatile Support: Works with regular functions, async functions, generators, and even class methods.
- Optional Metadata: You can add extra context by passing a dictionary to the decorator. Metadata keys are validated to ensure they’re strings.
- Data Serialization: Automatically serializes input arguments and outputs (or each yield) into a JSON-friendly format. It even handles Pydantic models by calling model_dump() if needed.
- Intermediate Event Logging: The decorator sends tracking events to a reactive stream in three stages:
-- SPAN_START: When the function starts.
-- SPAN_CHUNK: For each yielded item (for generators).
-- SPAN_END: When the function returns (or ends in generators).

Usage Example:

```python
from aiq.profiler.decorators.function_tracking import track_function

@track_function(metadata={"action": "compute", "source": "api"})
def add_numbers(a, b):
    return a + b
```
Just apply `@track_function` to any function (or method) where you need a bit of extra insight into the inputs and outputs. Subscribe to the reactive stream to see the published events and get a better sense of the function’s flow.

## By Submitting this PR I confirm:
- I am familiar with the [Contributing Guidelines](https://github.com/NVIDIA/AgentIQ/blob/develop/docs/source/advanced/contributing.md).
- We require that all contributors "sign-off" on their commits. This certifies that the contribution is your original work, or you have rights to submit it under the same license, or a compatible license.
  - Any contribution which contains commits that are not Signed-Off will not be accepted.
- When the PR is ready for review, new or existing tests cover these changes.
- When the PR is ready for review, the documentation is up to date with these changes.

Authors:
  - Dhruv Nandakumar (https://github.com/dnandakumar-nv)

Approvers:
  - Michael Demoret (https://github.com/mdemoret-nv)

URL: NVIDIA#98
Signed-off-by: Yuchen Zhang <134643420+yczhang-nv@users.noreply.github.com>
AnuradhaKaruppiah pushed a commit to AnuradhaKaruppiah/oss-agentiq that referenced this pull request Aug 4, 2025
Feature: Function Tracking Decorator

This PR adds a straightforward decorator, `@track_function`, that makes it easy to log details about function calls. It works on synchronous and asynchronous functions, generators (sync and async), and class methods—all with a consistent interface.

Key Points:
- Versatile Support: Works with regular functions, async functions, generators, and even class methods.
- Optional Metadata: You can add extra context by passing a dictionary to the decorator. Metadata keys are validated to ensure they’re strings.
- Data Serialization: Automatically serializes input arguments and outputs (or each yield) into a JSON-friendly format. It even handles Pydantic models by calling model_dump() if needed.
- Intermediate Event Logging: The decorator sends tracking events to a reactive stream in three stages:
-- SPAN_START: When the function starts.
-- SPAN_CHUNK: For each yielded item (for generators).
-- SPAN_END: When the function returns (or ends in generators).

Usage Example:

```python
from aiq.profiler.decorators.function_tracking import track_function

@track_function(metadata={"action": "compute", "source": "api"})
def add_numbers(a, b):
    return a + b
```
Just apply `@track_function` to any function (or method) where you need a bit of extra insight into the inputs and outputs. Subscribe to the reactive stream to see the published events and get a better sense of the function’s flow.

## By Submitting this PR I confirm:
- I am familiar with the [Contributing Guidelines](https://github.com/NVIDIA/AgentIQ/blob/develop/docs/source/advanced/contributing.md).
- We require that all contributors "sign-off" on their commits. This certifies that the contribution is your original work, or you have rights to submit it under the same license, or a compatible license.
  - Any contribution which contains commits that are not Signed-Off will not be accepted.
- When the PR is ready for review, new or existing tests cover these changes.
- When the PR is ready for review, the documentation is up to date with these changes.

Authors:
  - Dhruv Nandakumar (https://github.com/dnandakumar-nv)

Approvers:
  - Michael Demoret (https://github.com/mdemoret-nv)

URL: NVIDIA#98
scheckerNV pushed a commit to scheckerNV/aiq-factory-reset that referenced this pull request Aug 22, 2025
Feature: Function Tracking Decorator

This PR adds a straightforward decorator, `@track_function`, that makes it easy to log details about function calls. It works on synchronous and asynchronous functions, generators (sync and async), and class methods—all with a consistent interface.

Key Points:
- Versatile Support: Works with regular functions, async functions, generators, and even class methods.
- Optional Metadata: You can add extra context by passing a dictionary to the decorator. Metadata keys are validated to ensure they’re strings.
- Data Serialization: Automatically serializes input arguments and outputs (or each yield) into a JSON-friendly format. It even handles Pydantic models by calling model_dump() if needed.
- Intermediate Event Logging: The decorator sends tracking events to a reactive stream in three stages:
-- SPAN_START: When the function starts.
-- SPAN_CHUNK: For each yielded item (for generators).
-- SPAN_END: When the function returns (or ends in generators).

Usage Example:

```python
from aiq.profiler.decorators.function_tracking import track_function

@track_function(metadata={"action": "compute", "source": "api"})
def add_numbers(a, b):
    return a + b
```
Just apply `@track_function` to any function (or method) where you need a bit of extra insight into the inputs and outputs. Subscribe to the reactive stream to see the published events and get a better sense of the function’s flow.

## By Submitting this PR I confirm:
- I am familiar with the [Contributing Guidelines](https://github.com/NVIDIA/AgentIQ/blob/develop/docs/source/advanced/contributing.md).
- We require that all contributors "sign-off" on their commits. This certifies that the contribution is your original work, or you have rights to submit it under the same license, or a compatible license.
  - Any contribution which contains commits that are not Signed-Off will not be accepted.
- When the PR is ready for review, new or existing tests cover these changes.
- When the PR is ready for review, the documentation is up to date with these changes.

Authors:
  - Dhruv Nandakumar (https://github.com/dnandakumar-nv)

Approvers:
  - Michael Demoret (https://github.com/mdemoret-nv)

URL: NVIDIA#98
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

improvement Improvement to existing functionality non-breaking Non-breaking change

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants