aiqtoolkit 1.2.0a20250707__py3-none-any.whl → 1.2.0rc1__py3-none-any.whl

aiq/data_models/its_strategy.py

@@ -0,0 +1,30 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-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 typing
+
+from .common import BaseModelRegistryTag
+from .common import TypedBaseModel
+
+
+class ITSStrategyBaseConfig(TypedBaseModel, BaseModelRegistryTag):
+    """
+    Base configuration class for Inference Time Scaling (ITS) strategy.
+    This class is used to define the structure of ITS strategy configurations.
+    """
+    pass
+
+
+ITSStrategyBaseConfigT = typing.TypeVar("ITSStrategyBaseConfigT", bound=ITSStrategyBaseConfig)

aiq/data_models/llm.py

@@ -20,6 +20,7 @@ from .common import TypedBaseModel
 
 
 class LLMBaseConfig(TypedBaseModel, BaseModelRegistryTag):
+    """Base configuration for LLM providers."""
     pass
 
 

aiq/data_models/memory.py

@@ -20,6 +20,7 @@ from .common import TypedBaseModel
 
 
 class MemoryBaseConfig(TypedBaseModel, BaseModelRegistryTag):
+    """ The base level config object for a memory object. Memories provide an interface for storing and retrieving. """
     pass
 
 

aiq/data_models/object_store.py

@@ -0,0 +1,44 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-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 typing
+
+from .common import BaseModelRegistryTag
+from .common import TypedBaseModel
+
+
+class ObjectStoreBaseConfig(TypedBaseModel, BaseModelRegistryTag):
+    pass
+
+
+ObjectStoreBaseConfigT = typing.TypeVar("ObjectStoreBaseConfigT", bound=ObjectStoreBaseConfig)
+
+
+class KeyAlreadyExistsError(Exception):
+
+    def __init__(self, key: str, additional_message: str | None = None):
+        parts = [f"Key already exists: {key}."]
+        if additional_message:
+            parts.append(additional_message)
+        super().__init__(" ".join(parts))
+
+
+class NoSuchKeyError(Exception):
+
+    def __init__(self, key: str, additional_message: str | None = None):
+        parts = [f"No object found with key: {key}."]
+        if additional_message:
+            parts.append(additional_message)
+        super().__init__(" ".join(parts))

aiq/data_models/retry_mixin.py

@@ -0,0 +1,35 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-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.
+
+from pydantic import BaseModel
+from pydantic import Field
+
+
+class RetryMixin(BaseModel):
+    """Mixin class for retry configuration."""
+    do_auto_retry: bool = Field(default=True,
+                                description="Whether to automatically retry method calls"
+                                " that fail with a retryable error.",
+                                exclude=True)
+    num_retries: int = Field(default=5,
+                             description="Number of times to retry a method call that fails"
+                             " with a retryable error.",
+                             exclude=True)
+    retry_on_status_codes: list[int | str] = Field(default_factory=lambda: [429, 500, 502, 503, 504],
+                                                   description="List of HTTP status codes that should trigger a retry.",
+                                                   exclude=True)
+    retry_on_errors: list[str] | None = Field(default_factory=lambda: ["Too Many Requests"],
+                                              description="List of error substrings that should trigger a retry.",
+                                              exclude=True)

aiq/data_models/span.py

