nvidia-nat 1.3.0.dev2__py3-none-any.whl → 1.3.0rc1__py3-none-any.whl

nat/observability/register.py

@@ -45,7 +45,7 @@ class FileTelemetryExporterConfig(TelemetryExporterBaseConfig, name="file"):
 
 
 @register_telemetry_exporter(config_type=FileTelemetryExporterConfig)
-async def file_telemetry_exporter(config: FileTelemetryExporterConfig, builder: Builder):  # pylint: disable=W0613
+async def file_telemetry_exporter(config: FileTelemetryExporterConfig, builder: Builder):
     """
     Build and return a FileExporter for file-based telemetry export with optional rolling.
     """
@@ -68,12 +68,14 @@ class ConsoleLoggingMethodConfig(LoggingBaseConfig, name="console"):
 
 
 @register_logging_method(config_type=ConsoleLoggingMethodConfig)
-async def console_logging_method(config: ConsoleLoggingMethodConfig, builder: Builder):  # pylint: disable=W0613
+async def console_logging_method(config: ConsoleLoggingMethodConfig, builder: Builder):
     """
     Build and return a StreamHandler for console-based logging.
     """
+    import sys
+
     level = getattr(logging, config.level.upper(), logging.INFO)
-    handler = logging.StreamHandler()
+    handler = logging.StreamHandler(stream=sys.stdout)
     handler.setLevel(level)
     yield handler
 
@@ -86,7 +88,7 @@ class FileLoggingMethod(LoggingBaseConfig, name="file"):
 
 
 @register_logging_method(config_type=FileLoggingMethod)
-async def file_logging_method(config: FileLoggingMethod, builder: Builder):  # pylint: disable=W0613
+async def file_logging_method(config: FileLoggingMethod, builder: Builder):
     """
     Build and return a FileHandler for file-based logging.
     """

nat/profiler/calc/calc_runner.py

@@ -442,7 +442,7 @@ class CalcRunner:
                 runtime_fit=self.linear_analyzer.wf_runtime_fit  # May be None
             )
         except Exception as e:
-            logger.exception("Failed to plot concurrency vs. time metrics: %s", e, exc_info=True)
+            logger.exception("Failed to plot concurrency vs. time metrics: %s", e)
             logger.warning("Skipping plot of concurrency vs. time metrics")
 
     def write_output(self, output_dir: Path, calc_runner_output: CalcRunnerOutput):
@@ -506,11 +506,10 @@ class CalcRunner:
                 continue
             try:
                 calc_output = CalcRunnerOutput.model_validate_json(calc_runner_output_path.read_text())
-            except ValidationError as e:
+            except ValidationError:
                 logger.exception("Failed to validate calc runner output file %s. Skipping job %s.",
                                  calc_runner_output_path,
-                                 e,
-                                 exc_info=True)
+                                 job_dir.name)
                 continue
 
             # Extract sizing metrics from calc_data

nat/profiler/callbacks/agno_callback_handler.py

@@ -144,7 +144,7 @@ class AgnoProfilerHandler(BaseProfilerCallback):
                 return result
 
             except Exception as e:
-                logger.exception("Tool execution error: %s", e)
+                logger.error("Tool execution error: %s", e)
                 raise
 
         return wrapped_tool_execute

nat/profiler/callbacks/langchain_callback_handler.py

@@ -53,7 +53,7 @@ def _extract_tools_schema(invocation_params: dict) -> list:
     return tools_schema
 
 
-class LangchainProfilerHandler(AsyncCallbackHandler, BaseProfilerCallback):  # pylint: disable=R0901
+class LangchainProfilerHandler(AsyncCallbackHandler, BaseProfilerCallback):
     """Callback Handler that tracks NIM info."""
 
     total_tokens: int = 0
@@ -106,7 +106,7 @@ class LangchainProfilerHandler(AsyncCallbackHandler, BaseProfilerCallback):  # p
         try:
             model_name = kwargs.get("metadata")["ls_model_name"]
         except Exception as e:
-            logger.exception("Error getting model name: %s", e, exc_info=True)
+            logger.exception("Error getting model name: %s", e)
 
         run_id = str(kwargs.get("run_id", str(uuid4())))
         self._run_id_to_model_name[run_id] = model_name
