Skip to content

Commit

Permalink
fix b/w compat chitchat implementation (#1109)
Browse files Browse the repository at this point in the history
  • Loading branch information
vangheem committed Jul 18, 2023
1 parent c0e1f1f commit 395577d
Show file tree
Hide file tree
Showing 2 changed files with 47 additions and 1 deletion.
40 changes: 40 additions & 0 deletions nucliadb/nucliadb/common/cluster/discovery/chitchat.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,10 +20,12 @@
from __future__ import annotations

import asyncio
from enum import Enum
from typing import Optional

import pydantic
from fastapi import APIRouter, FastAPI, Response
from nucliadb_protos.writer_pb2 import Member
from uvicorn.config import Config # type: ignore
from uvicorn.server import Server # type: ignore

Expand All @@ -39,10 +41,47 @@
api_router = APIRouter()


class MemberType(str, Enum):
IO = "Io"
SEARCH = "Search"
INGEST = "Ingest"
TRAIN = "Train"
UNKNOWN = "Unknown"

@staticmethod
def from_pb(node_type: Member.Type.ValueType):
if node_type == Member.Type.IO:
return MemberType.IO
elif node_type == Member.Type.SEARCH:
return MemberType.SEARCH
elif node_type == Member.Type.INGEST:
return MemberType.INGEST
elif node_type == Member.Type.TRAIN:
return MemberType.TRAIN
elif node_type == Member.Type.UNKNOWN:
return MemberType.UNKNOWN
else:
raise ValueError(f"incompatible node type '{node_type}'")

def to_pb(self) -> Member.Type.ValueType:
if self == MemberType.IO:
return Member.Type.IO
elif self == MemberType.SEARCH:
return Member.Type.SEARCH
elif self == MemberType.INGEST:
return Member.Type.INGEST
elif self == MemberType.TRAIN:
return Member.Type.TRAIN
else:
return Member.Type.UNKNOWN


class ClusterMember(pydantic.BaseModel):
node_id: str = pydantic.Field(alias="id")
listen_addr: str = pydantic.Field(alias="address")
shard_count: Optional[int]
type: MemberType = MemberType.UNKNOWN
is_self: bool = False

class Config:
allow_population_by_field_name = True
Expand All @@ -59,6 +98,7 @@ async def api_update_members(members: list[ClusterMember]) -> Response:
shard_count=member.shard_count or 0,
)
for member in members
if not member.is_self and member.type == MemberType.IO
]
)
return Response(status_code=204)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,13 @@ def make_client_fixture():
async def test_chitchat_monitor(chitchat_monitor_client):
INDEX_NODES.clear()
async with chitchat_monitor_client() as client:
member = dict(node_id=f"node", listen_addr=f"10.0.0.0", shard_count=20)
member = dict(
node_id=f"node",
listen_addr=f"10.0.0.0",
shard_count=20,
type="Io",
is_self=False,
)
response = await client.patch("/members", json=[member])
assert response.status_code == 204
assert len(INDEX_NODES) == 1

1 comment on commit 395577d

@github-actions
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Benchmark

Benchmark suite Current: 395577d Previous: a40981d Ratio
nucliadb/tests/benchmarks/test_search.py::test_search_returns_labels 80.5513816806325 iter/sec (stddev: 0.000437323264874243) 60.779932309336715 iter/sec (stddev: 0.0019119907918232523) 0.75
nucliadb/tests/benchmarks/test_search.py::test_search_relations 179.68322884175103 iter/sec (stddev: 0.0005850892900654587) 182.57436721258293 iter/sec (stddev: 0.0002220745559283828) 1.02

This comment was automatically generated by workflow using github-action-benchmark.

Please sign in to comment.