@@ -0,0 +1,187 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-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 time
+import uuid
+from enum import Enum
+from typing import Any
+
+from pydantic import BaseModel
+from pydantic import Field
+from pydantic import field_validator
+
+logger = logging.getLogger(__name__)
+
+
+class SpanKind(Enum):
+    LLM = "LLM"
+    TOOL = "TOOL"
+    WORKFLOW = "WORKFLOW"
+    TASK = "TASK"
+    FUNCTION = "FUNCTION"
+    CUSTOM = "CUSTOM"
+    SPAN = "SPAN"
+    EMBEDDER = "EMBEDDER"
+    RETRIEVER = "RETRIEVER"
+    AGENT = "AGENT"
+    RERANKER = "RERANKER"
+    GUARDRAIL = "GUARDRAIL"
+    EVALUATOR = "EVALUATOR"
+    UNKNOWN = "UNKNOWN"
+
+
+EVENT_TYPE_TO_SPAN_KIND_MAP = {
+    "LLM_START": SpanKind.LLM,
+    "LLM_END": SpanKind.LLM,
+    "LLM_NEW_TOKEN": SpanKind.LLM,
+    "TOOL_START": SpanKind.TOOL,
+    "TOOL_END": SpanKind.TOOL,
+    "WORKFLOW_START": SpanKind.WORKFLOW,
+    "WORKFLOW_END": SpanKind.WORKFLOW,
+    "TASK_START": SpanKind.TASK,
+    "TASK_END": SpanKind.TASK,
+    "FUNCTION_START": SpanKind.FUNCTION,
+    "FUNCTION_END": SpanKind.FUNCTION,
+    "CUSTOM_START": SpanKind.CUSTOM,
+    "CUSTOM_END": SpanKind.CUSTOM,
+    "SPAN_START": SpanKind.SPAN,
+    "SPAN_END": SpanKind.SPAN,
+    "EMBEDDER_START": SpanKind.EMBEDDER,
+    "EMBEDDER_END": SpanKind.EMBEDDER,
+    "RETRIEVER_START": SpanKind.RETRIEVER,
+    "RETRIEVER_END": SpanKind.RETRIEVER,
+    "AGENT_START": SpanKind.AGENT,
+    "AGENT_END": SpanKind.AGENT,
+    "RERANKER_START": SpanKind.RERANKER,
+    "RERANKER_END": SpanKind.RERANKER,
+    "GUARDRAIL_START": SpanKind.GUARDRAIL,
+    "GUARDRAIL_END": SpanKind.GUARDRAIL,
+    "EVALUATOR_START": SpanKind.EVALUATOR,
+    "EVALUATOR_END": SpanKind.EVALUATOR,
+}
+
+
+def event_type_to_span_kind(event_type: str) -> SpanKind:
+    """Convert an event type to a span kind.
+
+    Args:
+        event_type (str): The event type to convert.
+
+    Returns:
+        SpanKind: The span kind.
+    """
+    return EVENT_TYPE_TO_SPAN_KIND_MAP.get(event_type, SpanKind.UNKNOWN)
+
+
+class SpanAttributes(Enum):
+    AIQ_SPAN_KIND = "aiq.span.kind"
+    INPUT_VALUE = "input.value"
+    INPUT_MIME_TYPE = "input.mime_type"
+    LLM_TOKEN_COUNT_PROMPT = "llm.token_count.prompt"
+    LLM_TOKEN_COUNT_COMPLETION = "llm.token_count.completion"
+    LLM_TOKEN_COUNT_TOTAL = "llm.token_count.total"
+    OUTPUT_VALUE = "output.value"
+    OUTPUT_MIME_TYPE = "output.mime_type"
+    AIQ_USAGE_NUM_LLM_CALLS = "aiq.usage.num_llm_calls"
+    AIQ_USAGE_SECONDS_BETWEEN_CALLS = "aiq.usage.seconds_between_calls"
+    AIQ_USAGE_TOKEN_COUNT_PROMPT = "aiq.usage.token_count.prompt"
+    AIQ_USAGE_TOKEN_COUNT_COMPLETION = "aiq.usage.token_count.completion"
+    AIQ_USAGE_TOKEN_COUNT_TOTAL = "aiq.usage.token_count.total"
+    AIQ_EVENT_TYPE = "aiq.event_type"
+
+
+class MimeTypes(Enum):
+    TEXT = "text/plain"
+    JSON = "application/json"
+
+
+class SpanStatusCode(Enum):
+    OK = "OK"
+    ERROR = "ERROR"
+    UNSET = "UNSET"
+
+
+class SpanEvent(BaseModel):
+    timestamp: float = Field(default_factory=lambda: int(time.time() * 1e9), description="The timestamp of the event.")
+    name: str = Field(description="The name of the event.")
+    attributes: dict[str, Any] = Field(default_factory=dict, description="The attributes of the event.")
+
+
+class SpanStatus(BaseModel):
+    code: SpanStatusCode = Field(default=SpanStatusCode.OK, description="The status code of the span.")
+    message: str | None = Field(default=None, description="The status message of the span.")
+
+
+class SpanContext(BaseModel):
+    trace_id: int = Field(default_factory=lambda: uuid.uuid4().int, description="The 128-bit trace ID of the span.")
+    span_id: int = Field(default_factory=lambda: uuid.uuid4().int & ((1 << 64) - 1),
+                         description="The 64-bit span ID of the span.")
+
+
+class Span(BaseModel):
+    name: str = Field(description="The name of the span.")
+    context: SpanContext | None = Field(default=None, description="The context of the span.")
+    parent: "Span | None" = Field(default=None, description="The parent span of the span.")
+    start_time: int = Field(default_factory=lambda: int(time.time() * 1e9), description="The start time of the span.")
+    end_time: int | None = Field(default=None, description="The end time of the span.")
+    attributes: dict[str, Any] = Field(default_factory=dict, description="The attributes of the span.")
+    events: list[SpanEvent] = Field(default_factory=list, description="The events of the span.")
+    status: SpanStatus = Field(default_factory=SpanStatus, description="The status of the span.")
+
+    @field_validator('context', mode='before')
+    @classmethod
+    def set_default_context(cls, v: SpanContext | None) -> SpanContext:
+        """Set the default context if the context is not provided.
+
+        Args:
+            v (SpanContext | None): The context to set.
+
+        Returns:
+            SpanContext: The context.
+        """
+        if v is None:
+            return SpanContext()
+        return v
+
+    def set_attribute(self, key: str, value: Any) -> None:
+        """Set the attribute of the span.
+
+        Args:
+            key (str): The key of the attribute.
+            value (Any): The value of the attribute.
+        """
+        self.attributes[key] = value
+
+    def add_event(self, name: str, attributes: dict[str, Any] | None = None) -> None:
+        """Add an event to the span.
+
+        Args:
+            name (str): The name of the event.
+            attributes (dict[str, Any] | None): The attributes of the event.
+        """
+        if attributes is None:
+            attributes = {}
+        self.events = self.events + [SpanEvent(name=name, attributes=attributes)]
+
+    def end(self, end_time: int | None = None) -> None:
+        """End the span.
+
+        Args:
+            end_time (int | None): The end time of the span.
+        """
+        if end_time is None:
+            end_time = int(time.time() * 1e9)
+        self.end_time = end_time

