aiqtoolkit 1.1.0a20250515__py3-none-any.whl → 1.1.0a20251020__py3-none-any.whl

aiq/profiler/decorators/function_tracking.py

@@ -1,254 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import functools
-import inspect
-import uuid
-from typing import Any
-
-from pydantic import BaseModel
-
-from aiq.builder.context import AIQContext
-from aiq.builder.intermediate_step_manager import IntermediateStepManager
-from aiq.data_models.intermediate_step import IntermediateStepPayload
-from aiq.data_models.intermediate_step import IntermediateStepType
-from aiq.data_models.intermediate_step import TraceMetadata
-
-
-# --- Helper function to recursively serialize any object into JSON-friendly data ---
-def _serialize_data(obj: Any) -> Any:
-    """Convert `obj` into a structure that can be passed to `json.dumps(...)`."""
-    if isinstance(obj, BaseModel):
-        # Convert Pydantic model to dict
-        return obj.model_dump()
-
-    if isinstance(obj, dict):
-        return {str(k): _serialize_data(v) for k, v in obj.items()}
-    if isinstance(obj, (list, tuple, set)):
-        return [_serialize_data(item) for item in obj]
-
-    if isinstance(obj, (str, int, float, bool, type(None))):
-        return obj
-
-    # Fallback
-    return str(obj)
-
-
-def _prepare_serialized_args_kwargs(*args, **kwargs) -> tuple[list[Any], dict[str, Any]]:
-    """Serialize args and kwargs before calling the wrapped function."""
-    serialized_args = [_serialize_data(a) for a in args]
-    serialized_kwargs = {k: _serialize_data(v) for k, v in kwargs.items()}
-    return serialized_args, serialized_kwargs
-
-
-def push_intermediate_step(step_manager: IntermediateStepManager,
-                           identifier: str,
-                           function_name: str,
-                           event_type: IntermediateStepType,
-                           args: Any = None,
-                           kwargs: Any = None,
-                           output: Any = None,
-                           metadata: dict[str, Any] | None = None) -> None:
-    """Push an intermediate step to the AIQ Toolkit Event Stream."""
-
-    payload = IntermediateStepPayload(UUID=identifier,
-                                      event_type=event_type,
-                                      name=function_name,
-                                      metadata=TraceMetadata(
-                                          span_inputs=[args, kwargs],
-                                          span_outputs=output,
-                                          provided_metadata=metadata,
-                                      ))
-
-    step_manager.push_intermediate_step(payload)
-
-
-def track_function(func: Any = None, *, metadata: dict[str, Any] | None = None):
-    """
-    Decorator that can wrap any type of function (sync, async, generator,
-    async generator) and executes "tracking logic" around it.
-
-    - If the function is async, it will be wrapped in an async function.
-    - If the function is a generator, it will be wrapped in a generator function.
-    - If the function is an async generator, it will be wrapped in an async generator function.
-    - If the function is sync, it will be wrapped in a sync function.
-    """
-    function_name: str = func.__name__ if func else "<unknown_function>"
-
-    # If called as @track_function(...) but not immediately passed a function
-    if func is None:
-
-        def decorator_wrapper(actual_func):
-            return track_function(actual_func, metadata=metadata)
-
-        return decorator_wrapper
-
-    # --- 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.")
-
-    # --- Now detect the function type and wrap accordingly ---
-    if inspect.isasyncgenfunction(func):
-        # ---------------------
-        # ASYNC GENERATOR
-        # ---------------------
-
-        @functools.wraps(func)
-        async def async_gen_wrapper(*args, **kwargs):
-            step_manager: IntermediateStepManager = AIQContext.get().intermediate_step_manager
-            # 1) Serialize input
-            serialized_args, serialized_kwargs = _prepare_serialized_args_kwargs(*args, **kwargs)
-
-            invocation_id = str(uuid.uuid4())
-            push_intermediate_step(step_manager,
-                                   invocation_id,
-                                   function_name,
-                                   IntermediateStepType.SPAN_START,
-                                   args=serialized_args,
-                                   kwargs=serialized_kwargs,
-                                   metadata=metadata)
-
-            # 2) Call the original async generator
-            async for item in func(*args, **kwargs):
-                # 3) Serialize the yielded item before yielding it
-                serialized_item = _serialize_data(item)
-                push_intermediate_step(step_manager,
-                                       invocation_id,
-                                       function_name,
-                                       IntermediateStepType.SPAN_CHUNK,
-                                       args=serialized_args,
-                                       kwargs=serialized_kwargs,
-                                       output=serialized_item,
-                                       metadata=metadata)
-                yield item  # yield the original item
-
-            push_intermediate_step(step_manager,
-                                   invocation_id,
-                                   function_name,
-                                   IntermediateStepType.SPAN_END,
-                                   args=serialized_args,
-                                   kwargs=serialized_kwargs,
-                                   output=None,
-                                   metadata=metadata)
-
-            # 4) Post-yield logic if any
-
-        return async_gen_wrapper
-
-    if inspect.iscoroutinefunction(func):
-        # ---------------------
-        # ASYNC FUNCTION
-        # ---------------------
-        @functools.wraps(func)
-        async def async_wrapper(*args, **kwargs):
-            step_manager: IntermediateStepManager = AIQContext.get().intermediate_step_manager
-            serialized_args, serialized_kwargs = _prepare_serialized_args_kwargs(*args, **kwargs)
-            invocation_id = str(uuid.uuid4())
-            push_intermediate_step(step_manager,
-                                   invocation_id,
-                                   function_name,
-                                   IntermediateStepType.SPAN_START,
-                                   args=serialized_args,
-                                   kwargs=serialized_kwargs,
-                                   metadata=metadata)
-
-            result = await func(*args, **kwargs)
-
-            serialized_result = _serialize_data(result)
-            push_intermediate_step(step_manager,
-                                   invocation_id,
-                                   function_name,
-                                   IntermediateStepType.SPAN_END,
-                                   args=serialized_args,
-                                   kwargs=serialized_kwargs,
-                                   output=serialized_result,
-                                   metadata=metadata)
-
-            return result
-
-        return async_wrapper
-
-    if inspect.isgeneratorfunction(func):
-        # ---------------------
-        # SYNC GENERATOR
-        # ---------------------
-        @functools.wraps(func)
-        def sync_gen_wrapper(*args, **kwargs):
-            step_manager: IntermediateStepManager = AIQContext.get().intermediate_step_manager
-            serialized_args, serialized_kwargs = _prepare_serialized_args_kwargs(*args, **kwargs)
-            invocation_id = str(uuid.uuid4())
-            push_intermediate_step(step_manager,
-                                   invocation_id,
-                                   function_name,
-                                   IntermediateStepType.SPAN_START,
-                                   args=serialized_args,
-                                   kwargs=serialized_kwargs,
-                                   metadata=metadata)
-
-            for item in func(*args, **kwargs):
-                serialized_item = _serialize_data(item)
-                push_intermediate_step(step_manager,
-                                       invocation_id,
-                                       function_name,
-                                       IntermediateStepType.SPAN_CHUNK,
-                                       args=serialized_args,
-                                       kwargs=serialized_kwargs,
-                                       output=serialized_item,
-                                       metadata=metadata)
-
-                yield item  # yield the original item
-
-            push_intermediate_step(step_manager,
-                                   invocation_id,
-                                   function_name,
-                                   IntermediateStepType.SPAN_END,
-                                   args=serialized_args,
-                                   kwargs=serialized_kwargs,
-                                   output=None,
-                                   metadata=metadata)
-
-        return sync_gen_wrapper
-
-    @functools.wraps(func)
-    def sync_wrapper(*args, **kwargs):
-        step_manager: IntermediateStepManager = AIQContext.get().intermediate_step_manager
-        serialized_args, serialized_kwargs = _prepare_serialized_args_kwargs(*args, **kwargs)
-        invocation_id = str(uuid.uuid4())
-        push_intermediate_step(step_manager,
-                               invocation_id,
-                               function_name,
-                               IntermediateStepType.SPAN_START,
-                               args=serialized_args,
-                               kwargs=serialized_kwargs,
-                               metadata=metadata)
-
-        result = func(*args, **kwargs)
-
-        serialized_result = _serialize_data(result)
-        push_intermediate_step(step_manager,
-                               invocation_id,
-                               function_name,
-                               IntermediateStepType.SPAN_END,
-                               args=serialized_args,
-                               kwargs=serialized_kwargs,
-                               output=serialized_result,
-                               metadata=metadata)
-
-        return result
-
-    return sync_wrapper

