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