Skip to content

Commit e01f465

Browse files
committed
create individual workdir for kernel and initialize a dedicated virtual env
TODO: file up- and download does not work yet due to the 'workspace' dir not being mapped correctly
1 parent 99a6ba8 commit e01f465

File tree

4 files changed

+61
-34
lines changed

4 files changed

+61
-34
lines changed

.gitignore

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44

55
# Program related
66
process_pids/
7+
kernel.*
78
kernel_connection_file.json
89

910
# Python stuff

frontend/src/App.tsx

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,7 @@ function App() {
5353
])
5454
);
5555
let [waitingForSystem, setWaitingForSystem] = useState<WaitingStates>(
56-
WaitingStates.Idle
56+
WaitingStates.StartingKernel
5757
);
5858
const chatScrollRef = React.useRef<HTMLDivElement>(null);
5959

@@ -77,11 +77,8 @@ function App() {
7777

7878
const handleCommand = (command: string) => {
7979
if (command == "reset") {
80-
addMessage({
81-
text: "Restarting the kernel.",
82-
type: "message",
83-
role: "system",
84-
});
80+
addMessage({text: "Restarting the kernel.", type: "message", role: "system"});
81+
setWaitingForSystem(WaitingStates.StartingKernel);
8582

8683
fetch(`${Config.API_ADDRESS}/restart`, {
8784
method: "POST",

frontend/src/components/Chat.tsx

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,7 @@ function Message(props: {
7575

7676

7777
export enum WaitingStates {
78+
StartingKernel = "Starting Kernel",
7879
GeneratingCode = "Generating code",
7980
RunningCode = "Running code",
8081
UploadingFile = "Uploading file",

gpt_code_ui/kernel_program/kernel_manager.py

Lines changed: 56 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
import sys
22
import subprocess
33
import os
4+
import shutil
5+
import atexit
46
import queue
57
import json
68
import signal
@@ -167,30 +169,53 @@ def flush_kernel_msgs(kc, tries=1, timeout=0.2):
167169
logger.debug(f"{e} [{type(e)}")
168170

169171

170-
def start_kernel():
171-
kernel_connection_file = os.path.join(os.getcwd(), "kernel_connection_file.json")
172-
173-
if os.path.isfile(kernel_connection_file):
174-
os.remove(kernel_connection_file)
175-
if os.path.isdir(kernel_connection_file):
176-
os.rmdir(kernel_connection_file)
177-
178-
launch_kernel_script_path = os.path.join(
179-
pathlib.Path(__file__).parent.resolve(), "launch_kernel.py"
180-
)
181-
182-
os.makedirs('workspace/', exist_ok=True)
183-
172+
def start_kernel(id: str):
173+
cwd = pathlib.Path(os.getcwd())
174+
kernel_dir = cwd / f'kernel.{id}'
175+
kernel_venv = kernel_dir / 'venv'
176+
kernel_venv_bindir = kernel_venv / 'bin'
177+
kernel_python_executable = kernel_venv_bindir / os.path.basename(sys.executable)
178+
kernel_connection_file = kernel_dir / "kernel_connection_file.json"
179+
launch_kernel_script_path = pathlib.Path(__file__).parent.resolve() / "launch_kernel.py"
180+
kernel_env = os.environ.copy()
181+
kernel_env['PATH'] = str(kernel_venv_bindir) + os.pathsep + kernel_env['PATH']
182+
183+
# Cleanup potential leftovers
184+
shutil.rmtree(kernel_dir, ignore_errors=True)
185+
186+
os.makedirs(kernel_dir)
187+
188+
# create virtual env inside kernel_venv directory
189+
subprocess.run([sys.executable, '-m', 'venv', kernel_venv, '--upgrade-deps', '--system-site-packages'])
190+
# install wheel because some packages do not like being installed without
191+
subprocess.run([kernel_python_executable, '-m', 'pip', 'install', 'wheel>=0.41,<1.0'])
192+
# install all default packages into the venv
193+
default_packages = [
194+
"ipykernel>=6,<7",
195+
"numpy>=1.24,<1.25",
196+
"dateparser>=1.1,<1.2",
197+
"pandas>=1.5,<1.6",
198+
"geopandas>=0.13,<0.14",
199+
"tabulate>=0.9.0<1.0",
200+
"PyPDF2>=3.0,<3.1",
201+
"pdfminer>=20191125,<20191200",
202+
"pdfplumber>=0.9,<0.10",
203+
"matplotlib>=3.7,<3.8",
204+
]
205+
subprocess.run([kernel_python_executable, '-m', 'pip', 'install'] + default_packages)
206+
207+
# start the kernel using the virtual env python executable
184208
kernel_process = subprocess.Popen(
185209
[
186-
sys.executable,
210+
kernel_python_executable,
187211
launch_kernel_script_path,
188212
"--IPKernelApp.connection_file",
189213
kernel_connection_file,
190214
"--matplotlib=inline",
191215
"--quiet",
192216
],
193-
cwd='workspace/'
217+
cwd=kernel_dir,
218+
env=kernel_env,
194219
)
195220
# Write PID for caller to kill
196221
str_kernel_pid = str(kernel_process.pid)
@@ -200,25 +225,28 @@ def start_kernel():
200225

201226
# Wait for kernel connection file to be written
202227
while True:
203-
if not os.path.isfile(kernel_connection_file):
228+
try:
229+
with open(kernel_connection_file, 'r') as fp:
230+
json.load(fp)
231+
except (FileNotFoundError, json.JSONDecodeError):
232+
# Either file was not yet there or incomplete (then JSON parsing failed)
204233
sleep(0.1)
234+
pass
205235
else:
206-
# Keep looping if JSON parsing fails, file may be partially written
207-
try:
208-
with open(kernel_connection_file, 'r') as fp:
209-
json.load(fp)
210-
break
211-
except json.JSONDecodeError:
212-
pass
236+
break
213237

214238
# Client
215-
kc = BlockingKernelClient(connection_file=kernel_connection_file)
239+
kc = BlockingKernelClient(connection_file=str(kernel_connection_file))
216240
kc.load_connection_file()
217241
kc.start_channels()
218242
kc.wait_for_ready()
219-
return kc
243+
return kc, kernel_dir
220244

221245

222246
if __name__ == "__main__":
223-
kc = start_kernel()
224-
start_snakemq(kc)
247+
kc, kernel_dir = start_kernel(id='ffa907c8-1908-49e3-974b-c0c27cbac6e4') # This is just a random but fixed ID for now - to be prepared for later multi-kernel scenarios
248+
249+
# make sure the dir with the virtualenv will be deleted after kernel termination
250+
atexit.register(lambda: shutil.rmtree(kernel_dir, ignore_errors=True))
251+
252+
start_snakemq(kc)

0 commit comments

Comments
 (0)