chore: import upstream snapshot with attribution

python/ray/autoscaler/v2/instance_manager/instance_storage.py

@@ -0,0 +1,151 @@
+import copy
+import logging
+from typing import Dict, List, Optional, Set, Tuple
+
+from ray.autoscaler.v2.instance_manager.storage import Storage, StoreStatus
+from ray.core.generated.instance_manager_pb2 import Instance
+
+logger = logging.getLogger(__name__)
+
+
+class InstanceStorage:
+    """Instance storage stores the states of instances in the storage."""
+
+    def __init__(
+        self,
+        cluster_id: str,
+        storage: Storage,
+    ) -> None:
+        self._storage = storage
+        self._cluster_id = cluster_id
+        self._table_name = f"instance_table@{cluster_id}"
+
+    def batch_upsert_instances(
+        self,
+        updates: List[Instance],
+        expected_storage_version: Optional[int] = None,
+    ) -> StoreStatus:
+        """Upsert instances into the storage. If the instance already exists,
+        it will be updated. Otherwise, it will be inserted. If the
+        expected_storage_version is specified, the update will fail if the
+        current storage version does not match the expected version.
+
+        Note the version of the upserted instances will be set to the current
+        storage version.
+
+        Args:
+            updates: A list of instances to be upserted.
+            expected_storage_version: The expected storage version.
+
+        Returns:
+            StoreStatus: A tuple of (success, storage_version).
+        """
+        mutations = {}
+        version = self._storage.get_version()
+        # handle version mismatch
+        if expected_storage_version and expected_storage_version != version:
+            return StoreStatus(False, version)
+
+        for instance in updates:
+            instance = copy.deepcopy(instance)
+            # the instance version is set to 0, it will be
+            # populated by the storage entry's verion on read
+            instance.version = 0
+            mutations[instance.instance_id] = instance.SerializeToString()
+
+        result, version = self._storage.batch_update(
+            self._table_name, mutations, {}, expected_storage_version
+        )
+
+        return StoreStatus(result, version)
+
+    def upsert_instance(
+        self,
+        instance: Instance,
+        expected_instance_version: Optional[int] = None,
+        expected_storage_version: Optional[int] = None,
+    ) -> StoreStatus:
+        """Upsert an instance in the storage.
+        If the expected_instance_version is specified, the update will fail
+        if the current instance version does not match the expected version.
+        Similarly, if the expected_storage_version is
+        specified, the update will fail if the current storage version does not
+        match the expected version.
+
+        Note the version of the upserted instances will be set to the current
+        storage version.
+
+        Args:
+            instance: The instance to be updated.
+            expected_instance_version: The expected instance version.
+            expected_storage_version: The expected storage version.
+
+        Returns:
+            StoreStatus: A tuple of (success, storage_version).
+        """
+        instance = copy.deepcopy(instance)
+        # the instance version is set to 0, it will be
+        # populated by the storage entry's verion on read
+        instance.version = 0
+        result, version = self._storage.update(
+            self._table_name,
+            key=instance.instance_id,
+            value=instance.SerializeToString(),
+            expected_entry_version=expected_instance_version,
+            expected_storage_version=expected_storage_version,
+            insert_only=False,
+        )
+
+        return StoreStatus(result, version)
+
+    def get_instances(
+        self,
+        instance_ids: List[str] = None,
+        status_filter: Set[int] = None,
+    ) -> Tuple[Dict[str, Instance], int]:
+        """Get instances from the storage.
+
+        Args:
+            instance_ids: A list of instance ids to be retrieved. If empty, all
+                instances will be retrieved.
+            status_filter: Only instances with the specified status will be returned.
+
+        Returns:
+            Tuple[Dict[str, Instance], int]: A tuple of (instances, version).
+                The instances is a dictionary of (instance_id, instance) pairs.
+        """
+        instance_ids = instance_ids or []
+        status_filter = status_filter or set()
+        pairs, version = self._storage.get(self._table_name, instance_ids)
+        instances = {}
+        for instance_id, (instance_data, entry_version) in pairs.items():
+            instance = Instance()
+            instance.ParseFromString(instance_data)
+            instance.version = entry_version
+            if status_filter and instance.status not in status_filter:
+                continue
+            instances[instance_id] = instance
+        return instances, version
+
+    def batch_delete_instances(
+        self, instance_ids: List[str], expected_storage_version: Optional[int] = None
+    ) -> StoreStatus:
+        """Delete instances from the storage. If the expected_storage_version
+        is specified, the update will fail if the current storage version does
+        not match the expected version.
+
+        Args:
+            instance_ids: A list of instance ids to be deleted.
+            expected_storage_version: The expected storage version.
+
+        Returns:
+            StoreStatus: A tuple of (success, storage_version).
+        """
+        version = self._storage.get_version()
+        if expected_storage_version and expected_storage_version != version:
+            return StoreStatus(False, version)
+
+        result = self._storage.batch_update(
+            self._table_name, {}, instance_ids, expected_storage_version
+        )
+        return result