Skip to content

add client dict to separate different sessions #485

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
liqul merged 2 commits into from
May 8, 2025
Merged

add client dict to separate different sessions #485

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
4110350
add client dict to separate different sessions
liqul May 8, 2025
29a8f35
Merge branch 'main' into liqun/fix_env_client
liqul May 8, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 9 additions & 9 deletions taskweaver/ces/environment.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -122,10 +122,10 @@ def __init__(
custom_image: Optional[str] = None,
) -> None:
self.session_dict: Dict[str, EnvSession] = {}
self.client_dict: Dict[str, BlockingKernelClient] = {}
self.id = get_id(prefix="env") if env_id is None else env_id
self.env_dir = env_dir if env_dir is not None else os.getcwd()
self.mode = env_mode
self._client: Optional[BlockingKernelClient] = None

if self.mode == EnvMode.Local:
self.multi_kernel_manager = TaskWeaverMultiKernelManager(
Expand Down Expand Up @@ -396,7 +396,7 @@ def update_session_var(
session.session_var.update(session_var)

def stop_session(self, session_id: str) -> None:
self._clean_client()
self._clean_client(session_id)
session = self._get_session(session_id)
if session is None:
# session not exist
Expand Down Expand Up @@ -490,8 +490,8 @@ def _get_client(
self,
session_id: str,
) -> BlockingKernelClient:
if self._client is not None:
return self._client
if session_id in self.client_dict:
return self.client_dict[session_id]
session = self._get_session(session_id)
connection_file = self._get_connection_file(session_id, session.kernel_id)
client = BlockingKernelClient(connection_file=connection_file)
Expand All @@ -507,13 +507,13 @@ def _get_client(
client.iopub_port = ports["iopub_port"]
client.wait_for_ready(timeout=30)
client.start_channels()
self._client = client
self.client_dict[session_id] = client
return client

def _clean_client(self):
if self._client is not None:
self._client.stop_channels()
self._client = None
def _clean_client(self, session_id: str):
if session_id in self.client_dict:
self.client_dict[session_id].stop_channels()
del self.client_dict[session_id]

def _execute_code_on_kernel(
self,
Expand Down