feat(BA-684): Create Raftify Client #3743

racheliee · 2025-02-18T07:38:18Z

about

This PR introduces the raftify client that supports key-value store.

setup

The raft-cluster-config.toml should contain the information of the raft clusters as such:

[[peers.other]]
host = "127.0.0.1"
port = 60061
node-id = 1
role = "leader"

[[peers.myself]]
host = "127.0.0.1"
port = 60062
node-id = 2
role = "voter"

[[peers.other]]
host = "127.0.0.1"
port = 60063
node-id = 3
role = "voter"

testing

The key-value store mechanism can be tested as such:

curl -XGET http://localhost:6025{node-id}/put/1/test
curl -XGET http://localhost:6025{node-id}/get/1
curl -XGET http://localhost:6025{node-id}/leader

Checklist: (if applicable)

[ x ] Mention to the original issue
Installer updates including:
- Fixtures for db schema changes
- New mandatory config options
API server-client counterparts (e.g., manager API -> client SDK)
Test case(s) to:
- Demonstrate the difference of before/after
- Demonstrate the flow of abstract/conceptual models with a concrete implementation
Documentation
- Contents in the docs directory
- docstrings in public interfaces and type annotations

HyeockJinKim · 2025-02-20T09:38:13Z

configs/manager/halfstack.toml

@@ -83,3 +86,10 @@ certfile = ""

 [debug]
 enabled = false
+
+[raft-kvs]


Suggested change

[raft-kvs]

[raft]

HyeockJinKim · 2025-02-20T09:55:59Z

src/ai/backend/common/lock.py

+    async def __aenter__(self) -> Union[EtcdCommunicator, RaftKVSCommunicator]:
+        if isinstance(self.etcd.etcd, EtcdClient):
+            self._etcd_client = self.etcd.etcd.with_lock(
+                EtcdLockOption(
+                    lock_name=self.lock_name.encode("utf-8"),
+                    timeout=self._timeout,
+                    ttl=int(self._lifetime) if self._lifetime is not None else None,
+                ),
+            )
+        elif isinstance(self.etcd.etcd, RaftKVSClient):
+            self._etcd_client = await self.etcd.etcd.with_lock(
+                RaftKVSLockOptions(
+                    lock_name=self.lock_name.encode("utf-8"),
+                    timeout=self._timeout,
+                    ttl=int(self._lifetime) if self._lifetime is not None else None,
+                ),
+            )
+        assert self._etcd_client is not None
        etcd_communicator = await self._etcd_client.__aenter__()


class EtcdLock(AbstractDistributedLock): _etcd_client: Optional[Union[EtcdClient, RaftKVSClient]] _etcd: AbstractKVStore _etcd_locked: ETCDLocked _debug: bool lock_name: str etcd: AbstractKVStore timeout: float default_timeout: float = 9600 # not allow infinite timeout for safety def __init__( self, lock_name: str, etcd: AbstractKVStore, *, timeout: Optional[float] = None, lifetime: Optional[float] = None, debug: bool = False, ) -> None: super().__init__(lifetime=lifetime) self.lock_name = lock_name self.etcd = etcd self._timeout = timeout if timeout is not None else self.default_timeout self._debug = debug self._etcd_client = None async def __aenter__(self) -> EtcdCommunicator: if self._etcd_locked is not None: raise RuntimeError("Already locked") locked = self.etcd.with_lock(...) self._etcd_locked = locked etcd_communicator = await locked.__aenter__() if self._debug: log.debug("etcd lock acquired") return etcd_communicator async def __aexit__(self, *exc_info) -> Optional[bool]: if self._etcd_locked is None: raise RuntimeError("Not locked") await self._etcd_client.__aexit__(*exc_info) if self._debug: log.debug("etcd lock released") self._etcd_locked = None return None

HyeockJinKim · 2025-02-20T09:57:15Z

src/ai/backend/manager/cli/context.py

+        assert local_config["etcd"]["user"] is not None
+        assert local_config["etcd"]["password"] is not None


Use raise

HyeockJinKim · 2025-02-20T09:59:03Z

src/ai/backend/manager/raft/client.py

+
+    async def with_lock(
+        self, lock_options: RaftKVSLockOptions, connect_options: Optional["ConnectOptions"] = None
+    ) -> "RaftKVSClient":


Use Self in typing library.

HyeockJinKim · 2025-02-20T10:01:30Z

src/ai/backend/manager/raft/utils.py

+def pickle_deserialize(data: bytes) -> str | None:
+    if data == b"":
+        return None
+
+    if pickle.PROTO in data:
+        r = pickle.loads(data[data.index(pickle.PROTO) :])
+        return r
+
+    # Not pickle data
+    return None
+
+
+def register_raft_custom_deserializer() -> None:
+    """
+    Initialize the custom deserializers.
+    """
+
+    set_confchange_context_deserializer(pickle_deserialize)
+    set_confchangev2_context_deserializer(pickle_deserialize)
+    set_entry_context_deserializer(pickle_deserialize)
+    set_entry_data_deserializer(pickle_deserialize)
+    set_message_context_deserializer(pickle_deserialize)
+    set_snapshot_data_deserializer(pickle_deserialize)
+    set_log_entry_deserializer(pickle_deserialize)
+    set_fsm_deserializer(pickle_deserialize)