@@ -144,7 +144,7 @@ class LangchainProfilerHandler(AsyncCallbackHandler, BaseProfilerCallback):  # p
         try:
             model_name = metadata["ls_model_name"] if metadata else kwargs.get("metadata")["ls_model_name"]
         except Exception as e:
-            logger.exception("Error getting model name: %s", e, exc_info=True)
+            logger.exception("Error getting model name: %s", e)
 
         run_id = str(run_id)
         self._run_id_to_model_name[run_id] = model_name
@@ -173,13 +173,13 @@ class LangchainProfilerHandler(AsyncCallbackHandler, BaseProfilerCallback):  # p
         try:
             model_name = self._run_id_to_model_name.get(str(kwargs.get("run_id", "")), "")
         except Exception as e:
-            logger.exception("Error getting model name: %s", e, exc_info=True)
+            logger.exception("Error getting model name: %s", e)
 
         usage_metadata = {}
         try:
             usage_metadata = kwargs.get("chunk").message.usage_metadata if kwargs.get("chunk") else {}
         except Exception as e:
-            logger.exception("Error getting usage metadata: %s", e, exc_info=True)
+            logger.exception("Error getting usage metadata: %s", e)
 
         stats = IntermediateStepPayload(
             event_type=IntermediateStepType.LLM_NEW_TOKEN,
@@ -206,7 +206,7 @@ class LangchainProfilerHandler(AsyncCallbackHandler, BaseProfilerCallback):  # p
             try:
                 model_name = self._run_id_to_model_name.get(str(kwargs.get("run_id", "")), "")
             except Exception as e_inner:
-                logger.exception("Error getting model name: %s from outer error %s", e_inner, e, exc_info=True)
+                logger.exception("Error getting model name: %s from outer error %s", e_inner, e)
 
         try:
             generation = response.generations[0][0]

nat/profiler/callbacks/llama_index_callback_handler.py

@@ -94,7 +94,7 @@ class LlamaIndexProfilerHandler(BaseCallbackHandler, BaseProfilerCallback):
             try:
                 model_name = payload.get(EventPayload.SERIALIZED)['model']
             except Exception as e:
-                logger.exception("Error getting model name: %s", e, exc_info=True)
+                logger.exception("Error getting model name: %s", e)
 
             llm_text_input = " ".join(prompts_or_messages) if prompts_or_messages else ""
 
@@ -159,13 +159,13 @@ class LlamaIndexProfilerHandler(BaseCallbackHandler, BaseProfilerCallback):
                     for block in response.message.blocks:
                         llm_text_output += block.text
                 except Exception as e:
-                    logger.exception("Error getting LLM text output: %s", e, exc_info=True)
+                    logger.exception("Error getting LLM text output: %s", e)
 
                 model_name = ""
                 try:
                     model_name = response.raw.model
                 except Exception as e:
-                    logger.exception("Error getting model name: %s", e, exc_info=True)
+                    logger.exception("Error getting model name: %s", e)
 
                 # Append usage data to NAT usage stats
                 with self._lock:

nat/profiler/callbacks/semantic_kernel_callback_handler.py

@@ -86,7 +86,7 @@ class SemanticKernelProfilerHandler(BaseProfilerCallback):
 
         # Gather the appropriate modules/functions based on your builder config
         for llm in self._builder_llms:
-            if self._builder_llms[llm].provider_type == 'openai':  # pylint: disable=consider-using-in
+            if self._builder_llms[llm].provider_type == 'openai':
                 functions_to_patch.extend(["openai_non_streaming", "openai_streaming"])
 
         # Grab original reference for the tool call
@@ -132,7 +132,7 @@ class SemanticKernelProfilerHandler(BaseProfilerCallback):
                         if "text" in item:
                             model_input += item["text"]
             except Exception as e:
-                logger.exception("Error in getting model input: %s", e, exc_info=True)
+                logger.exception("Error in getting model input: %s", e)
 
             input_stats = IntermediateStepPayload(event_type=IntermediateStepType.LLM_START,
                                                   framework=LLMFrameworkEnum.SEMANTIC_KERNEL,
@@ -232,7 +232,7 @@ class SemanticKernelProfilerHandler(BaseProfilerCallback):
                 return result
 
             except Exception as e:
-                logger.exception("ToolUsage._use error: %s", e)
+                logger.error("ToolUsage._use error: %s", e)
                 raise
 
         return patched_tool_call

nat/profiler/data_frame_row.py

@@ -42,7 +42,7 @@ class DataFrameRow(BaseModel):
     framework: str | None
 
     @field_validator('llm_text_input', 'llm_text_output', 'llm_new_token', mode='before')
-    def cast_to_str(cls, v):  # pylint: disable=no-self-argument
+    def cast_to_str(cls, v):
         if v is None:
             return v
         try:

nat/profiler/decorators/framework_wrapper.py

@@ -13,12 +13,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# pylint disable=ungrouped-imports
-
 from __future__ import annotations
 
 import functools
 import logging
+from collections.abc import AsyncIterator
 from collections.abc import Callable
 from contextlib import AbstractAsyncContextManager as AsyncContextManager
 from contextlib import asynccontextmanager
@@ -34,35 +33,55 @@ _library_instrumented = {
     "crewai": False,
     "semantic_kernel": False,
     "agno": False,
+    "adk": False,
 }
 
 callback_handler_var: ContextVar[Any | None] = ContextVar("callback_handler_var", default=None)
 
 
 def set_framework_profiler_handler(
-    workflow_llms: dict = None,
-    frameworks: list[LLMFrameworkEnum] = None,
+    workflow_llms: dict | None = None,
+    frameworks: list[LLMFrameworkEnum] | None = None,
 ) -> Callable[[Callable[..., AsyncContextManager[Any]]], Callable[..., AsyncContextManager[Any]]]:
     """
     Decorator that wraps an async context manager function to set up framework-specific profiling.
+
+    Args:
+        workflow_llms (dict | None): A dictionary of workflow LLM configurations.
+        frameworks (list[LLMFrameworkEnum] | None): A list of LLM frameworks used in the workflow functions.
+
+    Returns:
+        Callable[[Callable[..., AsyncContextManager[Any]]], Callable[..., AsyncContextManager[Any]]]:
+        A decorator that wraps the original function with profiling setup.
     """
 
     def decorator(func: Callable[..., AsyncContextManager[Any]]) -> Callable[..., AsyncContextManager[Any]]:
+        """The actual decorator that wraps the function.
+
+        Args:
+            func (Callable[..., AsyncContextManager[Any]]): The function to wrap.
+
+        Returns:
+            Callable[..., AsyncContextManager[Any]]: The wrapped function.
+        """
 
         @functools.wraps(func)
         @asynccontextmanager
         async def wrapper(workflow_config, builder):
 
-            if LLMFrameworkEnum.LANGCHAIN in frameworks and not _library_instrumented["langchain"]:
-                from langchain_core.tracers.context import register_configure_hook
-
+            if LLMFrameworkEnum.LANGCHAIN in frameworks:
+                # Always set a fresh handler in the current context so callbacks
+                # route to the active run. Only register the hook once globally.
                 from nat.profiler.callbacks.langchain_callback_handler import LangchainProfilerHandler
 
                 handler = LangchainProfilerHandler()
                 callback_handler_var.set(handler)
-                register_configure_hook(callback_handler_var, inheritable=True)
-                _library_instrumented["langchain"] = True
-                logger.debug("Langchain callback handler registered")
+
+                if not _library_instrumented["langchain"]:
+                    from langchain_core.tracers.context import register_configure_hook
+                    register_configure_hook(callback_handler_var, inheritable=True)
+                    _library_instrumented["langchain"] = True
+                    logger.debug("LangChain/LangGraph callback hook registered")
 
             if LLMFrameworkEnum.LLAMA_INDEX in frameworks:
                 from llama_index.core import Settings
@@ -75,8 +94,7 @@ def set_framework_profiler_handler(
                 logger.debug("LlamaIndex callback handler registered")
 
             if LLMFrameworkEnum.CREWAI in frameworks and not _library_instrumented["crewai"]:
-                from nat.plugins.crewai.crewai_callback_handler import \
-                    CrewAIProfilerHandler  # pylint: disable=ungrouped-imports,line-too-long  # noqa E501
+                from nat.plugins.crewai.crewai_callback_handler import CrewAIProfilerHandler
 
                 handler = CrewAIProfilerHandler()
                 handler.instrument()
@@ -99,6 +117,20 @@ def set_framework_profiler_handler(
                 _library_instrumented["agno"] = True
                 logger.info("Agno callback handler registered")
 
+            if LLMFrameworkEnum.ADK in frameworks and not _library_instrumented["adk"]:
+                try:
+                    from nat.plugins.adk.adk_callback_handler import ADKProfilerHandler
+                except ImportError as e:
+                    logger.warning(
+                        "ADK profiler not available. " +
+                        "Install NAT with ADK extras: pip install 'nvidia-nat[adk]'. Error: %s",
+                        e)
+                else:
+                    handler = ADKProfilerHandler()
+                    handler.instrument()
+                    _library_instrumented["adk"] = True
+                    logger.debug("ADK callback handler registered")
+
             # IMPORTANT: actually call the wrapped function as an async context manager
             async with func(workflow_config, builder) as result:
                 yield result
@@ -117,11 +149,28 @@ def chain_wrapped_build_fn(
     Convert an original build function into an async context manager that
     wraps it with a single call to set_framework_profiler_handler, passing
     all frameworks at once.
+
+    Args:
+        original_build_fn (Callable[..., AsyncContextManager]): The original build function to wrap.
+        workflow_llms (dict): A dictionary of workflow LLM configurations.
+        function_frameworks (list[LLMFrameworkEnum]): A list of LLM frameworks used in the workflow functions.
+
+    Returns:
+        Callable[..., AsyncContextManager]: The wrapped build function.
     """
 
     # Define a base async context manager that simply calls the original build function.
     @asynccontextmanager
-    async def base_fn(*args, **kwargs):
+    async def base_fn(*args, **kwargs) -> AsyncIterator[Any]:
+        """Base async context manager that calls the original build function.
+
+        Args:
+            *args: Positional arguments to pass to the original build function.
+            **kwargs: Keyword arguments to pass to the original build function.
+
+        Yields:
+            The result of the original build function.
+        """
         async with original_build_fn(*args, **kwargs) as w:
             yield w
 

nat/profiler/decorators/function_tracking.py

@@ -16,7 +16,11 @@
 import functools
 import inspect
 import uuid
+from collections.abc import Callable
 from typing import Any
+from typing import TypeVar
+from typing import cast
+from typing import overload
 
 from pydantic import BaseModel
 
@@ -36,10 +40,10 @@ def _serialize_data(obj: Any) -> Any:
 
     if isinstance(obj, dict):
         return {str(k): _serialize_data(v) for k, v in obj.items()}
-    if isinstance(obj, (list, tuple, set)):
+    if isinstance(obj, list | tuple | set):
         return [_serialize_data(item) for item in obj]
 
-    if isinstance(obj, (str, int, float, bool, type(None))):
+    if isinstance(obj, str | int | float | bool | type(None)):
         return obj
 
     # Fallback
@@ -75,7 +79,24 @@ def push_intermediate_step(step_manager: IntermediateStepManager,
     step_manager.push_intermediate_step(payload)
 
 
-def track_function(func: Any = None, *, metadata: dict[str, Any] | None = None):
+# Type variable for overloads
+F = TypeVar('F', bound=Callable[..., Any])
+
+
+# Overloads for different function types
+@overload
+def track_function(func: F, *, metadata: dict[str, Any] | None = None) -> F:
+    """Overload for when a function is passed directly."""
+    ...
+
+
+@overload
+def track_function(*, metadata: dict[str, Any] | None = None) -> Callable[[F], F]:
+    """Overload for decorator factory usage (when called with parentheses)."""
+    ...
+
+
+def track_function(func: Any = None, *, metadata: dict[str, Any] | None = None) -> Any:
     """
     Decorator that can wrap any type of function (sync, async, generator,
     async generator) and executes "tracking logic" around it.
@@ -252,3 +273,139 @@ def track_function(func: Any = None, *, metadata: dict[str, Any] | None = None):
         return result
 
     return sync_wrapper
+
+
+# Overloads for track_unregistered_function
+@overload
+def track_unregistered_function(func: F, *, name: str | None = None, metadata: dict[str, Any] | None = None) -> F:
+    """Overload for when a function is passed directly."""
+    ...
+
+
+@overload
+def track_unregistered_function(*, name: str | None = None, metadata: dict[str, Any] | None = None) -> Callable[[F], F]:
+    """Overload for decorator factory usage (when called with parentheses)."""
+    ...
+
+
+def track_unregistered_function(func: Callable[..., Any] | None = None,
+                                *,
+                                name: str | None = None,
+                                metadata: dict[str, Any] | None = None) -> Callable[..., Any]:
+    """
+    Decorator that wraps any function with scope management and automatic tracking.
+
+    - Sets active function context using the function name
+    - Leverages Context.push_active_function for built-in tracking
+    - Avoids duplicate tracking entries by relying on the library's built-in systems
+    - Supports sync/async functions and generators
+
+    Args:
+        func: The function to wrap (auto-detected when used without parentheses)
+        name: Custom name to use for tracking instead of func.__name__
+        metadata: Additional metadata to include in tracking
+    """
+
+    # If called with parameters: @track_unregistered_function(name="...", metadata={...})
+    if func is None:
+
+        def decorator_wrapper(actual_func: Callable[..., Any]) -> Callable[..., Any]:
+            # Cast to ensure type checker understands this returns a callable
+            return cast(Callable[..., Any], track_unregistered_function(actual_func, name=name, metadata=metadata))
+
+        return decorator_wrapper
+
+    # Direct decoration: @track_unregistered_function or recursive call with actual function
+    function_name: str = name if name else func.__name__
+
+    # --- Validate metadata ---
+    if metadata is not None:
+        if not isinstance(metadata, dict):
+            raise TypeError("metadata must be a dict[str, Any].")
+        if any(not isinstance(k, str) for k in metadata.keys()):
+            raise TypeError("All metadata keys must be strings.")
+
+    trace_metadata = TraceMetadata(provided_metadata=metadata)
+
+    # --- Now detect the function type and wrap accordingly ---
+    if inspect.isasyncgenfunction(func):
+        # ---------------------
+        # ASYNC GENERATOR
+        # ---------------------
+
+        @functools.wraps(func)
+        async def async_gen_wrapper(*args, **kwargs):
+            context = Context.get()
+            input_data = (
+                *args,
+                kwargs,
+            )
+            # Only do context management - let push_active_function handle tracking
+            with context.push_active_function(function_name, input_data=input_data, metadata=trace_metadata) as manager:
+                final_outputs = []
+                async for item in func(*args, **kwargs):
+                    final_outputs.append(item)
+                    yield item
+
+                manager.set_output(final_outputs)
+
+        return async_gen_wrapper
+
+    if inspect.iscoroutinefunction(func):
+        # ---------------------
+        # ASYNC FUNCTION
+        # ---------------------
+        @functools.wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            context = Context.get()
+            input_data = (
+                *args,
+                kwargs,
+            )
+
+            # Only do context management - let push_active_function handle tracking
+            with context.push_active_function(function_name, input_data=input_data, metadata=trace_metadata) as manager:
+                result = await func(*args, **kwargs)
+                manager.set_output(result)
+                return result
+
+        return async_wrapper
+
+    if inspect.isgeneratorfunction(func):
+        # ---------------------
+        # SYNC GENERATOR
+        # ---------------------
+        @functools.wraps(func)
+        def sync_gen_wrapper(*args, **kwargs):
+            context = Context.get()
+            input_data = (
+                *args,
+                kwargs,
+            )
+
+            # Only do context management - let push_active_function handle tracking
+            with context.push_active_function(function_name, input_data=input_data, metadata=trace_metadata) as manager:
+                final_outputs = []
+                for item in func(*args, **kwargs):
+                    final_outputs.append(item)
+                    yield item
+
+                manager.set_output(final_outputs)
+
+        return sync_gen_wrapper
+
+    @functools.wraps(func)
+    def sync_wrapper(*args, **kwargs):
+        context = Context.get()
+        input_data = (
+            *args,
+            kwargs,
+        )
+
+        # Only do context management - let push_active_function handle tracking
+        with context.push_active_function(function_name, input_data=input_data, metadata=trace_metadata) as manager:
+            result = func(*args, **kwargs)
+            manager.set_output(result)
+            return result
+
+    return sync_wrapper

nat/profiler/forecasting/models/forecasting_base_model.py

@@ -15,7 +15,9 @@
 
 # forecasting/models/base_model.py
 
-from abc import ABC, abstractmethod
+from abc import ABC
+from abc import abstractmethod
+
 import numpy as np
 
 

nat/profiler/inference_optimization/bottleneck_analysis/simple_stack_analysis.py

@@ -195,7 +195,7 @@ def profile_workflow_bottlenecks(all_steps: list[list[IntermediateStep]]) -> Sim
         c_max = 0
         for ts, delta in events_sub:
             c_curr += delta
-            if c_curr > c_max:  # pylint: disable=consider-using-max-builtin
+            if c_curr > c_max:  # noqa: PLR1730 - don't use max built-in
                 c_max = c_curr
         max_concurrency_by_name[op_name] = c_max
 

nat/profiler/inference_optimization/data_models.py

@@ -172,7 +172,7 @@ class CallNode(BaseModel):
         if not self.children:
             return self.duration
 
-        intervals = [(c.start_time, c.end_time) for c in self.children]  # pylint: disable=not-an-iterable
+        intervals = [(c.start_time, c.end_time) for c in self.children]
         # Sort by start time
         intervals.sort(key=lambda x: x[0])
 
@@ -204,7 +204,7 @@ class CallNode(BaseModel):
         This ensures no overlap double-counting among children.
         """
         total = self.compute_self_time()
-        for c in self.children:  # pylint: disable=not-an-iterable
+        for c in self.children:
             total += c.compute_subtree_time()
         return total
 
@@ -216,7 +216,7 @@ class CallNode(BaseModel):
         info = (f"{indent}- {self.operation_type} '{self.operation_name}' "
                 f"(uuid={self.uuid}, start={self.start_time:.2f}, "
                 f"end={self.end_time:.2f}, dur={self.duration:.2f})")
-        child_strs = [child._repr(level + 1) for child in self.children]  # pylint: disable=not-an-iterable
+        child_strs = [child._repr(level + 1) for child in self.children]
         return "\n".join([info] + child_strs)
 
 

nat/profiler/inference_optimization/experimental/prefix_span_analysis.py

@@ -228,7 +228,7 @@ def run_prefixspan(sequences_map: dict[int, list[PrefixCallNode]],
     else:
         abs_min_support = min_support
 
-    freq_patterns = ps.frequent(abs_min_support)  # pylint: disable=not-callable
+    freq_patterns = ps.frequent(abs_min_support)
     # freq_patterns => [(count, [item1, item2, ...])]
 
     results = []
@@ -321,13 +321,12 @@ def compute_coverage_and_duration(sequences_map: dict[int, list[PrefixCallNode]]
 # --------------------------------------------------------------------------------
 
 
-def prefixspan_subworkflow_with_text(  # pylint: disable=too-many-positional-arguments
-        all_steps: list[list[IntermediateStep]],
-        min_support: int | float = 2,
-        top_k: int = 10,
-        min_coverage: float = 0.0,
-        max_text_len: int = 700,
-        prefix_list: list[str] = None) -> PrefixSpanSubworkflowResult:
+def prefixspan_subworkflow_with_text(all_steps: list[list[IntermediateStep]],
+                                     min_support: int | float = 2,
+                                     top_k: int = 10,
+                                     min_coverage: float = 0.0,
+                                     max_text_len: int = 700,
+                                     prefix_list: list[str] = None) -> PrefixSpanSubworkflowResult:
     """
     1) Build sequences of calls for each example (with llm_text_input).
     2) Convert to token lists, run PrefixSpan with min_support.

nat/profiler/inference_optimization/token_uniqueness.py

@@ -66,7 +66,7 @@ def compute_inter_query_token_uniqueness_by_llm(all_steps: list[list[Intermediat
     # 2) Group by (llm_name, example_number), then sort each group
     grouped = cdf.groupby(['llm_name', 'example_number'], as_index=False, group_keys=True)
 
-    for (llm, ex_num), group_df in grouped:  # pylint: disable=unused-variable
+    for (llm, ex_num), group_df in grouped:
         # Sort by event_timestamp
         group_df = group_df.sort_values('event_timestamp', ascending=True)