diff --git a/release/nightly_tests/dask_on_ray/dask_on_ray_sort.py b/release/nightly_tests/dask_on_ray/dask_on_ray_sort.py index a512c4055..92999e61a 100644 --- a/release/nightly_tests/dask_on_ray/dask_on_ray_sort.py +++ b/release/nightly_tests/dask_on_ray/dask_on_ray_sort.py @@ -53,7 +53,19 @@ def load_dataset(client, data_dir, s3_bucket, nbytes, npartitions): f"s3://{s3_bucket}/df-{num_bytes_per_partition}-{i}.parquet.gzip" for i in range(npartitions) ] - df = dd.read_parquet(filenames) + + df = None + max_retry = 3 + retry = 0 + while not df and retry < max_retry: + try: + df = dd.read_parquet(filenames) + except FileNotFoundError as e: + print(f"Failed to load a file. {e}") + # Wait a little bit before retrying. + time.sleep(30) + retry += 1 + return df