chore: import upstream snapshot with attribution

python/ray/experimental/internal_kv.py

@@ -0,0 +1,127 @@
+from typing import List, Optional, Union
+
+from ray._private.client_mode_hook import client_mode_hook
+from ray._raylet import GcsClient
+
+_initialized = False
+global_gcs_client = None
+
+
+def _internal_kv_reset():
+    global global_gcs_client, _initialized
+    global_gcs_client = None
+    _initialized = False
+
+
+def internal_kv_get_gcs_client():
+    return global_gcs_client
+
+
+def _initialize_internal_kv(gcs_client: GcsClient):
+    """Initialize the internal KV for use in other function calls."""
+    global global_gcs_client, _initialized
+    assert gcs_client is not None
+    global_gcs_client = gcs_client
+    _initialized = True
+
+
+@client_mode_hook
+def _internal_kv_initialized():
+    return global_gcs_client is not None
+
+
+@client_mode_hook
+def _internal_kv_get(
+    key: Union[str, bytes], *, namespace: Optional[Union[str, bytes]] = None
+) -> bytes:
+    """Fetch the value of a binary key."""
+
+    if isinstance(key, str):
+        key = key.encode()
+    if isinstance(namespace, str):
+        namespace = namespace.encode()
+    assert isinstance(key, bytes)
+    return global_gcs_client.internal_kv_get(key, namespace)
+
+
+@client_mode_hook
+def _internal_kv_exists(
+    key: Union[str, bytes], *, namespace: Optional[Union[str, bytes]] = None
+) -> bool:
+    """Check key exists or not."""
+
+    if isinstance(key, str):
+        key = key.encode()
+    if isinstance(namespace, str):
+        namespace = namespace.encode()
+    assert isinstance(key, bytes)
+    return global_gcs_client.internal_kv_exists(key, namespace)
+
+
+@client_mode_hook
+def _pin_runtime_env_uri(uri: str, *, expiration_s: int) -> None:
+    """Pin a runtime_env URI for expiration_s."""
+    return global_gcs_client.pin_runtime_env_uri(uri, expiration_s)
+
+
+@client_mode_hook
+def _internal_kv_put(
+    key: Union[str, bytes],
+    value: Union[str, bytes],
+    overwrite: bool = True,
+    *,
+    namespace: Optional[Union[str, bytes]] = None
+) -> bool:
+    """Globally associates a value with a given binary key.
+
+    This only has an effect if the key does not already have a value.
+
+    Args:
+        key: The binary key to associate the value with.
+        value: The binary value to store under the key.
+        overwrite: Whether to overwrite an existing value for the key.
+        namespace: Optional namespace under which the key is scoped.
+
+    Returns:
+        Whether the value already exists.
+    """
+
+    if isinstance(key, str):
+        key = key.encode()
+    if isinstance(value, str):
+        value = value.encode()
+    if isinstance(namespace, str):
+        namespace = namespace.encode()
+    assert (
+        isinstance(key, bytes)
+        and isinstance(value, bytes)
+        and isinstance(overwrite, bool)
+    )
+    return global_gcs_client.internal_kv_put(key, value, overwrite, namespace) == 0
+
+
+@client_mode_hook
+def _internal_kv_del(
+    key: Union[str, bytes],
+    *,
+    del_by_prefix: bool = False,
+    namespace: Optional[Union[str, bytes]] = None
+) -> int:
+    if isinstance(key, str):
+        key = key.encode()
+    if isinstance(namespace, str):
+        namespace = namespace.encode()
+    assert isinstance(key, bytes)
+    return global_gcs_client.internal_kv_del(key, del_by_prefix, namespace)
+
+
+@client_mode_hook
+def _internal_kv_list(
+    prefix: Union[str, bytes], *, namespace: Optional[Union[str, bytes]] = None
+) -> List[bytes]:
+    """List all keys in the internal KV store that start with the prefix."""
+    if isinstance(prefix, str):
+        prefix = prefix.encode()
+    if isinstance(namespace, str):
+        namespace = namespace.encode()
+    return global_gcs_client.internal_kv_keys(prefix, namespace)