wehub-resource-sync
69 lines
2.2 KiB
Python
69 lines
2.2 KiB
Python
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())
|