def test_external_kv_gcs(): kv_store = RayGcsKVStore( "namespace", bucket="jiao-test", prefix="/checkpoint", ) _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 {"gs", "s3", "file", "custom"}: raise ValueError( f"Checkpoint must be one of `{DEFAULT_CHECKPOINT_PATH}`, " "`file://path...`, `gs://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 == "gs": bucket = parsed_url.netloc # We need to strip leading "/" in path as right key to use in # gcs. Ex: gs://bucket/folder/file.zip -> key = "folder/file.zip" prefix = parsed_url.path.lstrip("/") logger.info("Using Ray GCS KVStore for controller checkpoint and" " recovery: " f"bucket={bucket} checkpoint_path={checkpoint_path}") return RayGcsKVStore( namespace, bucket=bucket, prefix=prefix, ) 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.")