nvidia-nat 1.2.0a20250813__py3-none-any.whl

nat/observability/register.py

@@ -0,0 +1,96 @@
+# 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
+
+from pydantic import Field
+
+from nat.builder.builder import Builder
+from nat.cli.register_workflow import register_logging_method
+from nat.cli.register_workflow import register_telemetry_exporter
+from nat.data_models.logging import LoggingBaseConfig
+from nat.data_models.telemetry_exporter import TelemetryExporterBaseConfig
+from nat.observability.mixin.file_mode import FileMode
+
+logger = logging.getLogger(__name__)
+
+
+class FileTelemetryExporterConfig(TelemetryExporterBaseConfig, name="file"):
+    """A telemetry exporter that writes runtime traces to local files with optional rolling."""
+
+    output_path: str = Field(description="Output path for logs. When rolling is disabled: exact file path. "
+                             "When rolling is enabled: directory path or file path (directory + base name).")
+    project: str = Field(description="Name to affiliate with this application.")
+    mode: FileMode = Field(
+        default=FileMode.APPEND,
+        description="File write mode: 'append' to add to existing file or 'overwrite' to start fresh.")
+    enable_rolling: bool = Field(default=False, description="Enable rolling log files based on size limits.")
+    max_file_size: int = Field(
+        default=10 * 1024 * 1024,  # 10MB
+        description="Maximum file size in bytes before rolling to a new file.")
+    max_files: int = Field(default=5, description="Maximum number of rolled files to keep.")
+    cleanup_on_init: bool = Field(default=False, description="Clean up old files during initialization.")
+
+
+@register_telemetry_exporter(config_type=FileTelemetryExporterConfig)
+async def file_telemetry_exporter(config: FileTelemetryExporterConfig, builder: Builder):  # pylint: disable=W0613
+    """
+    Build and return a FileExporter for file-based telemetry export with optional rolling.
+    """
+
+    from nat.observability.exporter.file_exporter import FileExporter
+
+    yield FileExporter(output_path=config.output_path,
+                       project=config.project,
+                       mode=config.mode,
+                       enable_rolling=config.enable_rolling,
+                       max_file_size=config.max_file_size,
+                       max_files=config.max_files,
+                       cleanup_on_init=config.cleanup_on_init)
+
+
+class ConsoleLoggingMethodConfig(LoggingBaseConfig, name="console"):
+    """A logger to write runtime logs to the console."""
+
+    level: str = Field(description="The logging level of console logger.")
+
+
+@register_logging_method(config_type=ConsoleLoggingMethodConfig)
+async def console_logging_method(config: ConsoleLoggingMethodConfig, builder: Builder):  # pylint: disable=W0613
+    """
+    Build and return a StreamHandler for console-based logging.
+    """
+    level = getattr(logging, config.level.upper(), logging.INFO)
+    handler = logging.StreamHandler()
+    handler.setLevel(level)
+    yield handler
+
+
+class FileLoggingMethod(LoggingBaseConfig, name="file"):
+    """A logger to write runtime logs to a file."""
+
+    path: str = Field(description="The file path to save the logging output.")
+    level: str = Field(description="The logging level of file logger.")
+
+
+@register_logging_method(config_type=FileLoggingMethod)
+async def file_logging_method(config: FileLoggingMethod, builder: Builder):  # pylint: disable=W0613
+    """
+    Build and return a FileHandler for file-based logging.
+    """
+    level = getattr(logging, config.level.upper(), logging.INFO)
+    handler = logging.FileHandler(filename=config.path, mode="a", encoding="utf-8")
+    handler.setLevel(level)
+    yield handler

nat/observability/utils/__init__.py

@@ -0,0 +1,14 @@
+# 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.

nat/observability/utils/dict_utils.py

