chore: import upstream snapshot with attribution

python/ray/tests/test_client_warnings.py

@@ -0,0 +1,38 @@
+import sys
+import unittest
+
+import numpy as np
+import pytest
+
+from ray.util.client.ray_client_helpers import ray_start_client_server
+from ray.util.debug import _logged
+
+
+@pytest.fixture(autouse=True)
+def reset_debug_logs():
+    """Resets internal state of ray.util.debug so that warning can be tested
+    more than once in the same process"""
+    _logged.clear()
+
+
+class LoggerSuite(unittest.TestCase):
+    """Test client warnings are raised when many tasks are scheduled"""
+
+    def testOutboundMessageSizeWarning(self):
+        with self.assertWarns(UserWarning) as cm, ray_start_client_server() as ray:
+            large_argument = np.random.rand(100, 100, 100)
+
+            @ray.remote
+            def f(some_arg):
+                return some_arg[0][0][0]
+
+            for _ in range(50):
+                f.remote(large_argument)
+        assert (
+            "More than 10MB of messages have been created to "
+            "schedule tasks on the server." in cm.warning.args[0]
+        )
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))