workspaces/workspaces.py at master · acadnet-dev/workspaces · GitHub

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
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
from fastapi import FastAPI, UploadFile, File, HTTPException
import uvicorn
import time
from uuid import uuid4
import os
from os import path
import yaml

from config import Config
from utils import upload_textfile_to_pod

from kubernetes import client, config

app_config = Config("config.json")
app = FastAPI()


class Workspace:
    def __init__(self, workspace_id):
        self.id = workspace_id
        self.pod_name = "vscode-server-" + self.id

    def add_file(self, file: UploadFile):
        with open(self.path + "/" + file.filename, "wb") as f:
            f.write(file.file.read())

    def create_pod(self):
        v1 = client.CoreV1Api()

        with open(path.join(path.dirname(__file__), "vscode-server-deployment.yaml")) as f:
            pod = yaml.safe_load(f)

            pod_name = pod["metadata"]["name"] + "-" + self.id
            pod["metadata"]["name"] = pod_name

            resp = v1.create_namespaced_pod(
                body=pod, namespace="acadnet")
            print(f"Pod created with name {pod_name}")

        # wait for sandbox to start maxim 5 minutes (pod creation takes a while, even more if scaling up)
        # 5 minutes = 60 * 5 = 300 seconds (poll every 5 seconds 60 times)
        ok = False
        for i in range(60):
            pod_status = v1.read_namespaced_pod_status(pod_name, "acadnet")
            if pod_status.status.phase == "Running":
                ok = True
                break
            time.sleep(5)

        if not ok:
            raise Exception("Sandbox failed to start")

        self.pod_name = pod_name

    def get_pod_endpoint(self):
        v1 = client.CoreV1Api()

        # if pod does not exist, return None
        try:
            status = v1.read_namespaced_pod_status(self.pod_name, "acadnet")

            if status.status.phase != "Running":
                return None

            return f"http://{status.status.pod_ip}:3000"
        except Exception as e:
            return None

    def add_file(self, file: UploadFile, problem_name: str):
        v1 = client.CoreV1Api()

        upload_textfile_to_pod(v1, self.pod_name, file, "/home/workspace", problem_name)

# creates and then returns the endpoint for workspace
@app.post("/workspace/create/")
async def create_workspace(id: str, problem_name: str, files: list[UploadFile]):
    try:
        # create new workspace
        workspace = Workspace(id)

        # get pod endpoint
        endpoint = workspace.get_pod_endpoint()

        # if pod does not exist, create it
        if endpoint == None:
            # create pod
            workspace.create_pod()
            # add files to pod
            for file in files:
                workspace.add_file(file, problem_name)
            # get pod endpoint
            endpoint = workspace.get_pod_endpoint()

        return {"endpoint": endpoint}
    except Exception as e:
        raise e
        return {"error": str(e)}

# returns the endpoint for workspace if it exists
@app.get("/workspace/get/")
async def get_workspace(id: str):
    # create new workspace
    workspace = Workspace(id)

    # get pod endpoint
    endpoint = workspace.get_pod_endpoint()

    # if pod does not exist, return error
    if endpoint == None:
        raise HTTPException(status_code=404, detail="Workspace not found")

    return {"endpoint": endpoint}

def run():
    uvicorn.run(app, host="0.0.0.0", port=app_config.port)

if __name__ == "__main__":
    if app_config.is_development():
        config.load_kube_config('/home/dimi/.kube/config', context='do-fra1-acadnet-dev-k8s')
    else:
        config.load_incluster_config()
    run()