"server/vscode:/vscode.git/clone" did not exist on "49b4b33e805d0ffee62688fe2607120b0c759e3d"
utils.py 1012 Bytes
Newer Older
Woosuk Kwon's avatar
Woosuk Kwon committed
1
import enum
Zhuohan Li's avatar
Zhuohan Li committed
2
import random
3
import psutil
Zhuohan Li's avatar
Zhuohan Li committed
4
5
6
7
8
9

import numpy as np
import torch

from cacheflow.parallel_utils.parallel_state import model_parallel_is_initialized
from cacheflow.parallel_utils.tensor_parallel import model_parallel_cuda_manual_seed
Woosuk Kwon's avatar
Woosuk Kwon committed
10
11
12
13
14
15
16
17
18
19
20
21


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
22
    def __next__(self) -> int:
Woosuk Kwon's avatar
Woosuk Kwon committed
23
24
25
26
27
28
        id = self.counter
        self.counter += 1
        return id

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

30

Zhuohan Li's avatar
Zhuohan Li committed
31
32
33
34
35
36
37
38
39
def set_random_seed(seed: int):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)

    if model_parallel_is_initialized():
        model_parallel_cuda_manual_seed(seed)
40
41
42
43
44
45
46
47


def get_gpu_memory(gpu: int = 0) -> int:
    return torch.cuda.get_device_properties(gpu).total_memory


def get_cpu_memory() -> int:
    return psutil.virtual_memory().total