def test_get_filesystem_s3(shutdown_only): root = "bucket/foo/bar" with simulate_storage("s3", root) as s3_uri: ray.init(storage=s3_uri) fs, prefix = storage.get_filesystem() assert path_eq(prefix, root), prefix assert isinstance(fs, pyarrow.fs.S3FileSystem), fs
def test_put_get(shutdown_only, tmp_path, storage_type): with simulate_storage(storage_type) as storage_uri: ray.init(storage=storage_uri, num_gpus=1) client = storage.get_client("ns") client2 = storage.get_client("ns2") assert client.get("foo/bar") is None client.put("foo/bar", b"hello") client.put("baz", b"goodbye") client2.put("baz", b"goodbye!") assert client.get("foo/bar") == b"hello" assert client.get("baz") == b"goodbye" assert client2.get("baz") == b"goodbye!" # delete file assert client.delete("baz") assert client.get("baz") is None assert not client.delete("non_existing") # delete dir n_files = 3 for i in range(n_files): assert client2.get(f"foo/bar{i}") is None for i in range(n_files): client2.put(f"foo/bar{i}", f"hello{i}".encode()) for i in range(n_files): assert client2.get(f"foo/bar{i}") == f"hello{i}".encode() assert client2.delete_dir("foo") for i in range(n_files): assert client2.get(f"foo/bar{i}") is None assert not client2.delete_dir("non_existing")
def test_external_kv_aws_s3(): with simulate_storage("s3", "serve-test") as uri: from urllib.parse import parse_qs, urlparse o = urlparse(uri) qs = parse_qs(o.query) region_name = qs["region"][0] endpoint_url = qs["endpoint_override"][0] import boto3 s3 = boto3.client( "s3", region_name=region_name, endpoint_url=endpoint_url, ) s3.create_bucket( Bucket="serve-test", CreateBucketConfiguration={"LocationConstraint": "us-west-2"}, ) kv_store = RayS3KVStore( "namespace", bucket="serve-test", prefix="checkpoint", region_name=region_name, endpoint_url=endpoint_url, ) _test_operations(kv_store)
def workflow_start_cluster(storage_type, request): # This code follows the design of "call_ray_start" fixture. with simulate_storage(storage_type) as storage_url: utils.clear_marks() parameter = getattr( request, "param", "ray start --head --num-cpus=1 --min-worker-port=0 " "--max-worker-port=0 --port 0 --storage=" + storage_url, ) command_args = parameter.split(" ") out = ray._private.utils.decode( subprocess.check_output(command_args, stderr=subprocess.STDOUT)) # Get the redis address from the output. address_prefix = "--address='" address_location = out.find(address_prefix) + len(address_prefix) address = out[address_location:] address = address.split("'")[0] yield address, storage_url # Disconnect from the Ray cluster. ray.shutdown() # Kill the Ray cluster. subprocess.check_call(["ray", "stop"])
def test_get_info_basic(shutdown_only, storage_type): with simulate_storage(storage_type) as storage_uri: ray.init(storage=storage_uri, num_gpus=1) client = storage.get_client("ns") client.put("foo/bar1", b"hello") assert client.get_info("foo/bar1").base_name == "bar1" assert client.get_info("foo/bar2") is None assert client.get_info("foo").base_name == "foo" assert client.get_info("").base_name == "ns"
def test_directory_traversal_attack(shutdown_only, storage_type): with simulate_storage(storage_type) as storage_uri: ray.init(storage=storage_uri, num_gpus=1) client = storage.get_client("foo") client.put("data", b"hello") client2 = storage.get_client("foo/bar") # Should not be able to access '../data'. with pytest.raises(ValueError): client2.get("../data")
def test_connecting_to_cluster(shutdown_only, storage_type): with simulate_storage(storage_type) as storage_uri: try: subprocess.check_call( ["ray", "start", "--head", "--storage", storage_uri]) ray.init(address="auto") from ray._private.storage import _storage_uri # make sure driver is using the same storage when connecting to a cluster assert _storage_uri == storage_uri finally: subprocess.check_call(["ray", "stop"])
def test_list_basic(shutdown_only, storage_type): with simulate_storage(storage_type) as storage_uri: ray.init(storage=storage_uri, num_gpus=1) client = storage.get_client("ns") client.put("foo/bar1", b"hello") client.put("foo/bar2", b"hello") client.put("baz/baz1", b"goodbye!") d1 = client.list("") assert sorted([f.base_name for f in d1]) == ["baz", "foo"], d1 d2 = client.list("foo") assert sorted([f.base_name for f in d2]) == ["bar1", "bar2"], d2 with pytest.raises(FileNotFoundError): client.list("invalid") with pytest.raises(NotADirectoryError): client.list("foo/bar1")
def workflow_start_regular_shared(storage_type, request): param = getattr(request, "param", {}) with simulate_storage(storage_type) as storage_url, _workflow_start( storage_url, True, **param) as res: yield res