@@ -0,0 +1,236 @@
+# 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 asyncio
+import logging
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
+from typing import Any
+from weakref import WeakKeyDictionary
+
+logger = logging.getLogger(__name__)
+
+
+class KeyedLock:
+    """
+    A lock manager that provides an asyncio-compatible lock for each unique key.
+
+    This allows for fine-grained locking based on arbitrary keys, so that
+    concurrent operations on different keys do not block each other.
+
+    Attributes:
+        _locks (AsyncDictionary): A dictionary to store locks per key.
+    """
+
+    def __init__(self):
+        """
+        Initialize the KeyedLock with an internal AsyncSafeWeakKeyDictionary to store locks per key.
+        """
+        self._locks: AsyncDictionary = AsyncDictionary()
+
+    @asynccontextmanager
+    async def get_lock(self, key: Any) -> AsyncGenerator[None]:
+        """
+        Async context manager to acquire a lock for a specific key.
+
+        Args:
+            key (Any): The key to lock on.
+
+        Yields:
+            None: Control is yielded while the lock is held.
+        """
+        lock = await self._locks.get(key)
+        if lock is None:
+            lock = asyncio.Lock()
+            await self._locks.set(key, lock)
+        async with lock:
+            yield
+
+    async def delete(self, key: Any) -> None:
+        """
+        Remove the lock associated with the given key, if it exists.
+
+        Args:
+            key (Any): The key whose lock should be removed.
+        """
+        await self._locks.delete(key)
+
+    async def clear(self) -> None:
+        """
+        Remove all locks managed by this KeyedLock instance.
+        """
+        await self._locks.clear()
+
+
+class AsyncDictionary:
+    """
+    An asyncio-safe dictionary.
+
+    This class wraps a regular dictionary with an asyncio.Lock to ensure
+    thread safety for concurrent async operations.
+
+    Attributes:
+        _dict (dict): A dictionary to store the key-value pairs.
+        _lock (asyncio.Lock): A lock to synchronize access to the dictionary.
+    """
+
+    def __init__(self):
+        """
+        Initialize the AsyncDictionary with a regular dictionary and an asyncio.Lock.
+        """
+        self._dict: dict = {}
+        self._lock = asyncio.Lock()
+
+    async def get(self, key: Any, default: Any | None = None) -> Any | None:
+        """
+        Get the value associated with the given key, or return default if not found.
+
+        Args:
+            key (Any): The key to look up.
+            default (Any | None, optional): The value to return if key is not found. Defaults to None.
+
+        Returns:
+            Any | None: The value associated with the key, or default.
+        """
+        async with self._lock:
+            return self._dict.get(key, default)
+
+    async def keys(self) -> list[Any]:
+        """
+        Get a list of all keys currently in the dictionary.
+
+        Returns:
+            list[Any]: A list of keys.
+        """
+        async with self._lock:
+            return list(self._dict.keys())
+
+    async def values(self) -> list[Any]:
+        """
+        Get a list of all values currently in the dictionary.
+
+        Returns:
+            list[Any]: A list of values.
+        """
+        async with self._lock:
+            return list(self._dict.values())
+
+    async def set(self, key: Any, value: Any) -> None:
+        """
+        Set the value for the given key, overwriting any existing value.
+
+        Args:
+            key (Any): The key to set.
+            value (Any): The value to associate with the key.
+        """
+        async with self._lock:
+            self._dict[key] = value
+
+    async def set_strict(self, key: Any, value: Any) -> None:
+        """
+        Set the value for the given key only if the key does not already exist.
+
+        Args:
+            key (Any): The key to set.
+            value (Any): The value to associate with the key.
+
+        Raises:
+            ValueError: If the key already exists in the dictionary.
+        """
+        async with self._lock:
+            if key in self._dict:
+                raise ValueError(f"Key '{key}' already exists")
+            self._dict[key] = value
+
+    async def delete(self, key: Any) -> None:
+        """
+        Remove the value associated with the given key, if it exists.
+
+        Args:
+            key (Any): The key to remove.
+        """
+        async with self._lock:
+            self._dict.pop(key, None)
+
+    async def delete_strict(self, key: Any) -> None:
+        """
+        Remove the value associated with the given key, raising an error if the key does not exist.
+
+        Args:
+            key (Any): The key to remove.
+
+        Raises:
+            ValueError: If the key does not exist in the dictionary.
+        """
+        async with self._lock:
+            if key not in self._dict:
+                raise ValueError(f"Key '{key}' does not exist")
+            self._dict.pop(key)
+
+    async def clear(self) -> None:
+        """
+        Remove all items from the dictionary.
+        """
+        async with self._lock:
+            self._dict.clear()
+
+    async def items(self) -> dict[Any, Any]:
+        """
+        Get a copy of the dictionary's items as a regular dict.
+
+        Returns:
+            dict[Any, Any]: A copy of the dictionary's items.
+        """
+        async with self._lock:
+            return dict(self._dict)  # Return a copy to prevent external modification
+
+
+class AsyncSafeWeakKeyDictionary(AsyncDictionary):
+    """
+    An asyncio-safe, weakly-referenced dictionary.
+
+    This class wraps a WeakKeyDictionary with an asyncio.Lock to ensure
+    thread safety for concurrent async operations.
+
+    Attributes:
+        _dict (WeakKeyDictionary): A dictionary to store the key-value pairs.
+        _lock (asyncio.Lock): A lock to synchronize access to the dictionary.
+    """
+
+    def __init__(self):
+        """
+        Initialize the AsyncSafeWeakKeyDictionary with a WeakKeyDictionary and an asyncio.Lock.
+        """
+        super().__init__()
+        self._dict: WeakKeyDictionary = WeakKeyDictionary()
+        self._lock = asyncio.Lock()
+
+
+def merge_dicts(dict1: dict, dict2: dict) -> dict:
+    """
+    Merge two dictionaries, prioritizing non-null values from the first dictionary.
+
+    Args:
+        dict1 (dict): First dictionary (higher priority)
+        dict2 (dict): Second dictionary (lower priority)
+
+    Returns:
+        dict: Merged dictionary with non-null values from dict1 taking precedence
+    """
+    result = dict2.copy()  # Start with a copy of the second dictionary
+    for key, value in dict1.items():
+        if value is not None:  # Only update if value is not None
+            result[key] = value
+    return result

nat/observability/utils/time_utils.py

@@ -0,0 +1,31 @@
+# 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
+
+logger = logging.getLogger(__name__)
+
+
+def ns_timestamp(seconds_float: float) -> int:
+    """
+    Convert a float timestamp in seconds to an integer nanosecond timestamp.
+
+    Args:
+        seconds_float (float): The timestamp in seconds (as a float).
+
+    Returns:
+        int: The timestamp in nanoseconds (as an integer).
+    """
+    return int(seconds_float * 1e9)

nat/plugins/.namespace

@@ -0,0 +1 @@
+Note: This is a python namespace package and this directory should remain empty. Do NOT add a `__init__.py` file or any other files to this directory. This file is also needed to ensure the directory exists in git.

nat/profiler/calc/__init__.py

@@ -0,0 +1,14 @@
+# 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.