Skip to content

Commit

Permalink
[Typing][B-100] Add type annotations for `python/paddle/distributed/p…
Browse files Browse the repository at this point in the history
…arallel_with_gloo.py` (PaddlePaddle#66473)
  • Loading branch information
enkilee authored and inaomIIsfarell committed Jul 31, 2024
1 parent b81ae39 commit 07f857f
Showing 1 changed file with 7 additions and 4 deletions.
11 changes: 7 additions & 4 deletions python/paddle/distributed/parallel_with_gloo.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations

import time
from multiprocessing import Manager, Process
Expand Down Expand Up @@ -38,12 +39,14 @@ def _start_kv_server(port, http_server_d, size):
http_server.stop()


def gloo_init_parallel_env(rank_id, rank_num, server_endpoint):
def gloo_init_parallel_env(
rank_id: int, rank_num: int, server_endpoint: str
) -> None:
"""
Initialize parallel environment with gloo for cpu only.
Args:
- rank_idint, required) - the index of current rank;
- rank_id (int, required) - the index of current rank;
- rank_num (int, required) - the number of ranks in this parallel env;
- server_endpoint (str, required) - endpoint of server to init gloo context in ip:port format;
Expand Down Expand Up @@ -135,7 +138,7 @@ def gloo_init_parallel_env(rank_id, rank_num, server_endpoint):
http_server_proc.join()


def gloo_barrier():
def gloo_barrier() -> None:
"""
Call barrier function with initialized gloo context.
Expand Down Expand Up @@ -193,7 +196,7 @@ def gloo_barrier():
_global_gloo_ctx.barrier()


def gloo_release():
def gloo_release() -> None:
"""
Release the parallel environment initialized by gloo
Expand Down

0 comments on commit 07f857f

Please sign in to comment.