Why use pickle data?

@jopemachine

We can inject a deserializer used internally in raft-rs to deserialize byte slices used for logging purposes.
Besides pickle, other deserializers can also be injected.

If no deserializer is injected, the raw byte slice will be displayed as-is.

HyeockJinKim · 2025-02-20T10:29:44Z

src/ai/backend/manager/server.py

+        register_raft_custom_deserializer()
+
+        raft_configs = root_ctx.local_config.get("raft-kvs")
+        assert raft_configs is not None, "Raft configuration missing in the manager.toml"
+
+        raft_cluster_configs = root_ctx.raft_cluster_config
+        assert raft_cluster_configs is not None
+
+        other_peers = [{**peer, "myself": False} for peer in raft_cluster_configs["peers"]["other"]]
+        my_peers = [{**peer, "myself": True} for peer in raft_cluster_configs["peers"]["myself"]]
+        all_peers = sorted([*other_peers, *my_peers], key=lambda x: x["node-id"])
+
+        initial_peers = Peers({
+            int(peer_config["node-id"]): Peer(
+                addr=f"{peer_config['host']}:{peer_config['port']}",
+                role=InitialRole.from_str(peer_config["role"]),
+            )
+            for peer_config in all_peers
+        })
+
+        raft_core_config = RaftConfig(
+            heartbeat_tick=raft_configs["heartbeat-tick"],
+            election_tick=raft_configs["election-tick"],
+            min_election_tick=raft_configs["min-election-tick"],
+            max_election_tick=raft_configs["max-election-tick"],
+            max_committed_size_per_ready=raft_configs["max-committed-size-per-ready"],
+            max_size_per_msg=raft_configs["max-size-per-msg"],
+            max_inflight_msgs=raft_configs["max-inflight-msgs"],
+            check_quorum=raft_configs["check-quorum"],
+            batch_append=raft_configs["batch-append"],
+            max_uncommitted_size=raft_configs["max-uncommitted-size"],
+            skip_bcast_commit=raft_configs["skip-bcast-commit"],
+            pre_vote=raft_configs["pre-vote"],
+            priority=raft_configs["priority"],
+        )
+
+        raft_cfg = Config(
+            log_dir=raft_configs["log-dir"],
+            save_compacted_logs=True,
+            compacted_log_dir=raft_configs["log-dir"],
+            restore_wal_from=raft_cluster_configs["restore-wal-from"],
+            restore_wal_snapshot_from=raft_cluster_configs["restore-wal-snapshot-from"],
+            initial_peers=initial_peers,
+            raft_config=raft_core_config,
+        )
+
+        node_id_offset = next((idx for idx, item in enumerate(all_peers) if item["myself"]), None)
+        assert node_id_offset is not None, '"peers.myself" not found in initial_peers!'
+        node_id = node_id_offset + aiotools.process_index.get() + 1
+
+        raft_addr = initial_peers.get(node_id).get_addr()
+
+        store = RaftHashStore()
+
+        raft_logger = RaftLogger(
+            logging.getLogger(f"{__spec__.name}.raft.node-{node_id}"),  # type: ignore
+        )
+
+        root_ctx.kvstore_ctx.raft = Raft.bootstrap(
+            node_id,
+            raft_addr,
+            store,  # type: ignore
+            raft_cfg,
+            raft_logger,  # type: ignore
+        )
+        raft_cluster = root_ctx.kvstore_ctx.raft
+        raft_cluster.run()  # type: ignore


In my opinion, raft setup should be separated to another setup context function.
example:

@actxmgr async def raft_ctx(root_ctx: RootContext) -> AsyncIterator[None]: ... @actxmgr async def etcd_shared_config_ctx(root_ctx: RootContext) -> AsyncIterator[None]: ... @actxmgr async def raft_shared_config_ctx(root_ctx: RootContext) -> AsyncIterator[None]: ...

HyeockJinKim · 2025-02-26T12:56:04Z

src/ai/backend/common/etcd.py

    """
    Abstract interface for Key-Value Store (KVS) operations
    Defines the basic operations that a KVS should support
    """

+    etcd: T


In my opinion, it seems better to change etcd to private field of EtcdClient and provide interface methods that are needed for use.

You don't need to use Generic then.

HyeockJinKim · 2025-02-26T12:57:37Z

src/ai/backend/common/etcd.py

+        self.scope_prefix_map = t.Dict({
+            t.Key(ConfigScopes.GLOBAL): t.String(allow_blank=True),
+            t.Key(ConfigScopes.SGROUP, optional=True): t.String,
+            t.Key(ConfigScopes.NODE, optional=True): t.String,
+        }).check(scope_prefix_map)


It seems that there's no problem even with only private fields, not public fields.

