-
Notifications
You must be signed in to change notification settings - Fork 2
/
utils.py
84 lines (65 loc) · 2.47 KB
/
utils.py
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
import sys
from subprocess import run
from time import sleep
import re
from flytekit.models.core import execution as core_execution_models
def get_recent_execution_id(workflow_name, remote):
recent_executions = remote.recent_executions()
executions = [
e for e in recent_executions if e.spec.launch_plan.name == workflow_name
]
return executions[0].id.name
def launch_vscode(llm_uri, remote):
command = [
"unionai",
"run",
"--remote",
"workflows/launch_vscode.py",
"main",
"--model",
llm_uri,
]
result = run(command, capture_output=True, text=True)
match = re.search(r"executions/(\w+) to see execution", result.stdout)
if not match:
raise RuntimeError("Please run launch_vscode again")
execution_id = match.group(1)
print("VSCode workflow launched! Waiting for VSCode server to start...")
execution = remote.fetch_execution(name=execution_id)
execution = remote.sync_execution(execution, sync_nodes=True)
vscode_logs = []
count = 0
while not vscode_logs:
count += 1
if count % 10 == 0:
print("Still waiting...")
execution = remote.sync_execution(execution, sync_nodes=True)
if execution.is_done and execution.error is not None:
print("There was an error launching VScode. Please try again.")
sys.exit(1)
if execution.closure.phase != core_execution_models.TaskExecutionPhase.RUNNING:
sleep(1)
continue
if not execution.node_executions or "n0" not in execution.node_executions:
sleep(1)
continue
node_exectuion = execution.node_executions["n0"]
task_execution = node_exectuion.task_executions[0]
logs = task_execution.closure.logs
vscode_logs = [l for l in logs if l.name.startswith("VSCode")]
sleep(1)
for _ in range(3):
print("Still waiting...")
sleep(10)
uri = vscode_logs[0].uri
endpoint = remote.config.platform.endpoint
print("✅ Your VSCode instance launched at:")
print(f"https://{endpoint}{uri}")
return execution.id.name
def stop_vscode(vscode_execution_id, remote):
vscode_execution = remote.fetch_execution(name=vscode_execution_id)
if not vscode_execution.is_done:
remote.terminate(vscode_execution, cause="Stop VSCode from Colab")
print("Stopped VSCode instance!")
else:
print("VSCode instance already stopped!")