File size: 2,074 Bytes
a98d805
0a7ac7a
a98d805
ad841e1
609a724
0a7ac7a
 
609a724
 
 
 
 
 
 
 
 
a98d805
 
dac0f46
ad841e1
 
5d2bc12
dac0f46
5d2bc12
dac0f46
0a7ac7a
dac0f46
 
ad841e1
a98d805
ad841e1
 
293a93b
ad841e1
 
 
533bd9d
a98d805
 
 
 
b50dddd
d93eb2d
a98d805
b50dddd
a98d805
609a724
 
 
d93eb2d
 
 
 
293a93b
b50dddd
a98d805
 
 
 
b50dddd
609a724
 
d93eb2d
609a724
 
 
 
 
 
d93eb2d
 
 
 
609a724
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
import requests
from huggingface_hub import HfApi
from modules.smtp import send_email
from utils import HF_DOMAIN, MS_DOMAIN, HF_TK, MS_HEADER, TIMEOUT, REPOS

HF_API = HfApi(token=HF_TK)


def restart_private_studio(repo: str):
    response = requests.put(
        f"{MS_DOMAIN}/api/v1/studio/{repo}/reset_restart",
        headers=MS_HEADER,
        timeout=TIMEOUT,
    )
    response.raise_for_status()
    return f"\n{response.text}\n"


def restart_repos():
    if REPOS:
        repos = REPOS.split(";")
        for repo in repos:
            private_repo = repo.strip()
            print(f"Restarting {private_repo} space & studio...")
            restart_private_studio(private_repo)
            HF_API.restart_space(private_repo)


def self_restart(me="kakamond/keeps_alive"):
    try:
        spaces = HF_API.list_spaces(author=me.split("/")[0])
        for space in spaces:
            if space.id != me and not space.disabled:
                HF_API.restart_space(space.id)

        HF_API.restart_space(me)

    except Exception as e:
        send_email(f"Failed to self-restart: {e}")


# UI func
def test_hf_restart(repo: str):
    status = "Success"
    logs = None
    try:
        if not repo:
            raise ValueError("Empty repo!")

        if f"{HF_DOMAIN}/spaces/" in repo:
            repo = repo.split("/spaces/")[1].split("/")[:2]
            repo = "/".join(repo)

        logs = f"\n{HF_API.restart_space(repo)}\n"
        logs += f"\nSuccessfully restart {HF_DOMAIN}/spaces/{repo}\n"

    except Exception as e:
        status = f"{e}"

    return status, logs


def test_ms_restart(repo: str):
    status = "Success"
    logs = None
    try:
        if not repo:
            raise ValueError("Empty repo!")

        if f"{MS_DOMAIN}/studios/" in repo:
            repo = repo.split("/studios/")[1].split("/")[:2]
            repo = "/".join(repo)

        logs = restart_private_studio(repo)
        logs += f"\nSuccessfully restart {MS_DOMAIN}/studios/{repo}\n"

    except Exception as e:
        status = f"{e}"

    return status, logs