set global envs on first execution

Author: mishushakov

template/server/contexts.py

@@ -8,7 +8,6 @@
 from consts import JUPYTER_BASE_URL
 from errors import ExecutionError
 from messaging import ContextWebSocket
-from envs import get_envs
 
 logger = logging.Logger(__name__)
 
@@ -52,7 +51,6 @@ async def create_context(client, websockets: dict, language: str, cwd: str) -> C
     session_data = response.json()
     session_id = session_data["id"]
     context_id = session_data["kernel"]["id"]
-    global_env_vars = await get_envs()
 
     logger.debug(f"Created context {context_id}")
 
@@ -69,12 +67,4 @@ async def create_context(client, websockets: dict, language: str, cwd: str) -> C
             status_code=500,
         )
 
-    try:
-        await ws.set_env_vars(global_env_vars)
-    except ExecutionError as e:
-        return PlainTextResponse(
-            "Failed to set environment variables",
-            status_code=500,
-        )
-
     return Context(language=language, id=context_id, cwd=cwd)

template/server/main.py

@@ -17,6 +17,7 @@
 from messaging import ContextWebSocket
 from stream import StreamingListJsonResponse
 from utils.locks import LockedMap
+from envs import get_envs
 
 logging.basicConfig(level=logging.DEBUG, stream=sys.stdout)
 logger = logging.Logger(__name__)
@@ -104,6 +105,11 @@ async def post_execute(request: ExecutionRequest):
             status_code=404,
         )
 
+    # set global env vars if not set on first execution
+    if not ws.global_env_vars:
+        ws.global_env_vars = await get_envs()
+        await ws.set_env_vars(ws.global_env_vars)
+
     return StreamingListJsonResponse(
         ws.execute(
             request.code,

template/server/messaging.py

@@ -3,10 +3,8 @@
 import logging
 import uuid
 import asyncio
-import subprocess
 
 from asyncio import Queue
-from envs import get_envs
 from typing import (
     Dict,
     Optional,
@@ -49,20 +47,20 @@ def __init__(self, in_background: bool = False):
 class ContextWebSocket:
     _ws: Optional[WebSocketClientProtocol] = None
     _receive_task: Optional[asyncio.Task] = None
+    global_env_vars: Optional[Dict[StrictStr, str]] = None
 
     def __init__(
         self,
         context_id: str,
         session_id: str,
         language: str,
-        cwd: str,
+        cwd: str
     ):
         self.language = language
         self.cwd = cwd
         self.context_id = context_id
         self.url = f"ws://localhost:8888/api/kernels/{context_id}/channels"
         self.session_id = session_id
-
         self._executions: Dict[str, Execution] = {}
         self._lock = asyncio.Lock()
 
@@ -167,15 +165,13 @@ async def set_env_vars(self, env_vars: Dict[StrictStr, str]):
                     raise ExecutionError(f"Error during execution: {item}")
 
     async def reset_env_vars(self, env_vars: Dict[StrictStr, str]):
-        global_env_vars = await get_envs()
-
         # Create a dict of vars to reset and a list of vars to remove
         vars_to_reset = {}
         vars_to_remove = []
 
         for key in env_vars:
-            if key in global_env_vars:
-                vars_to_reset[key] = global_env_vars[key]
+            if self.global_env_vars and key in self.global_env_vars:
+                vars_to_reset[key] = self.global_env_vars[key]
             else:
                 vars_to_remove.append(key)