chore: import upstream snapshot with attribution

python/ray/_private/async_utils.py

@@ -0,0 +1,65 @@
+# Adapted from [aiodebug](https://gitlab.com/quantlane/libs/aiodebug)
+
+# Copyright 2016-2022 Quantlane s.r.o.
+
+#    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.
+
+# Modifications:
+# - Removed the dependency to `logwood`.
+# - Renamed `monitor_loop_lag.enable()` to just `enable_monitor_loop_lag()`.
+# - Miscellaneous changes to make it work with Ray.
+
+import asyncio
+import asyncio.events
+from typing import Callable, Optional, Set
+
+# Strong references to background tasks to prevent them from being garbage
+# collected mid-execution. The asyncio event loop only keeps weak references
+# to tasks, so a task with no other strong references can be collected at
+# any time. See https://docs.python.org/3/library/asyncio-task.html#asyncio.create_task
+_BACKGROUND_TASKS: Set[asyncio.Task] = set()
+
+
+def enable_monitor_loop_lag(
+    callback: Callable[[float], None],
+    interval_s: float = 0.25,
+    loop: Optional[asyncio.AbstractEventLoop] = None,
+) -> None:
+    """Start logging event loop lags to the callback.
+
+    In ideal circumstances they should be very close to zero. Lags may increase
+    if event loop callbacks block for too long.
+
+    Note: this works for all event loops, including uvloop.
+
+    Args:
+        callback: Callback to call with the lag in seconds.
+        interval_s: How often to measure the lag, in seconds.
+        loop: The event loop to monitor. If None, uses the running loop.
+    """
+    if loop is None:
+        loop = asyncio.get_running_loop()
+    if loop is None:
+        raise ValueError("No provided loop, nor running loop found.")
+
+    async def monitor():
+        while loop.is_running():
+            t0 = loop.time()
+            await asyncio.sleep(interval_s)
+            lag = loop.time() - t0 - interval_s  # Should be close to zero.
+            callback(lag)
+
+    task = loop.create_task(monitor(), name="async_utils.monitor_loop_lag")
+    # Keep a strong reference so the task isn't garbage collected mid-execution.
+    _BACKGROUND_TASKS.add(task)
+    task.add_done_callback(_BACKGROUND_TASKS.discard)