"vscode:/vscode.git/clone" did not exist on "e4b5348ebf407f0886bc11caa7ca74d2fbf693ac"
utils.py 730 Bytes
Newer Older
Woosuk Kwon's avatar
Woosuk Kwon committed
1
import enum
Zhuohan Li's avatar
Zhuohan Li committed
2
import uuid
Zhuohan Li's avatar
Zhuohan Li committed
3

4
import psutil
Zhuohan Li's avatar
Zhuohan Li committed
5
6
import torch

Woosuk Kwon's avatar
Woosuk Kwon committed
7
8
9
10
11
12
13
14
15
16
17

class Device(enum.Enum):
    GPU = enum.auto()
    CPU = enum.auto()


class Counter:

    def __init__(self, start: int = 0) -> None:
        self.counter = start

Woosuk Kwon's avatar
Woosuk Kwon committed
18
    def __next__(self) -> int:
Woosuk Kwon's avatar
Woosuk Kwon committed
19
20
21
22
23
24
        id = self.counter
        self.counter += 1
        return id

    def reset(self) -> None:
        self.counter = 0
Zhuohan Li's avatar
Zhuohan Li committed
25

26
27

def get_gpu_memory(gpu: int = 0) -> int:
28
    """Returns the total memory of the GPU in bytes."""
29
30
31
32
    return torch.cuda.get_device_properties(gpu).total_memory


def get_cpu_memory() -> int:
33
    """Returns the total CPU memory of the node in bytes."""
34
    return psutil.virtual_memory().total
Zhuohan Li's avatar
Zhuohan Li committed
35
36
37
38


def random_uuid() -> str:
    return str(uuid.uuid4().hex)