Skip to content

refactor(): Draft context manager #1634

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
benflexcompute merged 10 commits into from
Dec 11, 2025
Merged

refactor(): Draft context manager #1634

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
b956a4a
Stage 1: Refactor EntityRegistry to type-based storage and add view()…
benflexcompute Dec 7, 2025
094d239
Stage 2: Add EntityRegistry.from_entity_info() for DraftContext workflow
benflexcompute Dec 7, 2025
0891a02
Stage 3: DraftContext entity isolation with deep copy and registry.vi…
benflexcompute Dec 7, 2025
c64b678
Stage 4: Clean entity merging on upload with DraftContext support
benflexcompute Dec 7, 2025
6547e3e
Stage 4.1: Added legacy entity status route compatability test with m…
benflexcompute Dec 9, 2025
aea1829
Stage 5: EntityRegistry replaces entity_info dict-based entity lookup
benflexcompute Dec 9, 2025
0b9275c
Refactor entity materialization: Replace entity_pool with EntityRegi…
benflexcompute Dec 10, 2025
ad0c292
Start self review
benflexcompute Dec 10, 2025
fd4e40f
Merge remote-tracking branch 'origin/main' into BenY/DraftContextManager
benflexcompute Dec 10, 2025
2170e27
self reviewed
benflexcompute Dec 10, 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
3 changes: 2 additions & 1 deletion flow360/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
from flow360.cli.api_set_func import configure_caller as configure
from flow360.component.case import Case
from flow360.component.geometry import Geometry
from flow360.component.project import Project
from flow360.component.project import Project, create_draft
from flow360.component.simulation import migration, services
from flow360.component.simulation import units as u
from flow360.component.simulation.entity_operation import Transformation
Expand Down Expand Up @@ -199,6 +199,7 @@
"GeometryRefinement",
"Env",
"Case",
"create_draft",
"AngleBasedRefinement",
"AspectRatioBasedRefinement",
"ProjectAnisoSpacing",
Expand Down
10 changes: 10 additions & 0 deletions flow360/component/geometry.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -583,6 +583,16 @@ def __getitem__(self, key: str):
Get the entity by name.
`key` is the name of the entity or the naming pattern if wildcard is used.
"""
# pylint: disable=import-outside-toplevel
from flow360.component.simulation.draft_context import get_active_draft

if get_active_draft() is not None:
log.warning(
"Accessing entities via asset[key] while a DraftContext is active. "
"Use draft.surfaces[key] or draft.body_groups[key] instead to ensure "
"modifications are tracked in the draft's entity_info."
)

if isinstance(key, str) is False:
raise Flow360ValueError(f"Entity naming pattern: {key} is not a string.")

Expand Down
88 changes: 88 additions & 0 deletions flow360/component/project.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,13 +24,19 @@
VolumeMeshInterfaceV2,
)
from flow360.component.project_utils import (
apply_geometry_grouping_overrides,
get_project_records,
set_up_params_for_uploading,
show_projects_with_keyword_filter,
upload_imported_surfaces_to_draft,
validate_params_with_context,
)
from flow360.component.resource_base import Flow360Resource
from flow360.component.simulation.draft_context.context import (
DraftContext,
get_active_draft,
)
from flow360.component.simulation.entity_info import GeometryEntityInfo
from flow360.component.simulation.folder import Folder
from flow360.component.simulation.simulation_params import SimulationParams
from flow360.component.simulation.unit_system import LengthType
Expand All @@ -51,6 +57,7 @@
from flow360.exceptions import (
Flow360ConfigError,
Flow360FileError,
Flow360RuntimeError,
Flow360ValueError,
Flow360WebError,
)
Expand Down Expand Up @@ -80,6 +87,82 @@ class RootType(Enum):
VOLUME_MESH = "VolumeMesh"


def create_draft(
*,
new_run_from: Union[Geometry, SurfaceMeshV2, VolumeMeshV2],
face_grouping: Optional[str] = None,
edge_grouping: Optional[str] = None,
) -> DraftContext:
"""Factory helper used by end users (`with fl.create_draft() as draft`).