aiq/profiler/forecasting/config.py

@@ -1,18 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-# If you have any global constants or defaults
-DEFAULT_MODEL_TYPE = "randomforest"
-DEFAULT_MATRIX_LENGTH = 10

aiq/profiler/forecasting/model_trainer.py

@@ -1,75 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-# forecasting/model_trainer.py
-
-import logging
-
-from aiq.profiler.forecasting.config import DEFAULT_MODEL_TYPE
-from aiq.profiler.forecasting.models import ForecastingBaseModel
-from aiq.profiler.forecasting.models import LinearModel
-from aiq.profiler.forecasting.models import RandomForestModel
-from aiq.profiler.intermediate_property_adapter import IntermediatePropertyAdaptor
-
-logger = logging.getLogger(__name__)
-
-
-def create_model(model_type: str) -> ForecastingBaseModel:
-    """
-    A simple factory method that returns a model instance
-    based on the input string. Extend this with more model
-    classes (e.g., PolynomialModel, RandomForestModel, etc.).
-    """
-    if model_type == "linear":
-        return LinearModel()
-    if model_type == "randomforest":
-        return RandomForestModel()
-
-    raise ValueError(f"Unsupported model_type: {model_type}")
-
-
-class ModelTrainer:
-    """
-    Orchestrates data preprocessing, training, and returning
-    a fitted model.
-
-    Parameters
-    ----------
-    model_type: str, default = "randomforest"
-        The type of model to train. Options include "linear" and "randomforest".
-    """
-
-    def __init__(self, model_type: str = DEFAULT_MODEL_TYPE):
-        self.model_type = model_type
-        self._model = create_model(self.model_type)
-
-    def train(self, raw_stats: list[list[IntermediatePropertyAdaptor]]) -> ForecastingBaseModel:
-        """
-        Train the model using the `raw_stats` training data.
-
-        Parameters
-        ----------
-        raw_stats: list[list[IntermediatePropertyAdaptor]]
-            Stats collected by the profiler.
-
-        Returns
-        -------
-        ForecastingBaseModel
-            A fitted model.
-        """
-
-        self._model.fit(raw_stats)
-
-        return self._model

