Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Kafka producer should raise an exception when it fails to connect to broker #636

Merged
merged 4 commits into from
May 14, 2020
Merged

Kafka producer should raise an exception when it fails to connect to broker #636

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
57cb981
Added exception to flush when produce fails with unit test or with me…
junhui096 Apr 15, 2020
ac6dbf2
Fix: Throw exception in callback
junhui096 May 6, 2020
b56747a
Removed error_count property in abstract_producer
junhui096 May 11, 2020
040deaf
Merge branch 'master' into exception_handling
woop May 13, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 10 additions & 24 deletions sdk/python/feast/loaders/abstract_producer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,8 +25,6 @@ class AbstractProducer:
def __init__(self, brokers: str, row_count: int, disable_progress_bar: bool):
self.brokers = brokers
self.row_count = row_count
self.error_count = 0
self.last_exception = ""

# Progress bar will always display average rate
self.pbar = tqdm(
Expand All @@ -45,8 +43,7 @@ def _inc_pbar(self, meta):
self.pbar.update(1)

def _set_error(self, exception: str):
self.error_count += 1
self.last_exception = exception
raise Exception(exception)

def print_results(self) -> None:
"""
Expand All @@ -62,24 +59,7 @@ def print_results(self) -> None:

print("Ingestion complete!")

failed_message = (
""
if self.error_count == 0
else f"\nFail: {self.error_count / self.row_count}"
)

last_exception_message = (
""
if self.last_exception == ""
else f"\nLast exception:\n{self.last_exception}"
)

print(
f"\nIngestion statistics:"
f"\nSuccess: {self.pbar.n}/{self.row_count}"
f"{failed_message}"
f"{last_exception_message}"
)
print(f"\nIngestion statistics:" f"\nSuccess: {self.pbar.n}/{self.row_count}")
return None


Expand Down Expand Up @@ -129,7 +109,10 @@ def flush(self, timeout: Optional[int]):
Returns:
int: Number of messages still in queue.
"""
return self.producer.flush(timeout=timeout)
messages = self.producer.flush(timeout=timeout)
Copy link
Collaborator

Choose a reason for hiding this comment

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

Might be better to fail fast using produce() callbacks rather than wait for the timeout/queue to clear.

if messages:
raise Exception("Not all Kafka messages are successfully delivered.")
return messages

def _delivery_callback(self, err: str, msg) -> None:
"""
Expand Down Expand Up @@ -200,7 +183,10 @@ def flush(self, timeout: Optional[int]):
KafkaTimeoutError: failure to flush buffered records within the
provided timeout
"""
return self.producer.flush(timeout=timeout)
messages = self.producer.flush(timeout=timeout)
if messages:
raise Exception("Not all Kafka messages are successfully delivered.")
return messages


def get_producer(
Expand Down
32 changes: 32 additions & 0 deletions sdk/python/tests/test_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -601,6 +601,38 @@ def test_feature_set_ingest_success(self, dataframe, test_client, mocker):
# Ingest data into Feast
test_client.ingest("driver-feature-set", dataframe)

@pytest.mark.parametrize(
"dataframe,test_client,exception",
[(dataframes.GOOD, pytest.lazy_fixture("client"), Exception)],
)
def test_feature_set_ingest_throws_exception_if_kafka_down(
self, dataframe, test_client, exception, mocker
):

test_client.set_project("project1")
driver_fs = FeatureSet(
"driver-feature-set",
source=KafkaSource(brokers="localhost:4412", topic="test"),
)
driver_fs.add(Feature(name="feature_1", dtype=ValueType.FLOAT))
driver_fs.add(Feature(name="feature_2", dtype=ValueType.STRING))
driver_fs.add(Feature(name="feature_3", dtype=ValueType.INT64))
driver_fs.add(Entity(name="entity_id", dtype=ValueType.INT64))

# Register with Feast core
test_client.apply(driver_fs)
driver_fs = driver_fs.to_proto()
driver_fs.meta.status = FeatureSetStatusProto.STATUS_READY

mocker.patch.object(
test_client._core_service_stub,
"GetFeatureSet",
return_value=GetFeatureSetResponse(feature_set=driver_fs),
)

with pytest.raises(exception):
test_client.ingest("driver-feature-set", dataframe)

@pytest.mark.parametrize(
"dataframe,exception,test_client",
[
Expand Down