Creates a DraftContext with a deep copy of the asset's entity_info,
providing entity isolation so modifications in the draft don't affect
the original asset.
"""

# region -----------------------------Private implementations Below-----------------------------

def _deep_copy_entity_info(entity_info):
"""Create a deep copy of entity_info via model_dump + model_validate.

This ensures DraftContext has an independent copy of entity_info,
so modifications don't affect the original asset.
"""
entity_info_dict = entity_info.model_dump(mode="json")
return type(entity_info).model_validate(entity_info_dict)

def _inform_grouping_selections(entity_info) -> None:
"""Inform the user about the grouping selections made on the entity provider cloud asset."""

if isinstance(entity_info, GeometryEntityInfo):
applied_grouping = {
"face": entity_info.face_group_tag,
"edge": entity_info.edge_group_tag,
"body": entity_info.body_group_tag,
}
if face_grouping is not None or edge_grouping is not None:
applied_grouping = apply_geometry_grouping_overrides(
entity_info, face_grouping, edge_grouping
)
# If tags were None, fall back to defaults for logging purposes.
# pylint:disable = protected-access
face_tag = applied_grouping.get("face") or entity_info._get_default_grouping_tag("face")
edge_tag = applied_grouping.get("edge")
if edge_tag is None and entity_info.edge_attribute_names:
edge_tag = entity_info._get_default_grouping_tag("edge")

log.info(
"Creating draft with geometry grouping:\n"
" faces: %s\n"
" edges: %s\n"
"To change grouping, call:\n"
" fl.create_draft(face_grouping='%s', edge_grouping='%s', ...)",
face_tag,
edge_tag,
face_tag,
edge_tag,
)
elif face_grouping is not None or edge_grouping is not None:
log.info(
"Grouping override ignored: only geometry assets support face/edge/body regrouping."
)

# endregion ------------------------------------------------------------------------------------

if not isinstance(new_run_from, AssetBase):
raise Flow360RuntimeError("create_draft expects a cloud asset instance as `new_run_from`.")

# Deep copy entity_info for draft isolation
entity_info_copy = _deep_copy_entity_info(new_run_from.entity_info)

# Apply grouping overrides to the copy (not the original)
_inform_grouping_selections(entity_info_copy)

return DraftContext(
entity_info=entity_info_copy,
)


class ProjectMeta(pd.BaseModel, extra="allow"):
"""
Metadata class for a project.
Expand Down Expand Up @@ -1400,6 +1483,10 @@ def _run(
if use_geometry_AI is True and use_beta_mesher is False:
raise Flow360ValueError("Enabling Geometry AI requires also enabling beta mesher.")

# Check if there's an active DraftContext and get its entity_info
active_draft = get_active_draft()
draft_entity_info = active_draft._entity_info if active_draft is not None else None

root_asset = self._root_asset
if interpolate_to_mesh is not None:
project_vm = Project.from_cloud(project_id=interpolate_to_mesh.project_id)
Expand All @@ -1411,6 +1498,7 @@ def _run(
length_unit=self.length_unit,
use_beta_mesher=use_beta_mesher,
use_geometry_AI=use_geometry_AI,
draft_entity_info=draft_entity_info,
)

params, errors = validate_params_with_context(
Expand Down
129 changes: 115 additions & 14 deletions flow360/component/project_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,11 @@
from flow360.cloud.rest_api import RestApi
from flow360.component.interfaces import ProjectInterface
from flow360.component.simulation import services
from flow360.component.simulation.entity_info import DraftEntityTypes, EntityInfoModel
from flow360.component.simulation.entity_info import (
DraftEntityTypes,
EntityInfoModel,
GeometryEntityInfo,
)
from flow360.component.simulation.framework.base_model import Flow360BaseModel
from flow360.component.simulation.framework.entity_base import EntityList
from flow360.component.simulation.framework.param_utils import AssetCache
Expand All @@ -32,10 +36,38 @@
from flow360.component.simulation.user_code.core.types import save_user_variables
from flow360.component.simulation.utils import model_attribute_unlock
from flow360.component.utils import parse_datetime
from flow360.exceptions import Flow360ConfigurationError
from flow360.exceptions import Flow360ConfigurationError, Flow360ValueError
from flow360.log import log


def apply_geometry_grouping_overrides(
entity_info: GeometryEntityInfo,
face_grouping: Optional[str],
edge_grouping: Optional[str],
) -> dict[str, Optional[str]]:
"""Apply explicit face/edge grouping overrides onto geometry entity info."""

def _validate_tag(tag: str, available: list[str], kind: str) -> str:
if available and tag not in available: # pylint:disable=unsupported-membership-test
raise Flow360ValueError(
f"Invalid {kind} grouping tag '{tag}'. Available tags: {available}."
)
return tag

if face_grouping is not None:
face_tag = _validate_tag(face_grouping, entity_info.face_attribute_names, "face")
entity_info._group_entity_by_tag("face", face_tag) # pylint:disable=protected-access
if edge_grouping is not None and entity_info.edge_attribute_names:
edge_tag = _validate_tag(edge_grouping, entity_info.edge_attribute_names, "edge")
entity_info._group_entity_by_tag("edge", edge_tag) # pylint:disable=protected-access

return {
"face": entity_info.face_group_tag,
"edge": entity_info.edge_group_tag,
"body": entity_info.body_group_tag, # Not used since customized body grouping is not supported yet
}


class AssetStatistics(pd.BaseModel):
"""Statistics for an asset"""

Expand Down Expand Up @@ -223,6 +255,9 @@ def _set_up_params_non_persistent_entity_info(entity_info, params: SimulationPar
"""
Setting up non-persistent entities (AKA draft entities) in params.
Add the ones used to the entity info.

LEGACY: This function is used for the legacy workflow (without DraftContext).
For DraftContext workflow, use _merge_draft_entities_from_params() instead.
"""

entity_registry = params.used_entity_registry
Expand All @@ -231,13 +266,56 @@ def _set_up_params_non_persistent_entity_info(entity_info, params: SimulationPar
draft_type_union = get_args(DraftEntityTypes)[0]
draft_type_list = get_args(draft_type_union)
for draft_type in draft_type_list:
draft_entities = entity_registry.find_by_type(draft_type)
draft_entities = list(entity_registry.view(draft_type))
for draft_entity in draft_entities:
if draft_entity not in entity_info.draft_entities:
entity_info.draft_entities.append(draft_entity)
return entity_info


def _merge_draft_entities_from_params(
entity_info: EntityInfoModel,
params: SimulationParams,
) -> EntityInfoModel:
"""
Collect draft entities from params.used_entity_registry and merge into entity_info.

This function implements the merging logic for the DraftContext workflow:
- If a draft entity already exists in entity_info (by ID), use entity_info version (source of truth)
- If a draft entity is new (not in entity_info), add it from params

This ensures that:
1. Entities managed by DraftContext retain their modifications
2. New entities created by the user during simulation setup are captured

Parameters:
entity_info: The entity_info to merge into (typically from DraftContext)
params: The SimulationParams containing used_entity_registry

Returns:
EntityInfoModel: The updated entity_info with merged draft entities
"""
used_registry = params.used_entity_registry

# Get all draft entity types from the DraftEntityTypes annotation
draft_type_union = get_args(DraftEntityTypes)[0]
draft_type_list = get_args(draft_type_union)

# Build a set of IDs already in entity_info for quick lookup (Draft entities have unique UUIDs)
existing_ids = {e.private_attribute_id for e in entity_info.draft_entities}

for draft_type in draft_type_list:
draft_entities_used = list(used_registry.view(draft_type))
for draft_entity in draft_entities_used:
# Only add if not already in entity_info (by ID)
# If already present, entity_info version is source of truth - keep it as is
if draft_entity.private_attribute_id not in existing_ids:
entity_info.draft_entities.append(draft_entity)
existing_ids.add(draft_entity.private_attribute_id)

return entity_info


def _update_entity_grouping_tags(entity_info, params: SimulationParams) -> EntityInfoModel:
"""
Update the entity grouping tags in params to resolve possible conflicts
Expand Down Expand Up @@ -373,15 +451,26 @@ def _set_up_monitor_output_from_stopping_criterion(params: SimulationParams):
return params


def set_up_params_for_uploading(
def set_up_params_for_uploading( # pylint: disable=too-many-arguments
root_asset,
length_unit: LengthType,
params: SimulationParams,
use_beta_mesher: bool,
use_geometry_AI: bool, # pylint: disable=invalid-name
draft_entity_info: Optional[EntityInfoModel] = None,
) -> SimulationParams:
"""
Set up params before submitting the draft.

Parameters:
root_asset: The root asset (Geometry, SurfaceMesh, or VolumeMesh).
length_unit: The project length unit.
params: The SimulationParams to set up.
use_beta_mesher: Whether to use the beta mesher.
use_geometry_AI: Whether to use Geometry AI.
draft_entity_info: Optional entity_info from DraftContext. When provided,
this is used as the source of truth for entities instead of root_asset.entity_info (legacy behavior).
This enables proper entity isolation for the DraftContext workflow.
"""

with model_attribute_unlock(params.private_attribute_asset_cache, "project_length_unit"):
Expand All @@ -397,20 +486,32 @@ def set_up_params_for_uploading(
use_geometry_AI if use_geometry_AI else False
)

# User may have made modifications to the entities which is recorded in asset's entity registry
# We need to reflect these changes.
root_asset.entity_info.update_persistent_entities(
asset_entity_registry=root_asset.internal_registry
)
if draft_entity_info is not None:
# New DraftContext workflow: use draft's entity_info as source of truth
# Merge draft entities from params.used_entity_registry into draft_entity_info
entity_info = _merge_draft_entities_from_params(draft_entity_info, params)

# Update entity grouping tags if needed
# (back compatibility, since the grouping should already have been captured in the draft_entity_info)
entity_info = _update_entity_grouping_tags(entity_info, params)
else:
# Legacy workflow (without DraftContext): use root_asset.entity_info
# User may have made modifications to the entities which is recorded in asset's entity registry
# We need to reflect these changes.
root_asset.entity_info.update_persistent_entities(
asset_entity_registry=root_asset.internal_registry
)

# Check if there are any new draft entities that have been added in the params by the user
entity_info = _set_up_params_non_persistent_entity_info(root_asset.entity_info, params)
# Check if there are any new draft entities that have been added in the params by the user
entity_info = _set_up_params_non_persistent_entity_info(root_asset.entity_info, params)

# If the customer just load the param without re-specify the same set of entity grouping tags,
# we need to update the entity grouping tags to the ones in the SimulationParams.
entity_info = _update_entity_grouping_tags(entity_info, params)
# If the customer just load the param without re-specify the same set of entity grouping tags,
# we need to update the entity grouping tags to the ones in the SimulationParams.
entity_info = _update_entity_grouping_tags(entity_info, params)

with model_attribute_unlock(params.private_attribute_asset_cache, "project_entity_info"):
# At this point the draft entity info has replaced the SimulationParams's entity info.
# So the validation afterwards does not require the access to the draft entity info anymore.
params.private_attribute_asset_cache.project_entity_info = entity_info
# Replace the ghost surfaces in the SimulationParams by the real ghost ones from asset metadata.
# This has to be done after `project_entity_info` is properly set.
Expand Down
8 changes: 8 additions & 0 deletions flow360/component/simulation/draft_context/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,8 @@
"""Public interface for the draft context subsystem."""

from flow360.component.simulation.draft_context.context import (
DraftContext,
get_active_draft,
)

__all__ = ["DraftContext", "get_active_draft"]
Loading
Loading