chore: import upstream snapshot with attribution

python/ray/tests/chaos/potato_passer.py

@@ -0,0 +1,68 @@
+import argparse
+import asyncio
+
+import ray
+
+ray.init()
+
+"""
+Potato passer is a test script that lets multiple actors call each other's methods.
+Actors are wired in a round-trip fashion: actor 0 calls actor 1, which calls actor 2.
+The last actor calls actor 0. In each call, the actor sleeps for a time, occationally
+prints, and calls next actor.
+
+Note the number of tasks on-the-fly can go up to `pass-times` because the next call is
+made before exiting current call.
+"""
+
+
+@ray.remote
+class PotatoPasser:
+    def __init__(self, name, next_name, sleep_secs):
+        self.count = 0
+        self.name = name
+        self.next_name = next_name
+        self.sleep_secs = sleep_secs
+        self.print_every = 100
+
+    async def pass_potato(self, potato: int, target: int):
+        self.count += 1
+        if potato % self.print_every == 0:
+            print(
+                f"running, name {self.name}, count {self.count}, "
+                f"potato {potato}, target {target}"
+            )
+        if potato >= target:
+            print(f"target reached! name = {self.name}, count = {self.count}")
+            return target
+        next_actor = ray.get_actor(self.next_name)
+        await asyncio.sleep(self.sleep_secs)
+        return await next_actor.pass_potato.remote(potato + 1, target)
+
+
+async def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--num-actors", type=int, help="Make this many actors")
+    parser.add_argument("--pass-times", type=int, help="Pass this many messages")
+    parser.add_argument(
+        "--sleep-secs",
+        type=float,
+        help="Sleep seconds before sending message to next actor",
+    )
+    args = parser.parse_args()
+
+    actors = []
+    for i in range(args.num_actors):
+        this_actor = "actor" + str(i)
+        next_actor = "actor" + str((i + 1) % args.num_actors)
+        actor = PotatoPasser.options(
+            name=this_actor, scheduling_strategy="SPREAD"
+        ).remote(this_actor, next_actor, args.sleep_secs)
+        actors.append(actor)
+
+    ret = await actors[0].pass_potato.remote(0, args.pass_times)
+    print(f"passed potato {ret} times! expected {args.pass_times} times.")
+    assert ret == args.pass_times
+
+
+asyncio.run(main())