add ms ck test

app.py

@@ -6,8 +6,8 @@ from datetime import datetime
 from modules.smtp import send_email
 from modules.activate import trigger
 from modules.times import fix_datetime, time_diff
-from modules.restart import self_restart, test_restart
-from utils import START_TIME, DELAY, PERIOD
+from modules.restart import self_restart, test_hf_restart, test_ms_restart
+from utils import START_TIME, PERIOD, DELAY
 
 
 def run_schedule():
@@ -51,18 +51,31 @@ def tasklst():
 if __name__ == "__main__":
     monitor()
     with gr.Blocks() as demo:
-        gr.Markdown("# Keep Spaces Alive")
+        gr.Markdown("# Keep Spaces/Studios Alive")
         with gr.Row():
             with gr.Column():
                 task_btn = gr.Button("See current task status")
                 once_btn = gr.Button("Trigger once manually")
                 smtp_btn = gr.Button("SMTP test")
-                re_txt = gr.Textbox(
-                    label="Restart a space with permissions",
-                    placeholder="username/repo",
-                    value="kakamond/cmd_inject",
-                )
-                re_btn = gr.Button("Restart to test cookie validity")
+                with gr.Tab("HF cookie test"):
+                    hf_txt = gr.Textbox(
+                        label="Restart a space with permissions",
+                        placeholder="username/repo",
+                        value="kakamond/cmd_inject",
+                    )
+                    with gr.Row():
+                        clhf_btn = gr.Button("Clear")
+                        hf_btn = gr.Button("Restart to test cookie validity")
+
+                with gr.Tab("MS cookie test"):
+                    ms_txt = gr.Textbox(
+                        label="Restart a studio with permissions",
+                        placeholder="username/repo",
+                        value="kakamond/cmd_inject",
+                    )
+                    with gr.Row():
+                        clms_btn = gr.Button("Clear")
+                        ms_btn = gr.Button("Restart to test cookie validity")
 
             with gr.Column():
                 status_bar = gr.Textbox(label="Status", show_copy_button=True)
@@ -72,10 +85,17 @@ if __name__ == "__main__":
             task_btn.click(fn=tasklst, outputs=[status_bar, logs_bar])
             once_btn.click(fn=trigger, outputs=[status_bar, logs_bar])
             smtp_btn.click(fn=send_email, outputs=[status_bar, logs_bar])
-            re_btn.click(
-                fn=test_restart,
-                inputs=[re_txt],
+            hf_btn.click(
+                fn=test_hf_restart,
+                inputs=[hf_txt],
+                outputs=[status_bar, logs_bar],
+            )
+            clhf_btn.click(fn=lambda: None, outputs=hf_txt)
+            ms_btn.click(
+                fn=test_ms_restart,
+                inputs=[ms_txt],
                 outputs=[status_bar, logs_bar],
             )
+            clms_btn.click(fn=lambda: None, outputs=ms_txt)
 
     demo.launch()

modules/restart.py

@@ -1,15 +1,22 @@
 import requests
 from modules.smtp import send_email
-from utils import HF_DOMAIN, HEADER, TIMEOUT, HF_CK
+from utils import HF_DOMAIN, MS_DOMAIN, HF_HEADER, MS_HEADER, TIMEOUT
+
+
+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_private_space(repo: str):
     response = requests.post(
         f"{HF_DOMAIN}/api/spaces/{repo}/restart",
-        headers={
-            "User-Agent": HEADER["User-Agent"],
-            "cookie": HF_CK,
-        },
+        headers=HF_HEADER,
         timeout=TIMEOUT,
     )
     response.raise_for_status()
@@ -24,10 +31,13 @@ def self_restart():
 
 
 # UI func
-def test_restart(repo):
+def test_hf_restart(repo):
     status = "Success"
     logs = None
     try:
+        if not repo:
+            raise ValueError("Empty repo!")
+
         logs = restart_private_space(repo)
         logs += f"\nSuccessfully restart {HF_DOMAIN}/spaces/{repo}\n"
 
@@ -35,3 +45,19 @@ def test_restart(repo):
         status = f"{e}"
 
     return status, logs
+
+
+def test_ms_restart(repo):
+    status = "Success"
+    logs = None
+    try:
+        if not repo:
+            raise ValueError("Empty 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

utils.py

@@ -27,6 +27,10 @@ MS_DOMAIN = "https://www.modelscope.cn"
 HEADER = {
     "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537"
 }
+HF_HEADER = {
+    "User-Agent": HEADER["User-Agent"],
+    "Cookie": HF_CK,
+}
 MS_HEADER = {
     "User-Agent": HEADER["User-Agent"],
     "Cookie": MS_CK,