saymrwulf/pytorch
1
0
Fork 0
mirror of https://github.com/saymrwulf/pytorch.git synced 2026-05-14 20:57:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add pytorch-distributed justknobs helper (#118568)

Browse source 
Summary:
Sets up a helper that checks any JKs relevent to pytorch distributed,
and propagates their values to ENV.

Test Plan: Added unit test

Differential Revision: D53192406

Pull Request resolved: https://github.com/pytorch/pytorch/pull/118568
Approved by: https://github.com/zdevito
This commit is contained in:
Will Constable 2024-01-30 08:13:52 +00:00 committed by PyTorch MergeBot
parent 3ecc2f3a0d
commit da0635d17c
2 changed files with 7 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
torch/_utils_internal.py
View file

@ -86,6 +86,10 @@ def print_graph(graph, msg: str):
pass
def set_pytorch_distributed_envs_from_justknobs():
pass
TEST_MASTER_ADDR = "127.0.0.1"
TEST_MASTER_PORT = 29500
# USE_GLOBAL_DEPS controls whether __init__.py tries to load

3
torch/distributed/distributed_c10d.py
View file

@ -40,6 +40,7 @@ from torch._C._distributed_c10d import (
_unregister_all_process_groups,
_unregister_process_group,
)
from torch._utils_internal import set_pytorch_distributed_envs_from_justknobs
from .constants import default_pg_timeout, default_pg_nccl_timeout
from .c10d_logger import _exception_logger, _time_logger
from .rendezvous import register_rendezvous_handler, rendezvous # noqa: F401
@ -1191,6 +1192,8 @@ def init_process_group(
"cpu:gloo,cuda:custom_backend".
"""
set_pytorch_distributed_envs_from_justknobs()
global _world
global _backend