Merge remote-tracking branch 'origin' into HTTP/2_Multi_Room_Lighthouse

Author: d4l3k

.lintrunner.toml

@@ -65,5 +65,6 @@ init_command = [
     'pip_init',
     '--dry-run={{DRYRUN}}',
     'pyre-check==0.9.23',
+    'click==8.1.0',
 ]
 is_formatter = false

proto/torchft.proto

@@ -73,7 +73,7 @@ service LighthouseService {
 }
 
 message ManagerQuorumRequest {
-    int64 rank = 1;
+    int64 group_rank = 1;
     int64 step = 2;
     string checkpoint_metadata = 3;
     bool shrink_only = 4;
@@ -84,12 +84,12 @@ message ManagerQuorumRequest {
 message ManagerQuorumResponse {
     int64 quorum_id = 1;
     string recover_src_manager_address = 2;
-    optional int64 recover_src_rank = 3;
-    repeated int64 recover_dst_ranks = 4;
+    optional int64 recover_src_replica_rank = 3;
+    repeated int64 recover_dst_replica_ranks = 4;
     string store_address = 5;
     // These are information for the replicas which are at the max step.
     int64 max_step = 6;
-    optional int64 max_rank = 7;
+    optional int64 max_replica_rank = 7;
     int64 max_world_size = 8;
     // These are information for all replicas including behind replicas.
     int64 replica_rank = 9;
@@ -108,7 +108,7 @@ message CheckpointMetadataResponse {
 
 message ShouldCommitRequest {
     bool should_commit = 1;
-    int64 rank = 2;
+    int64 group_rank = 2;
     int64 step = 3;
 }
 message ShouldCommitResponse {

pyproject.toml

@@ -22,7 +22,7 @@ Issues = "https://github.com/pytorch-labs/torchft/issues"
 
 [project.optional-dependencies]
 dev = [
-    "pytest",
+    "pytest==8.3.4",
     "pytest-timeout",
     "black",
     "pyre-check",

src/lib.rs

@@ -181,7 +181,7 @@ impl ManagerClient {
     fn _quorum(
         &self,
         py: Python<'_>,
-        rank: i64,
+        group_rank: i64,
         step: i64,
         checkpoint_metadata: String,
         shrink_only: bool,
@@ -191,7 +191,7 @@ impl ManagerClient {
     ) -> Result<QuorumResult, StatusError> {
         py.allow_threads(move || {
             let mut request = tonic::Request::new(ManagerQuorumRequest {
-                rank: rank,
+                group_rank: group_rank,
                 step: step,
                 checkpoint_metadata: checkpoint_metadata,
                 shrink_only: shrink_only,
@@ -210,11 +210,11 @@ impl ManagerClient {
                 replica_rank: resp.replica_rank,
                 replica_world_size: resp.replica_world_size,
                 recover_src_manager_address: resp.recover_src_manager_address,
-                recover_src_rank: resp.recover_src_rank,
-                recover_dst_ranks: resp.recover_dst_ranks,
+                recover_src_replica_rank: resp.recover_src_replica_rank,
+                recover_dst_replica_ranks: resp.recover_dst_replica_ranks,
                 store_address: resp.store_address,
                 max_step: resp.max_step,
-                max_rank: resp.max_rank,
+                max_replica_rank: resp.max_replica_rank,
                 max_world_size: resp.max_world_size,
                 heal: resp.heal,
             })
@@ -259,14 +259,14 @@ impl ManagerClient {
     fn should_commit(
         &self,
         py: Python<'_>,
-        rank: i64,
+        group_rank: i64,
         step: i64,
         should_commit: bool,
         timeout: Duration,
     ) -> Result<bool, StatusError> {
         py.allow_threads(move || {
             let mut request = tonic::Request::new(ShouldCommitRequest {
-                rank: rank,
+                group_rank: group_rank,
                 step: step,
                 should_commit: should_commit,
             });
@@ -290,11 +290,11 @@ struct QuorumResult {
     replica_rank: i64,
     replica_world_size: i64,
     recover_src_manager_address: String,
-    recover_src_rank: Option<i64>,
-    recover_dst_ranks: Vec<i64>,
+    recover_src_replica_rank: Option<i64>,
+    recover_dst_replica_ranks: Vec<i64>,
     store_address: String,
     max_step: i64,
-    max_rank: Option<i64>,
+    max_replica_rank: Option<i64>,
     max_world_size: i64,
     heal: bool,
 }
@@ -308,11 +308,11 @@ impl QuorumResult {
             replica_rank: 0,
             replica_world_size: 1,
             recover_src_manager_address: "".to_string(),
-            recover_src_rank: None,
-            recover_dst_ranks: Vec::new(),
+            recover_src_replica_rank: None,
+            recover_dst_replica_ranks: Vec::new(),
             store_address: "".to_string(),
             max_step: 0,
-            max_rank: None,
+            max_replica_rank: None,
             max_world_size: 1,
             heal: false,
         }