wehub-resource-sync
152 lines
5.6 KiB
Python
152 lines
5.6 KiB
Python
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
|