From 7f342eb3714be7c847a21b01a798bab8b9f5ba59 Mon Sep 17 00:00:00 2001 From: Eric Liang Date: Tue, 9 Feb 2021 20:47:41 -0800 Subject: [PATCH] Update example shuffle script (#14021) --- python/ray/experimental/shuffle.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/python/ray/experimental/shuffle.py b/python/ray/experimental/shuffle.py index 6b7936ddf..0a3f01656 100644 --- a/python/ray/experimental/shuffle.py +++ b/python/ray/experimental/shuffle.py @@ -169,8 +169,10 @@ def main(): parser.add_argument("--partition-size", type=float, default=200e6) args = parser.parse_args() - ray.init( - address=args.ray_address, object_store_memory=args.object_store_memory) + if args.ray_address: + ray.init(address=args.ray_address) + else: + ray.init(object_store_memory=args.object_store_memory) partition_size = int(args.partition_size) num_partitions = args.num_partitions