mirror of
https://github.com/langgenius/dify.git
synced 2026-01-14 06:07:33 +08:00
0da4d64d38
- Simplified the SandboxLayer initialization by removing unused parameters and consolidating sandbox creation logic. - Integrated SandboxManager for better lifecycle management of sandboxes during workflow execution. - Updated error handling to ensure proper initialization and cleanup of sandboxes. - Enhanced CommandNode to retrieve sandboxes from SandboxManager, improving sandbox availability checks. - Added unit tests to validate the new sandbox management approach and ensure robust error handling.
64 lines
2.1 KiB
Python
64 lines
2.1 KiB
Python
import logging
|
|
import threading
|
|
from typing import Final
|
|
|
|
from core.virtual_environment.__base.virtual_environment import VirtualEnvironment
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class SandboxManager:
|
|
_lock: Final[threading.Lock] = threading.Lock()
|
|
_sandboxes: dict[str, VirtualEnvironment] = {}
|
|
|
|
@classmethod
|
|
def register(cls, workflow_execution_id: str, sandbox: VirtualEnvironment) -> None:
|
|
if not workflow_execution_id:
|
|
raise ValueError("workflow_execution_id cannot be empty")
|
|
|
|
with cls._lock:
|
|
if workflow_execution_id in cls._sandboxes:
|
|
raise RuntimeError(
|
|
f"Sandbox already registered for workflow_execution_id={workflow_execution_id}. "
|
|
"Call unregister() first if you need to replace it."
|
|
)
|
|
cls._sandboxes[workflow_execution_id] = sandbox
|
|
logger.debug(
|
|
"Registered sandbox for workflow_execution_id=%s, sandbox_id=%s",
|
|
workflow_execution_id,
|
|
sandbox.metadata.id,
|
|
)
|
|
|
|
@classmethod
|
|
def get(cls, workflow_execution_id: str) -> VirtualEnvironment | None:
|
|
with cls._lock:
|
|
return cls._sandboxes.get(workflow_execution_id)
|
|
|
|
@classmethod
|
|
def unregister(cls, workflow_execution_id: str) -> VirtualEnvironment | None:
|
|
with cls._lock:
|
|
sandbox = cls._sandboxes.pop(workflow_execution_id, None)
|
|
if sandbox:
|
|
logger.debug(
|
|
"Unregistered sandbox for workflow_execution_id=%s, sandbox_id=%s",
|
|
workflow_execution_id,
|
|
sandbox.metadata.id,
|
|
)
|
|
return sandbox
|
|
|
|
@classmethod
|
|
def has(cls, workflow_execution_id: str) -> bool:
|
|
with cls._lock:
|
|
return workflow_execution_id in cls._sandboxes
|
|
|
|
@classmethod
|
|
def clear(cls) -> None:
|
|
with cls._lock:
|
|
cls._sandboxes.clear()
|
|
logger.debug("Cleared all registered sandboxes")
|
|
|
|
@classmethod
|
|
def count(cls) -> int:
|
|
with cls._lock:
|
|
return len(cls._sandboxes)
|