mirror of
https://github.com/langgenius/dify.git
synced 2026-02-27 20:02:42 -05:00
95 lines
3.6 KiB
Python
95 lines
3.6 KiB
Python
import logging
|
|
|
|
from core.sandbox import SandboxManager
|
|
from core.sandbox.storage.sandbox_storage import SandboxStorage
|
|
from core.virtual_environment.__base.virtual_environment import VirtualEnvironment
|
|
from core.workflow.graph_engine.layers.base import GraphEngineLayer
|
|
from core.workflow.graph_events.base import GraphEngineEvent
|
|
from core.workflow.graph_events.graph import GraphRunPausedEvent
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class SandboxInitializationError(Exception):
|
|
pass
|
|
|
|
|
|
class SandboxLayer(GraphEngineLayer):
|
|
def __init__(self, tenant_id: str, app_id: str, sandbox_id: str, sandbox_storage: SandboxStorage) -> None:
|
|
super().__init__()
|
|
self._tenant_id = tenant_id
|
|
self._app_id = app_id
|
|
self._sandbox_id = sandbox_id
|
|
self._sandbox_storage = sandbox_storage
|
|
|
|
@property
|
|
def sandbox(self) -> VirtualEnvironment:
|
|
sandbox = SandboxManager.get(self._sandbox_id)
|
|
if sandbox is None:
|
|
raise RuntimeError(f"Sandbox not found or not initialized for sandbox_id={self._sandbox_id}")
|
|
return sandbox
|
|
|
|
def on_graph_start(self) -> None:
|
|
try:
|
|
# Initialize sandbox
|
|
from core.sandbox import AppAssetsInitializer, DifyCliInitializer
|
|
from services.sandbox.sandbox_provider_service import SandboxProviderService
|
|
|
|
logger.info("Initializing sandbox for tenant_id=%s, app_id=%s", self._tenant_id, self._app_id)
|
|
|
|
builder = (
|
|
SandboxProviderService.create_sandbox_builder(self._tenant_id)
|
|
.initializer(DifyCliInitializer())
|
|
.initializer(AppAssetsInitializer(self._tenant_id, self._app_id))
|
|
)
|
|
sandbox = builder.build()
|
|
|
|
SandboxManager.register(self._sandbox_id, sandbox)
|
|
logger.info(
|
|
"Sandbox initialized, workflow_execution_id=%s, sandbox_id=%s, sandbox_arch=%s",
|
|
self._sandbox_id,
|
|
sandbox.metadata.id,
|
|
sandbox.metadata.arch,
|
|
)
|
|
|
|
# Check if sandbox is initialized
|
|
if self._sandbox_storage.mount(sandbox):
|
|
logger.info("Sandbox files restored, sandbox_id=%s", self._sandbox_id)
|
|
except Exception as e:
|
|
logger.exception("Failed to initialize sandbox")
|
|
raise SandboxInitializationError(f"Failed to initialize sandbox: {e}") from e
|
|
|
|
def on_event(self, event: GraphEngineEvent) -> None:
|
|
# TODO: handle graph run paused event
|
|
if not isinstance(event, GraphRunPausedEvent):
|
|
return
|
|
|
|
def on_graph_end(self, error: Exception | None) -> None:
|
|
if self._sandbox_id is None:
|
|
logger.debug("No workflow_execution_id set, nothing to release")
|
|
return
|
|
|
|
sandbox = SandboxManager.unregister(self._sandbox_id)
|
|
if sandbox is None:
|
|
logger.debug("No sandbox to release for sandbox_id=%s", self._sandbox_id)
|
|
return
|
|
|
|
sandbox_id = sandbox.metadata.id
|
|
logger.info(
|
|
"Releasing sandbox, workflow_execution_id=%s, sandbox_id=%s",
|
|
self._sandbox_id,
|
|
sandbox_id,
|
|
)
|
|
|
|
try:
|
|
self._sandbox_storage.unmount(sandbox)
|
|
logger.info("Sandbox files persisted, sandbox_id=%s", self._sandbox_id)
|
|
except Exception:
|
|
logger.exception("Failed to persist sandbox files, sandbox_id=%s", self._sandbox_id)
|
|
|
|
try:
|
|
sandbox.release_environment()
|
|
logger.info("Sandbox released, sandbox_id=%s", sandbox_id)
|
|
except Exception:
|
|
logger.exception("Failed to release sandbox, sandbox_id=%s", sandbox_id)
|