aiq/data_models/telemetry_exporter.py

@@ -15,8 +15,8 @@
 
 import typing
 
-from .common import BaseModelRegistryTag
-from .common import TypedBaseModel
+from aiq.data_models.common import BaseModelRegistryTag
+from aiq.data_models.common import TypedBaseModel
 
 
 class TelemetryExporterBaseConfig(TypedBaseModel, BaseModelRegistryTag):

aiq/embedder/nim_embedder.py

@@ -24,6 +24,7 @@ from aiq.builder.builder import Builder
 from aiq.builder.embedder import EmbedderProviderInfo
 from aiq.cli.register_workflow import register_embedder_provider
 from aiq.data_models.embedder import EmbedderBaseConfig
+from aiq.data_models.retry_mixin import RetryMixin
 
 allowed_truncate_values = ["NONE", "START", "END"]
 
@@ -37,7 +38,7 @@ def option_in_allowed_values(v):
 TruncationOption = typing.Annotated[str, AfterValidator(option_in_allowed_values)]
 
 
-class NIMEmbedderModelConfig(EmbedderBaseConfig, name="nim"):
+class NIMEmbedderModelConfig(EmbedderBaseConfig, RetryMixin, name="nim"):
     """A NVIDIA Inference Microservice (NIM) embedder provider to be used with an embedder client."""
 
     api_key: str | None = Field(default=None, description="NVIDIA API key to interact with hosted NIM.")

