nvidia-nat 1.2.0rc5__py3-none-any.whl

aiq/utils/optional_imports.py

@@ -0,0 +1,142 @@
+# 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 importlib
+import logging
+from types import ModuleType
+
+logger = logging.getLogger(__name__)
+
+
+class OptionalImportError(Exception):
+    """Raised when an optional import fails."""
+
+    def __init__(self, module_name: str, additional_message: str = ""):
+        super().__init__(f"Optional dependency '{module_name}' is not installed. {additional_message}")
+
+
+class TelemetryOptionalImportError(OptionalImportError):
+    """Raised when an optional import of telemetry dependencies fails."""
+
+    def __init__(self, module_name: str):
+        super().__init__(
+            module_name,
+            "But the configuration file contains tracing exporters. "
+            "If you want to use this feature, please install it with: uv pip install -e '.[telemetry]'",
+        )
+
+
+def optional_import(module_name: str) -> ModuleType:
+    """Attempt to import a module, raising OptionalImportError if it fails."""
+    try:
+        return importlib.import_module(module_name)
+    except ImportError as e:
+        raise OptionalImportError(module_name) from e
+
+
+def telemetry_optional_import(module_name: str) -> ModuleType:
+    """Attempt to import a module, raising TelemetryOptionalImportError if it fails."""
+    try:
+        return importlib.import_module(module_name)
+    except ImportError as e:
+        raise TelemetryOptionalImportError(module_name) from e
+
+
+def try_import_opentelemetry() -> ModuleType:
+    """Get the opentelemetry module if available."""
+    return telemetry_optional_import("opentelemetry")
+
+
+def try_import_phoenix() -> ModuleType:
+    """Get the phoenix module if available."""
+    return telemetry_optional_import("phoenix")
+
+
+# Dummy OpenTelemetry classes for when the package is not available
+class DummySpan:
+    """Dummy span class that does nothing when OpenTelemetry is not available."""
+
+    def __init__(self, *args, **kwargs):
+        pass
+
+    def end(self, *args, **kwargs):
+        pass
+
+    def set_attribute(self, *args, **kwargs):
+        pass
+
+
+class DummyTracer:
+    """Dummy tracer class that returns dummy spans."""
+
+    def start_span(self, *args, **kwargs):
+        return DummySpan()
+
+
+class DummyTracerProvider:
+    """Dummy tracer provider that returns dummy tracers."""
+
+    @staticmethod
+    def get_tracer(*args, **kwargs):
+        return DummyTracer()
+
+    @staticmethod
+    def add_span_processor(*args, **kwargs):
+        pass
+
+
+class DummyTrace:
+    """Dummy trace module that returns dummy tracer providers."""
+
+    @staticmethod
+    def get_tracer_provider():
+        return DummyTracerProvider()
+
+    @staticmethod
+    def set_tracer_provider(*args, **kwargs):
+        pass
+
+    @staticmethod
+    def get_tracer(*args, **kwargs):
+        return DummyTracer()
+
+
+class DummySpanExporter:
+    """Dummy span exporter that does nothing."""
+
+    @staticmethod
+    def export(*args, **kwargs):
+        pass
+
+    @staticmethod
+    def shutdown(*args, **kwargs):
+        pass
+
+
+class DummyBatchSpanProcessor:
+    """Dummy implementation of BatchSpanProcessor for when OpenTelemetry is not available."""
+
+    def __init__(self, *args, **kwargs):
+        pass
+
+    @staticmethod
+    def shutdown(*args, **kwargs):
+        pass
+
+
+# Dummy functions for when OpenTelemetry is not available
+def dummy_set_span_in_context(*args, **kwargs) -> None:
+    """Dummy function that does nothing."""
+    return None

aiq/utils/producer_consumer_queue.py

