mirror of
https://github.com/langgenius/dify.git
synced 2026-02-20 10:01:59 -05:00
18a589003e
- Introduced DraftAppAssetsInitializer for handling draft assets. - Updated SandboxLayer to conditionally set sandbox ID and storage based on workflow version. - Improved asset initialization logging and error handling. - Refactored ArchiveSandboxStorage to support exclusion patterns during archiving. - Modified command and LLM nodes to retrieve sandbox from workflow context, supporting draft workflows.
123 lines
4.9 KiB
Python
123 lines
4.9 KiB
Python
import logging
|
|
|
|
from core.sandbox import AppAssetsInitializer, DifyCliInitializer, SandboxManager
|
|
from core.sandbox.constants import APP_ASSETS_PATH
|
|
from core.sandbox.initializer.app_assets_initializer import DraftAppAssetsInitializer
|
|
from core.sandbox.storage.archive_storage import ArchiveSandboxStorage
|
|
from core.sandbox.vm import SandboxBuilder
|
|
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
|
|
from core.workflow.nodes.base.node import Node
|
|
from models.workflow import Workflow
|
|
from services.app_asset_service import AppAssetService
|
|
from services.sandbox.sandbox_provider_service import SandboxProviderService
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class SandboxInitializationError(Exception):
|
|
pass
|
|
|
|
|
|
class SandboxLayer(GraphEngineLayer):
|
|
def __init__(
|
|
self,
|
|
tenant_id: str,
|
|
app_id: str,
|
|
user_id: str,
|
|
workflow_version: str,
|
|
workflow_execution_id: str,
|
|
) -> None:
|
|
super().__init__()
|
|
self._tenant_id = tenant_id
|
|
self._app_id = app_id
|
|
self._user_id = user_id
|
|
self._workflow_version = workflow_version
|
|
self._workflow_execution_id = workflow_execution_id
|
|
is_draft = self._workflow_version == Workflow.VERSION_DRAFT
|
|
self._sandbox_id = SandboxBuilder.draft_id(self._user_id) if is_draft else self._workflow_execution_id
|
|
self._sandbox_storage = ArchiveSandboxStorage(
|
|
self._tenant_id, self._sandbox_id, exclude_patterns=[APP_ASSETS_PATH] if is_draft else None
|
|
)
|
|
|
|
def on_graph_start(self) -> None:
|
|
try:
|
|
is_draft = self._workflow_version == Workflow.VERSION_DRAFT
|
|
assets = AppAssetService.get_assets(self._tenant_id, self._app_id, self._user_id, is_draft=is_draft)
|
|
if not assets:
|
|
raise ValueError(
|
|
f"No assets found for tid={self._tenant_id}, app_id={self._app_id}, wf={self._workflow_version}"
|
|
)
|
|
|
|
self._assets_id = assets.id
|
|
|
|
if is_draft:
|
|
logger.info(
|
|
"Building draft assets for tenant_id=%s, app_id=%s, workflow_version=%s, assets_id=%s",
|
|
self._tenant_id,
|
|
self._app_id,
|
|
self._workflow_version,
|
|
assets.id,
|
|
)
|
|
AppAssetService.build_assets(self._tenant_id, self._app_id, assets)
|
|
|
|
assets_initializer = (
|
|
DraftAppAssetsInitializer(self._tenant_id, self._app_id, assets.id)
|
|
if is_draft
|
|
else AppAssetsInitializer(self._tenant_id, self._app_id, assets.id)
|
|
)
|
|
|
|
builder = (
|
|
SandboxProviderService.create_sandbox_builder(self._tenant_id)
|
|
.initializer(assets_initializer)
|
|
.initializer(DifyCliInitializer(self._tenant_id, self._user_id, self._app_id, assets.id))
|
|
)
|
|
try:
|
|
sandbox = builder.build()
|
|
logger.info(
|
|
"Sandbox initialized, workflow_execution_id=%s, sandbox_id=%s, sandbox_arch=%s",
|
|
self._sandbox_id,
|
|
sandbox.metadata.id,
|
|
sandbox.metadata.arch,
|
|
)
|
|
except Exception as e:
|
|
raise SandboxInitializationError(f"Failed to build sandbox: {e}") from e
|
|
|
|
SandboxManager.register(self._sandbox_id, sandbox)
|
|
|
|
# mount sandbox files from storage
|
|
mounted = self._sandbox_storage.mount(sandbox)
|
|
logger.info("Sandbox files mount status: %s", mounted)
|
|
|
|
except Exception as e:
|
|
logger.exception("Failed to initialize sandbox")
|
|
raise SandboxInitializationError(f"Failed to initialize sandbox: {e}") from e
|
|
|
|
def on_node_run_start(self, node: Node) -> None:
|
|
# FIXME(Mairuis): should read from workflow run context...
|
|
node.assets_id = self._assets_id
|
|
|
|
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:
|
|
sandbox = SandboxManager.unregister(self._sandbox_id)
|
|
if sandbox is None:
|
|
logger.debug("No sandbox to release for sandbox_id=%s", self._sandbox_id)
|
|
return
|
|
|
|
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", self._sandbox_id)
|
|
except Exception:
|
|
logger.exception("Failed to release sandbox, sandbox_id=%s", self._sandbox_id)
|