aiq/embedder/openai_embedder.py

@@ -21,9 +21,10 @@ from aiq.builder.builder import Builder
 from aiq.builder.embedder import EmbedderProviderInfo
 from aiq.cli.register_workflow import register_embedder_provider
 from aiq.data_models.embedder import EmbedderBaseConfig
+from aiq.data_models.retry_mixin import RetryMixin
 
 
-class OpenAIEmbedderModelConfig(EmbedderBaseConfig, name="openai"):
+class OpenAIEmbedderModelConfig(EmbedderBaseConfig, RetryMixin, name="openai"):
     """An OpenAI LLM provider to be used with an LLM client."""
 
     model_config = ConfigDict(protected_namespaces=())

aiq/eval/config.py

@@ -17,13 +17,18 @@ from pathlib import Path
 
 from pydantic import BaseModel
 
+from aiq.eval.evaluator.evaluator_model import EvalInput
+from aiq.eval.evaluator.evaluator_model import EvalOutput
+from aiq.eval.usage_stats import UsageStats
+from aiq.profiler.data_models import ProfilerResults
+
 
 class EvaluationRunConfig(BaseModel):
     """
     Parameters used for a single evaluation run.
     """
     config_file: Path
-    dataset: str | None  # dataset file path can be specified in the config file
+    dataset: str | None = None  # dataset file path can be specified in the config file
     result_json_path: str = "$"
     skip_workflow: bool = False
     skip_completed_entries: bool = False
@@ -31,6 +36,14 @@ class EvaluationRunConfig(BaseModel):
     endpoint_timeout: int = 300
     reps: int = 1
     override: tuple[tuple[str, str], ...] = ()
+    # If false, the output will not be written to the output directory. This is
+    # useful when running evaluation via another tool.
+    write_output: bool = True
+    # if true, the dataset is adjusted to a multiple of the concurrency
+    adjust_dataset_size: bool = False
+    # number of passes at each concurrency, if 0 the dataset is adjusted to a multiple of the
+    # concurrency. The is only used if adjust_dataset_size is true
+    num_passes: int = 0
 
 
 class EvaluationRunOutput(BaseModel):
@@ -40,3 +53,8 @@ class EvaluationRunOutput(BaseModel):
     workflow_output_file: Path | None
     evaluator_output_files: list[Path]
     workflow_interrupted: bool
+
+    eval_input: EvalInput
+    evaluation_results: list[tuple[str, EvalOutput]]
+    usage_stats: UsageStats | None = None
+    profiler_results: ProfilerResults

aiq/eval/dataset_handler/dataset_handler.py

@@ -14,6 +14,7 @@
 # limitations under the License.
 
 import json
+import math
 
 import pandas as pd
 