@@ -0,0 +1,178 @@
+# 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 asyncio
+import typing
+
+_T = typing.TypeVar("_T")
+
+
+class QueueClosed(Exception):
+    'Exception raised when the queue is closed'
+    pass
+
+
+class AsyncIOProducerConsumerQueue(asyncio.Queue, typing.Generic[_T]):
+    """
+    Custom queue.Queue implementation which supports closing and uses recursive locks
+    """
+
+    def __init__(self, maxsize=0) -> None:
+        super().__init__(maxsize=maxsize)
+
+        self._closed = asyncio.Event()
+        self._is_closed = False
+
+    async def __aiter__(self):
+        try:
+            while True:
+                yield await self.get()
+        except QueueClosed:
+            return
+
+    async def join(self):
+        """Block until all items in the queue have been gotten and processed.
+
+        The count of unfinished tasks goes up whenever an item is added to the
+        queue. The count goes down whenever a consumer calls task_done() to
+        indicate that the item was retrieved and all work on it is complete.
+        When the count of unfinished tasks drops to zero, join() unblocks.
+        """
+
+        # First wait for the closed flag to be set
+        await self._closed.wait()
+
+        if self._unfinished_tasks > 0:
+            await self._finished.wait()
+
+    async def put(self, item):
+        """Put an item into the queue.
+
+        Put an item into the queue. If the queue is full, wait until a free
+        slot is available before adding item.
+        """
+        while self.full() and not self._is_closed:
+            putter = self._get_loop().create_future()
+            self._putters.append(putter)
+            try:
+                await putter
+            except Exception:
+                putter.cancel()  # Just in case putter is not done yet.
+                try:
+                    # Clean self._putters from canceled putters.
+                    self._putters.remove(putter)
+                except ValueError:
+                    # The putter could be removed from self._putters by a
+                    # previous get_nowait call.
+                    pass
+                if not self.full() and not putter.cancelled():
+                    # We were woken up by get_nowait(), but can't take
+                    # the call.  Wake up the next in line.
+                    self._wakeup_next(self._putters)
+                raise
+
+        if (self._is_closed):
+            raise QueueClosed  # @IgnoreException
+
+        return self.put_nowait(item)
+
+    async def get(self) -> _T:
+        """Remove and return an item from the queue.
+
+        If queue is empty, wait until an item is available.
+        """
+        while self.empty() and not self._is_closed:
+            getter = self._get_loop().create_future()
+            self._getters.append(getter)
+            try:
+                await getter
+            except Exception:
+                getter.cancel()  # Just in case getter is not done yet.
+                try:
+                    # Clean self._getters from canceled getters.
+                    self._getters.remove(getter)
+                except ValueError:
+                    # The getter could be removed from self._getters by a
+                    # previous put_nowait call.
+                    pass
+                if not self.empty() and not getter.cancelled():
+                    # We were woken up by put_nowait(), but can't take
+                    # the call.  Wake up the next in line.
+                    self._wakeup_next(self._getters)
+                raise
+
+        if (self.empty() and self._is_closed):
+            raise QueueClosed  # @IgnoreException
+
+        return self.get_nowait()
+
+    def put_blocking(self, item: _T):
+        """
+        Synchronously block until the item can be put.
+        This method creates or uses an event loop internally to call the async put().
+        If the queue is closed, it raises QueueClosed.
+
+        NOTE: If you already have an event loop running in this same thread, calling
+              `run_until_complete` can cause conflicts or an error. Typically, you only
+              want to do this from a pure synchronous environment.
+        """
+
+        # If the queue is already closed, raise immediately
+        if self._is_closed:
+            raise QueueClosed("Queue is closed, cannot put more items.")
+
+        # Quick check: if there's space, just put_nowait() and exit
+        # (This covers the trivial case with no blocking)
+        if not self.full():
+            self.put_nowait(item)
+            return None
+
+        # If we do need to block, we run self.put(...) in an event loop
+        # We'll attempt to get the currently running loop if there is one,
+        # otherwise create a new one. If there's an existing loop, we might get
+        # an error if that loop is in the same thread. Adjust logic as needed.
+
+        try:
+            # If a loop is already running in this thread, get_running_loop() will succeed.
+            loop = asyncio.get_running_loop()
+        except RuntimeError:
+            # Means no running event loop in this thread -> create a new loop
+            loop = asyncio.new_event_loop()
+            try:
+                result = loop.run_until_complete(self.put(item))
+            finally:
+                loop.close()
+            return None
+
+        # If we got a running loop, but we aren't inside an async function,
+        # do a "blocking" wait by scheduling the put and waiting:
+        future = asyncio.run_coroutine_threadsafe(self.put(item), loop)
+        result = future.result()  # blocks until done
+        return result
+
+    async def close(self):
+        """Close the queue."""
+        if (not self._is_closed):
+            self._is_closed = True
+
+            # Hit the flag
+            self._closed.set()
+
+            self._wakeup_next(self._putters)
+            self._wakeup_next(self._getters)
+
+    def is_closed(self) -> bool:
+        """Check if the queue is closed."""
+        return self._is_closed

