Commit f9c62b77 authored by Yanghan Wang's avatar Yanghan Wang Committed by Facebook GitHub Bot
Browse files

initial implementation of build_auto_stream_train_loader

Summary: Add `build_auto_stream_train_loader`

Reviewed By: newstzpz

Differential Revision: D27343030

fbshipit-source-id: a79d3ed1ac41fc159d10bb6ff1db74549b645a1c
parent bf8d84b2
...@@ -2,6 +2,9 @@ ...@@ -2,6 +2,9 @@
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
import os
import socket
import uuid
from functools import wraps from functools import wraps
from tempfile import TemporaryDirectory from tempfile import TemporaryDirectory
...@@ -28,12 +31,20 @@ def skip_if_no_gpu(func): ...@@ -28,12 +31,20 @@ def skip_if_no_gpu(func):
def enable_ddp_env(func): def enable_ddp_env(func):
@wraps(func) @wraps(func)
def wrapper(*args, **kwargs): def wrapper(*args, **kwargs):
def find_free_port() -> str:
s = socket.socket()
s.bind(("localhost", 0)) # Bind to a free port provided by the host.
return str(s.getsockname()[1])
os.environ["MASTER_ADDR"] = "localhost"
os.environ["MASTER_PORT"] = find_free_port()
dist.init_process_group( dist.init_process_group(
"gloo", "gloo",
rank=0, rank=0,
world_size=1, world_size=1,
init_method="file:///tmp/detectron2go_test_ddp_init", init_method="file:///tmp/detectron2go_test_ddp_init_{}".format(
uuid.uuid4().hex
),
) )
ret = func(*args, **kwargs) ret = func(*args, **kwargs)
dist.destroy_process_group() dist.destroy_process_group()
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment