This repository has been archived by the owner on Aug 30, 2022. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 28
WIP: test coordinator::core::Service #350
Closed
Closed
Changes from all commits
Commits
Show all changes
7 commits
Select commit
Hold shift + click to select a range
406ba68
add a mock the aggregator rpc client
little-dude b648461
add a mock for coordinator::rpc::Client
little-dude aa51acc
make coordinator::core::Service Send
little-dude 4c73729
disallow dead code
little-dude 753404e
wip add tests
little-dude a5f3a9d
aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaHHHHHHHHHHHHHHHHHHHHHHHHH
little-dude 30a298f
cjeiwofjweioewj
little-dude File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,5 +1,3 @@ | ||
#![allow(dead_code)] | ||
|
||
pub mod api; | ||
pub mod py_aggregator; | ||
pub mod rpc; | ||
|
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -10,3 +10,6 @@ extern crate serde; | |
pub mod aggregator; | ||
pub mod common; | ||
pub mod coordinator; | ||
|
||
#[cfg(test)] | ||
mod tests; |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,77 @@ | ||
use crate::{ | ||
coordinator::{core::Service, models::HeartBeatResponse, settings::FederatedLearningSettings}, | ||
tests::mocks::{ | ||
coordinator::{MaxSelector, ServiceHandle}, | ||
rpc::aggregator::{Client, MockClient}, | ||
}, | ||
}; | ||
use futures::future; | ||
use tokio::time::{delay_for, Duration}; | ||
|
||
#[cfg(logging)] | ||
use crate::common::{logging, settings::LoggingSettings}; | ||
#[cfg(logging)] | ||
use tracing_subscriber::filter::EnvFilter; | ||
|
||
#[tokio::test] | ||
async fn test_rendez_vous_accept() { | ||
#[cfg(logging)] | ||
logging::configure(LoggingSettings { | ||
telemetry: None, | ||
filter: EnvFilter::try_new("trace").unwrap(), | ||
}); | ||
|
||
let rpc_client: Client = MockClient::default().into(); | ||
let aggregator_url = "http://localhost:8082".to_string(); | ||
|
||
let (service_handle, service_requests) = ServiceHandle::new(); | ||
|
||
let service = Service::new( | ||
MaxSelector, | ||
FederatedLearningSettings { | ||
rounds: 1, | ||
participants_ratio: 1.0, | ||
min_clients: 1, | ||
heartbeat_timeout: 10, | ||
}, | ||
aggregator_url.clone(), | ||
rpc_client.clone(), | ||
service_requests, | ||
); | ||
let _join_handle = tokio::spawn(service); | ||
|
||
let id = service_handle.rendez_vous_accepted().await; | ||
let round = service_handle.heartbeat_selected(id).await; | ||
assert_eq!(round, 0); | ||
|
||
rpc_client | ||
.mock() | ||
.expect_select() | ||
.returning(|_, _| future::ready(Ok(Ok(())))); | ||
|
||
let (url, _token) = service_handle.start_training_accepted(id).await; | ||
assert_eq!(url, aggregator_url); | ||
|
||
// pretend the client trained and sent its weights to the | ||
// aggregator. The aggregator now sends an end training requests | ||
// to the coordinator RPC server that we fake with the | ||
// service_handle. The service should then trigger the aggregation | ||
// and reject subsequent heartbeats and rendez-vous | ||
rpc_client | ||
.mock() | ||
.expect_aggregate() | ||
.returning(|_| future::ready(Ok(Ok(())))); | ||
|
||
service_handle.end_training(id, true).await; | ||
loop { | ||
match service_handle.heartbeat(id).await { | ||
HeartBeatResponse::StandBy => sleep_ms(10).await, | ||
HeartBeatResponse::Finish => break, | ||
_ => panic!("expected StandBy or Finish"), | ||
} | ||
} | ||
} | ||
|
||
async fn sleep_ms(ms: u64) { | ||
delay_for(Duration::from_millis(ms)).await | ||
} |
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
First working test :)