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 test_external_kv_aws_s3():
kv_store = RayS3KVStore(
"namespace",
bucket="jiao-test",
s3_path="/checkpoint",
aws_access_key_id=os.environ.get("AWS_ACCESS_KEY_ID", None),
aws_secret_access_key=os.environ.get("AWS_SECRET_ACCESS_KEY", None),
aws_session_token=os.environ.get("AWS_SESSION_TOKEN", None),
)
_test_operations(kv_store)
def make_kv_store(checkpoint_path, namespace):
"""Create KVStore instance based on checkpoint_path configuration"""
if checkpoint_path == DEFAULT_CHECKPOINT_PATH:
logger.info("Using RayInternalKVStore for controller "
"checkpoint and recovery.")
return RayInternalKVStore(namespace)
else:
parsed_url = urlparse(checkpoint_path)
if parsed_url.scheme not in {"s3", "file", "custom"}:
raise ValueError(
f"Checkpoint must be one of `{DEFAULT_CHECKPOINT_PATH}`, "
"`file://path...`, `s3://path...` or "
"`custom://my_module.ClassName?arg1=val1`. But it is "
f"{checkpoint_path}")
if parsed_url.scheme == "file":
db_path = parsed_url.netloc + parsed_url.path
logger.info("Using RayLocalKVStore for controller "
f"checkpoint and recovery: path={db_path}")
return RayLocalKVStore(namespace, db_path)
if parsed_url.scheme == "s3":
bucket = parsed_url.netloc
# We need to strip leading "/" in path as right key to use in
# boto3. Ex: s3://bucket/folder/file.zip -> key = "folder/file.zip"
prefix = parsed_url.path.lstrip("/")
logger.info(
"Using Ray S3 KVStore for controller checkpoint and recovery: "
f"bucket={bucket} checkpoint_path={checkpoint_path}")
return RayS3KVStore(
namespace,
bucket=bucket,
prefix=prefix,
)
if parsed_url.scheme == "custom":
kwargs = dict(parse_qsl(parsed_url.query))
# Prepare the parameters to initialize imported class.
checkpoint_provider = parsed_url.netloc
KVStoreClass = import_attr(checkpoint_provider)
if not issubclass(KVStoreClass, KVStoreBase):
raise ValueError(
f"{KVStoreClass} doesn't inherit from "
"`ray.serve.storage.kv_store_base.KVStoreBase`.")
logger.info(
f"Using {checkpoint_provider} for controller checkpoint and "
f"recovery: kwargs={kwargs}")
return KVStoreClass(namespace=namespace, **kwargs)
raise RuntimeError("This shouldn't be reachable.")