aiq/utils/reactive/base/observable_base.py

@@ -0,0 +1,65 @@
+# 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 typing
+from abc import ABC
+from abc import abstractmethod
+from collections.abc import Callable
+from typing import Generic
+from typing import TypeVar
+
+from aiq.utils.reactive.base.observer_base import ObserverBase
+from aiq.utils.reactive.subscription import Subscription
+
+# Covariant type param: An Observable producing type X can also produce
+# a subtype of X.
+_T_out_co = TypeVar("_T_out_co", covariant=True)  # pylint: disable=invalid-name
+_T = TypeVar("_T")  # pylint: disable=invalid-name
+
+OnNext = Callable[[_T], None]
+OnError = Callable[[Exception], None]
+OnComplete = Callable[[], None]
+
+
+class ObservableBase(Generic[_T_out_co], ABC):
+    """
+    Abstract base class for an Observable that can be subscribed to.
+    Produces items of type _T_out for its subscribers.
+    """
+
+    @typing.overload
+    def subscribe(self, on_next: ObserverBase[_T_out_co]) -> Subscription:
+        ...
+
+    @typing.overload
+    def subscribe(self,
+                  on_next: OnNext[_T_out_co] | None = None,
+                  on_error: OnError | None = None,
+                  on_complete: OnComplete | None = None) -> Subscription:
+        ...
+
+    @abstractmethod
+    def subscribe(self,
+                  on_next: ObserverBase[_T_out_co] | OnNext[_T_out_co] | None = None,
+                  on_error: OnError | None = None,
+                  on_complete: OnComplete | None = None) -> Subscription:
+        """
+        Subscribes an Observer or callbacks to this Observable.
+
+        If an Observer is provided, it will be subscribed to this Observable.
+        If callbacks are provided, they will be wrapped into an Observer and
+        subscribed to this Observable.
+        """
+        pass

aiq/utils/reactive/base/observer_base.py

@@ -0,0 +1,55 @@
+# 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.
+
+from abc import ABC
+from abc import abstractmethod
+from typing import Generic
+from typing import TypeVar
+
+# Contravariant type param: An Observer that can accept type X can also
+# accept any supertype of X.
+_T_in_contra = TypeVar("_T_in_contra", contravariant=True)  # pylint: disable=invalid-name
+
+
+class ObserverBase(Generic[_T_in_contra], ABC):
+    """
+    Abstract base class for an Observer that can receive events of type _T_in.
+
+    Once on_error or on_complete is called, the observer is considered stopped.
+    """
+
+    @abstractmethod
+    def on_next(self, value: _T_in_contra) -> None:
+        """
+        Called when a new item is produced. If the observer is stopped,
+        this call should be ignored or raise an error.
+        """
+        pass
+
+    @abstractmethod
+    def on_error(self, exc: Exception) -> None:
+        """
+        Called when the producer signals an unrecoverable error.
+        After this call, the observer is stopped.
+        """
+        pass
+
+    @abstractmethod
+    def on_complete(self) -> None:
+        """
+        Called when the producer signals completion (no more items).
+        After this call, the observer is stopped.
+        """
+        pass

aiq/utils/reactive/base/subject_base.py

@@ -0,0 +1,79 @@
+# 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 typing
+from abc import abstractmethod
+from collections.abc import Callable
+from typing import TypeVar
+
+from .observable_base import ObservableBase
+from .observer_base import ObserverBase
+
+if typing.TYPE_CHECKING:
+    from aiq.utils.reactive.subscription import Subscription
+
+T = TypeVar("T")
+
+OnNext = Callable[[T], None]
+OnError = Callable[[Exception], None]
+OnComplete = Callable[[], None]
+
+
+class SubjectBase(ObserverBase[T], ObservableBase[T]):
+    """
+    Minimal interface we expect from the Subject for unsubscribing logic.
+    """
+
+    @abstractmethod
+    def _unsubscribe_observer(self, observer: object) -> None:
+        pass
+
+    @abstractmethod
+    def subscribe(self,
+                  on_next: ObserverBase[T] | OnNext[T] | None = None,
+                  on_error: OnError | None = None,
+                  on_complete: OnComplete | None = None) -> "Subscription":
+        """
+        Subscribes an Observer or callbacks to this Observable.
+
+        If an Observer is provided, it will be subscribed to this Observable.
+        If callbacks are provided, they will be wrapped into an Observer and
+        subscribed to this Observable.
+        """
+        pass
+
+    @abstractmethod
+    def on_next(self, value: T) -> None:
+        """
+        Called when a new item is produced. If the observer is stopped,
+        this call should be ignored or raise an error.
+        """
+        pass
+
+    @abstractmethod
+    def on_error(self, exc: Exception) -> None:
+        """
+        Called when the producer signals an unrecoverable error.
+        After this call, the observer is stopped.
+        """
+        pass
+
+    @abstractmethod
+    def on_complete(self) -> None:
+        """
+        Called when the producer signals completion (no more items).
+        After this call, the observer is stopped.
+        """
+        pass

aiq/utils/reactive/observable.py

@@ -0,0 +1,59 @@
+# 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.
+
+from collections.abc import Callable
+from typing import TypeVar
+
+from aiq.utils.reactive.base.observable_base import ObservableBase
+from aiq.utils.reactive.base.observer_base import ObserverBase
+from aiq.utils.reactive.observer import Observer
+from aiq.utils.reactive.subscription import Subscription
+from aiq.utils.type_utils import override
+
+# Covariant type param: An Observable producing type X can also produce
+# a subtype of X.
+_T_out_co = TypeVar("_T_out_co", covariant=True)  # pylint: disable=invalid-name
+_T = TypeVar("_T")  # pylint: disable=invalid-name
+
+OnNext = Callable[[_T], None]
+OnError = Callable[[Exception], None]
+OnComplete = Callable[[], None]
+
+
+class Observable(ObservableBase[_T_out_co]):
+    """
+    Concrete base Observable that implements subscribe, deferring actual hooking
+    logic to _subscribe_core.
+    """
+
+    __slots__ = ()
+
+    def _subscribe_core(self, observer: ObserverBase) -> Subscription:
+        """
+        By default, does nothing. Subclasses should override this to
+        attach the observer to their emission logic.
+        """
+        raise NotImplementedError("Observable._subscribe_core must be implemented by subclasses")
+
+    @override
+    def subscribe(self,
+                  on_next: ObserverBase[_T_out_co] | OnNext[_T_out_co] | None = None,
+                  on_error: OnError | None = None,
+                  on_complete: OnComplete | None = None) -> "Subscription":
+
+        if isinstance(on_next, ObserverBase):
+            return self._subscribe_core(on_next)
+
+        return self._subscribe_core(Observer(on_next, on_error, on_complete))

aiq/utils/reactive/observer.py

@@ -0,0 +1,76 @@
+# 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 collections.abc import Callable
+from typing import TypeVar
+
+from aiq.utils.reactive.base.observer_base import ObserverBase
+
+logger = logging.getLogger(__name__)
+
+# Contravariant type param: An Observer that can accept type X can also
+# accept any supertype of X.
+_T_in_contra = TypeVar("_T_in_contra", contravariant=True)  # pylint: disable=invalid-name
+_T = TypeVar("_T")  # pylint: disable=invalid-name
+
+OnNext = Callable[[_T], None]
+OnError = Callable[[Exception], None]
+OnComplete = Callable[[], None]
+
+
+class Observer(ObserverBase[_T_in_contra]):
+    """
+    Concrete Observer that wraps user-provided callbacks into an ObserverBase.
+    """
+
+    def __init__(
+        self,
+        on_next: OnNext | None = None,
+        on_error: OnError | None = None,
+        on_complete: OnComplete | None = None,
+    ) -> None:
+        self._on_next = on_next
+        self._on_error = on_error
+        self._on_complete = on_complete
+        self._stopped = False
+
+    def on_next(self, value: _T) -> None:
+        if self._stopped:
+            return
+        if self._on_next is None:
+            return
+        try:
+            self._on_next(value)
+        except Exception as exc:
+            # If the callback itself raises, treat that as an error
+            self.on_error(exc)
+
+    def on_error(self, exc: Exception) -> None:
+        if not self._stopped:
+            if self._on_error:
+                try:
+                    self._on_error(exc)
+                except Exception as e:
+                    logger.exception("Error in on_error callback: %s", e, exc_info=True)
+
+    def on_complete(self) -> None:
+        if not self._stopped:
+            self._stopped = True
+            if self._on_complete:
+                try:
+                    self._on_complete()
+                except Exception as e:
+                    logger.exception("Error in on_complete callback: %s", e, exc_info=True)