File size: 1,457 Bytes
608ca3b
 
 
442695b
 
608ca3b
 
 
c6d0ab0
 
11f21b8
608ca3b
 
 
 
 
 
a2cd106
442695b
 
a2cd106
 
644a7e4
608ca3b
50111d4
 
 
500868d
50111d4
 
 
 
 
 
 
 
 
 
 
 
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
import os
from fastapi import FastAPI
from lightning_sdk import Studio
import logging
from io import StringIO

app = FastAPI()

log_stream = StringIO()
logging.basicConfig(stream=log_stream, level=logging.INFO)
    
@app.post("/start-studio/")
def start_studio(userid: str, apikey: str, name: str, teamspace: str, user: str):
    try:
        os.environ["LIGHTNING_USER_ID"] = userid
        os.environ["LIGHTNING_API_KEY"] = apikey
        s = Studio(name=name, teamspace=teamspace, user=user)
        s.start()
        log_content = log_stream.getvalue().strip()
        print(log_content)
        os.environ.pop("LIGHTNING_USER_ID", None)
        os.environ.pop("LIGHTNING_API_KEY", None)
        return {"status": "success", "message": log_content}
    except Exception as e:
        return {"status": "error", "message": str(e)}

@app.post("/stop-studio/")
def stop_studio(userid: str, apikey: str, name: str, teamspace: str, user: str):
    try:
        os.environ["LIGHTNING_USER_ID"] = userid
        os.environ["LIGHTNING_API_KEY"] = apikey
        s = Studio(name=name, teamspace=teamspace, user=user)
        s.stop()
        log_content = log_stream.getvalue().strip()
        print(log_content)
        os.environ.pop("LIGHTNING_USER_ID", None)
        os.environ.pop("LIGHTNING_API_KEY", None)
        return {"status": "success", "message": log_content}
    except Exception as e:
        return {"status": "error", "message": str(e)}