"cuda/cuda_stream_manager.h" did not exist on "27d8beaa0ea7e20a8ee3a692004072017929dcf9"
helper.py 1.59 KB
Newer Older
facebook-github-bot's avatar
facebook-github-bot committed
1
2
3
4
#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved


5
6
7
import os
import socket
import uuid
facebook-github-bot's avatar
facebook-github-bot committed
8
from functools import wraps
Kai Zhang's avatar
Kai Zhang committed
9
from tempfile import TemporaryDirectory
facebook-github-bot's avatar
facebook-github-bot committed
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33

import torch
import torch.distributed as dist


def skip_if_no_gpu(func):
    """Decorator that can be used to skip GPU tests on non-GPU machines."""
    func.skip_if_no_gpu = True

    @wraps(func)
    def wrapper(*args, **kwargs):
        if not torch.cuda.is_available():
            return
        if torch.cuda.device_count() <= 0:
            return

        return func(*args, **kwargs)

    return wrapper


def enable_ddp_env(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
34
35
36
37
        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])
facebook-github-bot's avatar
facebook-github-bot committed
38

39
40
        os.environ["MASTER_ADDR"] = "localhost"
        os.environ["MASTER_PORT"] = find_free_port()
facebook-github-bot's avatar
facebook-github-bot committed
41
42
43
44
        dist.init_process_group(
            "gloo",
            rank=0,
            world_size=1,
45
46
47
            init_method="file:///tmp/detectron2go_test_ddp_init_{}".format(
                uuid.uuid4().hex
            ),
facebook-github-bot's avatar
facebook-github-bot committed
48
49
50
51
52
53
        )
        ret = func(*args, **kwargs)
        dist.destroy_process_group()
        return ret

    return wrapper
Kai Zhang's avatar
Kai Zhang committed
54

Yanghan Wang's avatar
Yanghan Wang committed
55

Kai Zhang's avatar
Kai Zhang committed
56
57
58
59
60
61
62
63
64
def tempdir(func):
    """ A decorator for creating a tempory directory that is cleaned up after function execution. """

    @wraps(func)
    def wrapper(self, *args, **kwargs):
        with TemporaryDirectory() as temp:
            return func(self, temp, *args, **kwargs)

    return wrapper