nvidia-nat 1.1.0a20251020__py3-none-any.whl

nat/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

nat/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 nat.utils.reactive.base.observer_base import ObserverBase
+from nat.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)
+_T = TypeVar("_T")
+
+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

nat/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)
+
+
+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

nat/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 nat.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

nat/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 nat.utils.reactive.base.observable_base import ObservableBase
+from nat.utils.reactive.base.observer_base import ObserverBase
+from nat.utils.reactive.observer import Observer
+from nat.utils.reactive.subscription import Subscription
+from nat.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)
+_T = TypeVar("_T")
+
+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))

nat/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 nat.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)
+_T = TypeVar("_T")
+
+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)
+
+    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)

nat/utils/reactive/subject.py

@@ -0,0 +1,131 @@
+# 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 threading
+from collections.abc import Callable
+from typing import TypeVar
+
+from nat.utils.reactive.base.subject_base import SubjectBase
+from nat.utils.reactive.observable import Observable
+from nat.utils.reactive.observer import Observer
+from nat.utils.reactive.subscription import Subscription
+
+T = TypeVar("T")
+
+OnNext = Callable[[T], None]
+OnError = Callable[[Exception], None]
+OnComplete = Callable[[], None]
+
+
+class Subject(Observable[T], Observer[T], SubjectBase[T]):
+    """
+    A Subject is both an Observer (receives events) and an Observable (sends events).
+    - Maintains a list of ObserverBase[T].
+    - No internal buffering or replay; events are only delivered to current subscribers.
+    - Thread-safe via a lock.
+
+    Once on_error or on_complete is called, the Subject is closed.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self._lock = threading.RLock()
+        self._closed = False
+        self._error: Exception | None = None
+        self._observers: list[Observer[T]] = []
+        self._disposed = False
+
+    # ==========================================================================
+    # Observable[T] - for consumers
+    # ==========================================================================
+    def _subscribe_core(self, observer: Observer[T]) -> Subscription:
+        """
+        Subscribe to this subject. If disposed, returns a dummy subscription.
+        Otherwise, registers the given observer.
+        """
+        with self._lock:
+            if self._disposed:
+                # Already disposed => no subscription
+                return Subscription(self, None)
+
+            self._observers.append(observer)
+            return Subscription(self, observer)
+
+    # ==========================================================================
+    # ObserverBase[T] - for producers
+    # ==========================================================================
+    def on_next(self, value: T) -> None:
+        """
+        Called by producers to emit an item. Delivers synchronously to each observer.
+        If closed or disposed, do nothing.
+        """
+        with self._lock:
+            if self._closed or self._disposed:
+                return
+            # Copy the current observers to avoid mutation issues
+            current_observers = list(self._observers)
+
+        # Deliver outside the lock
+        for obs in current_observers:
+            obs.on_next(value)
+
+    def on_error(self, exc: Exception) -> None:
+        """
+        Called by producers to signal an error. Notifies all observers.
+        """
+        with self._lock:
+            if self._closed or self._disposed:
+                return
+            current_obs = list(self._observers)
+
+        for obs in current_obs:
+            obs.on_error(exc)
+
+    def on_complete(self) -> None:
+        """
+        Called by producers to signal completion. Notifies all observers, then
+        clears them. Subject is closed.
+        """
+        with self._lock:
+            if self._closed or self._disposed:
+                return
+            current_observers = list(self._observers)
+            self.dispose()
+
+        for obs in current_observers:
+            obs.on_complete()
+
+    # ==========================================================================
+    # SubjectBase - internal unsubscribing
+    # ==========================================================================
+    def _unsubscribe_observer(self, observer: Observer[T]) -> None:
+        with self._lock:
+            if not self._disposed and observer in self._observers:
+                self._observers.remove(observer)
+
+    # ==========================================================================
+    # Disposal
+    # ==========================================================================
+    def dispose(self) -> None:
+        """
+        Immediately close the Subject. No future on_next, on_error, or on_complete.
+        Clears all observers.
+        """
+        with self._lock:
+            if not self._disposed:
+                self._disposed = True
+                self._observers.clear()
+                self._closed = True
+                self._error = None