Skip to content

fix(replication): safely shutdown transfer service #105

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

Merged
kkkapu merged 1 commit into from
May 5, 2026
Merged

fix(replication): safely shutdown transfer service #105

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
1 change: 1 addition & 0 deletions src/ml_flashpoint/replication/replication_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -503,6 +503,7 @@ def shutdown(self):
_LOGGER.info("Shutting down ReplicationManager and TransferService...")
if self._transfer_service:
self._transfer_service.shutdown()
self._transfer_service = None

@g-husam g-husam May 5, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

was discussing with @Leahlijuan last week - not sure if this truly fixes the problem, since there will still be ongoing requests, this just wont allow new requests.

We need to deterministically prevent any replication on train end, or wait for it all to finish, to fully fix this I think

_LOGGER.info("TransferService shut down.")

# TODO: Use the ip address return from c++ transfer service to avoid duplication.
Expand Down
14 changes: 14 additions & 0 deletions tests/replication/test_replication_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -398,3 +398,17 @@ def test_async_replicate_single_node_skips(replication_manager, mocker):
replication_manager._checkpoint_object_manager.close_buffer.assert_called_once_with(
buffer_io, skip_close_if_symlink=True
)


def test_shutdown_clears_transfer_service(replication_manager):
"""Tests that shutdown calls transfer_service.shutdown() and sets it to None."""
# Given
mock_transfer_service = replication_manager._transfer_service

# When
replication_manager.shutdown()

# Then
mock_transfer_service.shutdown.assert_called_once()

assert replication_manager._transfer_service is None
Loading