aiq/profiler/forecasting/models/__init__.py

@@ -1,22 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-# forecasting/models/__init__.py
-
-from .forecasting_base_model import ForecastingBaseModel
-from .linear_model import LinearModel
-from .random_forest_regressor import RandomForestModel
-
-__all__ = ["ForecastingBaseModel", "LinearModel", "RandomForestModel"]

aiq/profiler/forecasting/models/forecasting_base_model.py

@@ -1,40 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-# forecasting/models/base_model.py
-
-from abc import ABC, abstractmethod
-import numpy as np
-
-
-class ForecastingBaseModel(ABC):
-    """
-    Abstract base class for all models in this package.
-    """
-
-    @abstractmethod
-    def fit(self, raw_stats):
-        """
-        Train/fine-tune the model on the provided dataset.
-        """
-        pass
-
-    @abstractmethod
-    def predict(self, raw_stats) -> np.ndarray:
-        """
-        Predict using the trained model.
-        Returns a np.ndarray, shape = (N, 4).
-        """
-        pass

aiq/profiler/forecasting/models/linear_model.py

@@ -1,196 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import logging
-
-import numpy as np
-
-from aiq.profiler.forecasting.models.forecasting_base_model import ForecastingBaseModel
-from aiq.profiler.intermediate_property_adapter import IntermediatePropertyAdaptor
-
-logger = logging.getLogger(__name__)
-
-
-class LinearModel(ForecastingBaseModel):
-    """
-    A linear regression model that conforms to the BaseModel interface.
-    """
-
-    def __init__(self):
-        super().__init__()
-
-        try:
-            from sklearn.linear_model import LinearRegression
-        except ImportError:
-            logger.error("scikit-learn is not installed. Please install scikit-learn to use the LinearModel "
-                         "profiling model or install `aiq[profiler]` to install all necessary profiling packages.")
-
-            raise
-
-        self.model = LinearRegression()
-        self.matrix_length = None
-
-    def fit(self, raw_stats: list[list[IntermediatePropertyAdaptor]]):
-        """
-        X: shape (N, M)  # M = matrix_length * 4
-        y: shape (N, 4)
-        """
-        x_flat, y_flat = self._prep_for_model_training(raw_stats)
-
-        logger.info("Training dataset size: X=%s, y=%s", x_flat.shape, y_flat.shape)
-
-        # 3) Fit
-        self.model.fit(x_flat, y_flat)
-
-    def predict(self, raw_stats: list[list[IntermediatePropertyAdaptor]]) -> np.ndarray:
-        """
-        Predict using the fitted linear model.
-        Returns shape (N, 4)
-        """
-        X = self._prep_single(raw_stats)
-        return self.model.predict(X)
-
-    def _prep_single(self, raw_stats: list[list[IntermediatePropertyAdaptor]]) -> np.ndarray:
-        arr, _ = self._extract_token_usage_meta(raw_stats)
-        arr = arr[0]
-        n_rows = arr.shape[0]
-
-        matrix_length = self.matrix_length
-
-        assert matrix_length is not None, "matrix_length must be set before calling _prep_single"
-
-        if n_rows >= matrix_length:
-            # Keep the latest matrix_length rows
-            x_mat = arr[-matrix_length:, :]
-        else:
-            # Pad with zeros at the top
-            pad_size = matrix_length - n_rows
-            pad_block = np.zeros((pad_size, arr.shape[1]), dtype=arr.dtype)
-            x_mat = np.vstack([pad_block, arr])
-
-        return x_mat
-
-    def _prep_for_model_training(self, raw_stats: list[list[IntermediatePropertyAdaptor]]):
-        raw_matrices, matrix_length = self._extract_token_usage_meta(raw_stats)
-
-        self.matrix_length = matrix_length
-
-        x_list = []
-        y_list = []
-        for arr in raw_matrices:
-            samples = self._preprocess_for_forecasting(arr, matrix_length)
-            for (x_mat, y_mat) in samples:
-                x_list.append(x_mat)
-                y_list.append(y_mat)
-
-        # 2) Flatten features
-        x_flat, y_flat = self._flatten_features(x_list, y_list)
-
-        return x_flat, y_flat
-
-    def _extract_token_usage_meta(self, all_requests_data: list[list[IntermediatePropertyAdaptor]]):
-        import math
-
-        all_run_data = []
-        call_stack_sizes = []
-
-        for prompt in all_requests_data:
-            run_data = []
-            seconds_between_call_map = {}
-
-            for stat in prompt:
-                if stat.event_type.value == "LLM_START":
-                    seconds_between_call_map[stat.UUID] = stat.seconds_between_calls
-
-                if stat.event_type.value == "LLM_END":
-                    step_data = [
-                        seconds_between_call_map[stat.UUID],
-                        stat.token_usage.prompt_tokens,
-                        stat.token_usage.completion_tokens
-                    ]
-
-                    run_data.append(step_data)
-
-            all_run_data.append(run_data)
-            call_stack_sizes.append(len(run_data))
-
-        all_run_data = [np.array(run) for run in all_run_data]
-        recommended_matrix_length = math.ceil(sum(call_stack_sizes) / len(call_stack_sizes))
-
-        return all_run_data, recommended_matrix_length
-
-    def _preprocess_for_forecasting(self, arr: np.ndarray, matrix_length: int):
-        """
-        Given a 2D NumPy array `arr` of shape (n_rows, 4), generate a list of
-        (input_array, output_array) pairs for forecasting, each of shape:
-
-        - input_array: (matrix_length, 4) after padding/trimming
-        - output_array: (1, 4)
-        """
-        n_rows = arr.shape[0]
-
-        # partial_sums[i] = sum of arr[i:] per column
-        partial_sums = np.flip(np.cumsum(np.flip(arr, axis=0), axis=0), axis=0)
-
-        samples = []
-        for i in range(n_rows):
-            x_untrimmed = arr[:i + 1, :]
-            # Trim or pad
-            current_len = x_untrimmed.shape[0]
-            if current_len > matrix_length:
-                x_mat = x_untrimmed[-matrix_length:, :]
-            elif current_len < matrix_length:
-                pad_size = matrix_length - current_len
-                pad_block = np.zeros((pad_size, x_untrimmed.shape[1]), dtype=arr.dtype)
-                x_mat = np.vstack([pad_block, x_untrimmed])
-            else:
-                x_mat = x_untrimmed
-
-            # Compute output
-            if i == n_rows - 1:
-                y_vec = np.array([0, 0, 0, 0], dtype=arr.dtype)
-            else:
-                n_below = n_rows - (i + 1)
-                sum_below = partial_sums[i + 1]
-                avg_col0 = sum_below[0] / n_below
-                sum_rest = sum_below[1:]
-                y_vec = np.concatenate(([avg_col0], sum_rest))
-
-            samples.append((x_mat, y_vec.reshape(1, 4)))
-
-        return samples
-
-    def _flatten_features(self, x_list, y_list):
-        """
-        x_list: list of arrays, each of shape (matrix_length, 4)
-        y_list: list of arrays, each of shape (1, 4)
-
-        Returns:
-            x_flat: np.array of shape (N, matrix_length*4)
-            y_flat: np.array of shape (N, 4)
-        """
-        flattened_x = []
-        flattened_y = []
-
-        for x_mat, y_mat in zip(x_list, y_list):
-            x_1d = x_mat.flatten()  # shape -> (matrix_length*4,)
-            y_1d = y_mat.flatten()  # shape -> (4,)
-            flattened_x.append(x_1d)
-            flattened_y.append(y_1d)
-
-        x_flat = np.array(flattened_x)
-        y_flat = np.array(flattened_y)
-        logger.debug("Flattened features to shapes: %s (X), %s (y).", x_flat.shape, y_flat.shape)
-        return x_flat, y_flat