HyeockJinKim · 2025-02-26T13:02:09Z

src/ai/backend/manager/BUILD

@@ -15,6 +15,7 @@ python_sources(
 visibility_private_component(
    allowed_dependents=[
        "//src/ai/backend/testutils/**",
+        "//src/ai/backend/common/**",


HyeockJinKim · 2025-02-26T13:04:36Z

src/ai/backend/manager/api/context.py

@@ -29,6 +33,38 @@ class BaseContext:
    pass


+class KVStoreKind(CIStrEnum):


Can't we just receive lowercase values for configuration?

we could but i felt that using the enum would make it safer

What I meant was that StrEnum seems to be sufficient, rather than CIStrEnum.

ohh i see gotcha

HyeockJinKim · 2025-02-26T13:07:36Z

src/ai/backend/manager/raft/client.py

+
+class ConnectOptions:
+    def __init__(self) -> None: ...
+    def with_user(self, user: str, password: str) -> "ConnectOptions":


User Self in typing package

HyeockJinKim · 2025-02-26T13:08:05Z

src/ai/backend/manager/raft/client.py

+        self,
+        raft_node: RaftNode,
+        endpoints: list[str],
+        connect_options: Optional["ConnectOptions"] = None,


Why use "ConnectOptions"?

HyeockJinKim · 2025-02-26T13:09:14Z

src/ai/backend/manager/raft/client.py

+        self.raft_node = raft_node
+        self.endpoints = endpoints
+        self.connect_options = connect_options


I would prefer it if internal fields were not accessed from outside as much as possible, but why are this fields used as public?

HyeockJinKim · 2025-02-26T13:10:46Z

src/ai/backend/manager/server.py

+@click.option(
+    "--raft-cluster-config-path",
+    "--raft-cluster-config",
+    type=Path,
+    default=None,
+    help="The raft cluster config file path. (default: ./raft-cluster-config.toml and /etc/backend.ai/raft-cluster-config.toml)",
+)


Can't we only inject it from a configuration file without inputting it in CLI?

it is injected from the configuration file via the default path if no configuration file is specified. this is just an option to specify the raft configuration file path if you'd like to use a file from a different place

Wouldn’t it be better to receive the raft config path from the manager’s config file instead of a CLI option?

racheliee added 5 commits February 14, 2025 14:07

init

ab74973

bootstrap?

2fe5bda

allow common???

0f7753b

some fixes

bb28b5a

type type type

7d7e168

github-actions bot assigned racheliee Feb 18, 2025

fix?

cdf1d6c

racheliee changed the title ~~feat(BA-686): Implement Raftify KVS~~ feat(BA-684): Create Raftify Client Feb 18, 2025

racheliee added 4 commits February 18, 2025 17:54

toml is such a cute word to say

d246b6a

client??? please work

457f4a7

easter for manager!

982a9ce

fix etcd context

ca1e437

HyeockJinKim reviewed Feb 20, 2025

View reviewed changes

racheliee added 6 commits February 21, 2025 11:43

change key to raft

93bc292

minimize pr

0226c92

include raft node intialization

3408bef

check put/get/leader done

717231b

clean up server

7c7342f

remove pickle

cf74e38

racheliee marked this pull request as ready for review February 21, 2025 08:56

changes

a664376

HyeockJinKim reviewed Feb 26, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat(BA-684): Create Raftify Client #3743

feat(BA-684): Create Raftify Client #3743

racheliee commented Feb 18, 2025 •

edited

Loading

HyeockJinKim Feb 20, 2025

HyeockJinKim Feb 20, 2025

HyeockJinKim Feb 20, 2025

HyeockJinKim Feb 20, 2025

HyeockJinKim Feb 20, 2025

HyeockJinKim Feb 20, 2025

jopemachine Feb 20, 2025 •

edited

Loading

HyeockJinKim Feb 20, 2025

HyeockJinKim Feb 26, 2025

HyeockJinKim Feb 26, 2025

HyeockJinKim Feb 26, 2025

HyeockJinKim Feb 26, 2025

HyeockJinKim Feb 26, 2025

racheliee Feb 27, 2025

HyeockJinKim Feb 27, 2025

racheliee Feb 27, 2025

HyeockJinKim Feb 26, 2025

HyeockJinKim Feb 26, 2025

HyeockJinKim Feb 26, 2025

HyeockJinKim Feb 26, 2025

racheliee Feb 27, 2025

HyeockJinKim Feb 27, 2025

		assert local_config["etcd"]["user"] is not None
		assert local_config["etcd"]["password"] is not None

		@@ -29,6 +33,38 @@ class BaseContext:
		pass


		class KVStoreKind(CIStrEnum):

feat(BA-684): Create Raftify Client #3743

Are you sure you want to change the base?

feat(BA-684): Create Raftify Client #3743

Conversation

racheliee commented Feb 18, 2025 • edited Loading

about

setup

testing

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jopemachine Feb 20, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

racheliee commented Feb 18, 2025 •

edited

Loading

jopemachine Feb 20, 2025 •

edited

Loading