@@ -33,12 +34,23 @@ class DatasetHandler:
     One DatasetHandler object is needed for each dataset to be evaluated.
     """
 
-    def __init__(self, dataset_config: EvalDatasetConfig, reps: int):
+    def __init__(self,
+                 dataset_config: EvalDatasetConfig,
+                 reps: int,
+                 concurrency: int,
+                 num_passes: int | None = None,
+                 adjust_dataset_size: bool = False):
         from aiq.eval.intermediate_step_adapter import IntermediateStepAdapter
 
         self.dataset_config = dataset_config
         self.dataset_filter = DatasetFilter(dataset_config.filter)
         self.reps = reps
+
+        # number of passes at specific concurrency
+        self.concurrency = concurrency
+        self.num_passes = num_passes
+        self.adjust_dataset_size = adjust_dataset_size
+
         # Helpers
         self.intermediate_step_adapter = IntermediateStepAdapter()
 
@@ -109,6 +121,63 @@ class DatasetHandler:
 
         return input_df
 
+    def adjust_dataset(self, input_df: pd.DataFrame) -> pd.DataFrame:
+        """
+        Adjust the dataset so its length is a multiple of concurrency.
+
+        If num_passes > 0:
+            dataset size is adjusted to concurrency * num_passes
+        else:
+            dataset size is adjusted to the largest multiple of concurrency
+            that is less than or equal to the current dataset size
+        """
+        if self.concurrency <= 0:
+            raise ValueError("Concurrency must be > 0")
+
+        if self.num_passes < 0:
+            raise ValueError("num_passes must be >= 0")
+
+        original_size = input_df.shape[0]
+
+        # Calculate target size
+        if self.num_passes > 0:
+            # When num_passes is specified, always use concurrency * num_passes
+            # This respects the user's intent for exact number of passes
+            target_size = self.concurrency * self.num_passes
+        else:
+            # When num_passes = 0, use the largest multiple of concurrency <= original_size
+            # If original_size < concurrency, we need at least concurrency rows
+            if original_size >= self.concurrency:
+                target_size = (original_size // self.concurrency) * self.concurrency
+            else:
+                target_size = self.concurrency
+
+        if target_size == 0:
+            raise ValueError("Input dataset too small for even one batch at given concurrency.")
+
+        id_col = self.dataset_config.id_key
+
+        # If we need more rows than we have, replicate the dataset
+        if original_size < target_size:
+            # Clean existing _rep suffix if present
+            input_df[id_col] = input_df[id_col].astype(str).str.replace(r"_rep\d+$", "", regex=True)
+
+            # Calculate how many complete copies we need
+            copies_needed = math.ceil(target_size / original_size)
+
+            # Create the replicated dataframe
+            replicated_dfs = []
+            for i in range(copies_needed):
+                df_copy = input_df.copy()
+                if i > 0:  # Add suffix to all but the first copy
+                    df_copy[id_col] = df_copy[id_col].astype(str) + f"_rep{i}"
+                replicated_dfs.append(df_copy)
+
+            input_df = pd.concat(replicated_dfs, ignore_index=True)
+
+        # Return exactly the target size
+        return input_df.head(target_size)
+
     def get_eval_input_from_dataset(self, dataset: str) -> EvalInput:
         # read the dataset and convert it to EvalInput
 
@@ -127,9 +196,14 @@ class DatasetHandler:
         input_df = self.dataset_filter.apply_filters(input_df)
         input_df.drop_duplicates(subset=[self.dataset_config.id_key], inplace=True)
 
+        if self.reps > 1 and self.adjust_dataset_size:
+            raise ValueError("reps and adjust_dataset_size are mutually exclusive")
+
         # If more than one repetition is needed, replicate the rows
         if self.reps > 1:
             input_df = self.setup_reps(input_df)
+        elif self.adjust_dataset_size:
+            input_df = self.adjust_dataset(input_df)
 
         # Convert the DataFrame to a list of EvalInput objects
         return self.get_eval_input_from_df(input_df)

aiq/eval/evaluate.py

@@ -99,12 +99,34 @@ class EvaluationRun:  # pylint: disable=too-many-public-methods
             max_timestamp = max(step.event_timestamp for step in item.trajectory)
             runtime = max_timestamp - min_timestamp
         else:
+            min_timestamp = 0.0
+            max_timestamp = 0.0
             runtime = 0.0
 
+        # find llm latency by calculating p95 of all llm calls
+        llm_latencies = []
+        previous_llm_start_time = None
+        for step in steps:
+            if step.event_type == "LLM_START":
+                previous_llm_start_time = step.event_timestamp
+            elif step.event_type == "LLM_END" and previous_llm_start_time is not None:
+                llm_latencies.append(step.event_timestamp - previous_llm_start_time)
+                previous_llm_start_time = None
+
+        # Calculate p95 LLM latency (or 0 if no LLM calls)
+        if llm_latencies:
+            import numpy as np
+            llm_latency = float(np.percentile(llm_latencies, 95))
+        else:
+            llm_latency = 0.0
+
         # add the usage stats to the usage stats dict
         self.usage_stats.usage_stats_items[item.id] = UsageStatsItem(usage_stats_per_llm=usage_stats_per_llm,
                                                                      runtime=runtime,
-                                                                     total_tokens=total_tokens)
+                                                                     total_tokens=total_tokens,
+                                                                     min_timestamp=min_timestamp,
+                                                                     max_timestamp=max_timestamp,
+                                                                     llm_latency=llm_latency)
         return self.usage_stats.usage_stats_items[item.id]
 
     async def run_workflow_local(self, session_manager: AIQSessionManager):
@@ -221,7 +243,9 @@ class EvaluationRun:  # pylint: disable=too-many-public-methods
         for input_item in self.eval_input.eval_input_items:
             all_stats.append(input_item.trajectory)
 
-        profiler_runner = ProfilerRunner(self.eval_config.general.profiler, self.eval_config.general.output_dir)
+        profiler_runner = ProfilerRunner(self.eval_config.general.profiler,
+                                         self.eval_config.general.output_dir,
+                                         write_output=self.config.write_output)
 
         return await profiler_runner.run(all_stats)
 
@@ -308,6 +332,11 @@ class EvaluationRun:  # pylint: disable=too-many-public-methods
             self.evaluator_output_files.append(output_file)
             logger.info("Evaluation results written to %s", output_file)
 
+    def publish_output(self, dataset_handler: DatasetHandler, profiler_results: ProfilerResults):
+        """Publish the output"""
+        if self.config.write_output:
+            self.write_output(dataset_handler, profiler_results)
+
         if self.workflow_interrupted:
             # Issue a warning if the workflow was not completed on all datasets
             msg = ("Workflow execution was interrupted due to an error. The results may be incomplete. "
@@ -415,7 +444,11 @@ class EvaluationRun:  # pylint: disable=too-many-public-methods
                 workflow_interrupted=self.workflow_interrupted,
             )
 
-        dataset_handler = DatasetHandler(dataset_config=dataset_config, reps=self.config.reps)
+        dataset_handler = DatasetHandler(dataset_config=dataset_config,
+                                         reps=self.config.reps,
+                                         concurrency=self.eval_config.general.max_concurrency,
+                                         num_passes=self.config.num_passes,
+                                         adjust_dataset_size=self.config.adjust_dataset_size)
         self.eval_input = dataset_handler.get_eval_input_from_dataset(self.config.dataset)
         if not self.eval_input.eval_input_items:
             logger.info("Dataset is empty. Nothing to evaluate.")
@@ -447,8 +480,16 @@ class EvaluationRun:  # pylint: disable=too-many-public-methods
         # Profile the workflow
         profiler_results = await self.profile_workflow()
 
-        # Write the results to the output directory
-        self.write_output(dataset_handler, profiler_results)
+        # compute total runtime
+        if self.usage_stats.usage_stats_items:
+            self.usage_stats.total_runtime = max(self.usage_stats.usage_stats_items.values(),
+                                                 key=lambda x: x.max_timestamp).max_timestamp - \
+                min(self.usage_stats.usage_stats_items.values(), key=lambda x: x.min_timestamp).min_timestamp
+        else:
+            self.usage_stats.total_runtime = 0.0
+
+        # Publish the results
+        self.publish_output(dataset_handler, profiler_results)
 
         # Run custom scripts and upload evaluation outputs to S3
         if self.eval_config.general.output:
@@ -456,8 +497,10 @@ class EvaluationRun:  # pylint: disable=too-many-public-methods
             output_uploader.run_custom_scripts()
             await output_uploader.upload_directory()
 
-        return EvaluationRunOutput(
-            workflow_output_file=self.workflow_output_file,
-            evaluator_output_files=self.evaluator_output_files,
-            workflow_interrupted=self.workflow_interrupted,
-        )
+        return EvaluationRunOutput(workflow_output_file=self.workflow_output_file,
+                                   evaluator_output_files=self.evaluator_output_files,
+                                   workflow_interrupted=self.workflow_interrupted,
+                                   eval_input=self.eval_input,
+                                   evaluation_results=self.evaluation_results,
+                                   usage_stats=self.usage_stats,
+                                   profiler_results=profiler_results)