wehub-resource-sync
78 lines
2.2 KiB
Python
78 lines
2.2 KiB
Python
# Regular ray application that user wrote and runs on local cluster.
|
|
# intermediate status are dumped to GCS
|
|
import argparse
|
|
import time
|
|
|
|
import ray
|
|
import ray.experimental.internal_kv as ray_kv
|
|
|
|
|
|
@ray.remote
|
|
class StepActor:
|
|
def __init__(self, interval_s=1, total_steps=3):
|
|
self.interval_s = interval_s
|
|
self.stopped = False
|
|
self.current_step = 1
|
|
self.total_steps = total_steps
|
|
|
|
worker = ray._private.worker.global_worker
|
|
worker_id = worker.core_worker.get_actor_id()
|
|
ray_kv._internal_kv_put(f"JOB:{worker_id}", self.current_step, overwrite=True)
|
|
|
|
def run(self):
|
|
worker = ray._private.worker.global_worker
|
|
worker_id = worker.core_worker.get_actor_id()
|
|
|
|
while self.current_step <= self.total_steps:
|
|
if not self.stopped:
|
|
print(
|
|
f"Sleeping {self.interval_s} secs to executing "
|
|
f"step {self.current_step}"
|
|
)
|
|
time.sleep(self.interval_s)
|
|
self.current_step += 1
|
|
ray_kv._internal_kv_put(
|
|
f"JOB:{worker_id}", self.current_step, overwrite=True
|
|
)
|
|
else:
|
|
print("Stop called or reached final step.")
|
|
break
|
|
|
|
self.stopped = True
|
|
ray_kv._internal_kv_put(f"JOB:{worker_id}", "DONE", overwrite=True)
|
|
return "DONE"
|
|
|
|
def get_step(self):
|
|
return self.current_step
|
|
|
|
def stop(self):
|
|
self.stopped = True
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument(
|
|
"--interval-s",
|
|
required=False,
|
|
type=int,
|
|
default=1,
|
|
help="Address to use to connect to Ray",
|
|
)
|
|
parser.add_argument(
|
|
"--total-steps",
|
|
required=False,
|
|
type=int,
|
|
default=3,
|
|
help="Password for connecting to Redis",
|
|
)
|
|
args, _ = parser.parse_known_args()
|
|
|
|
ray.init()
|
|
step_actor = StepActor.remote(
|
|
interval_s=args.interval_s, total_steps=args.total_steps
|
|
)
|
|
ref = step_actor.run.remote()
|
|
print(ray.get([ref]))
|
|
job_key = ray_kv._internal_kv_list("JOB:")[0]
|
|
print(f"{job_key}, {ray_kv._internal_kv_get(job_key)}")
|