Dynamic step config

Author: schustmi

src/zenml/models/v2/core/pipeline_snapshot.py

@@ -208,7 +208,6 @@ class PipelineSnapshotUpdate(BaseUpdate):
     remove_tags: Optional[List[str]] = Field(
         default=None, title="Tags to remove from the snapshot."
     )
-    add_steps: Optional[Dict[str, Step]] = None
 
     @field_validator("name")
     @classmethod

src/zenml/models/v2/core/step_run.py

@@ -28,7 +28,7 @@
 
 from pydantic import ConfigDict, Field
 
-from zenml.config.step_configurations import StepConfiguration, StepSpec
+from zenml.config.step_configurations import Step, StepConfiguration, StepSpec
 from zenml.constants import STR_FIELD_MAX_LENGTH, TEXT_FIELD_MAX_LENGTH
 from zenml.enums import (
     ArtifactSaveType,
@@ -148,6 +148,10 @@ class StepRunRequest(ProjectScopedRequest):
         default=None,
         title="The exception information of the step run.",
     )
+    dynamic_config: Optional["Step"] = Field(
+        title="The dynamic configuration of the step run.",
+        default=None,
+    )
 
     model_config = ConfigDict(protected_namespaces=())
 

src/zenml/orchestrators/step_launcher.py

@@ -108,6 +108,7 @@ def __init__(
         snapshot: PipelineSnapshotResponse,
         step: Step,
         orchestrator_run_id: str,
+        dynamic: bool = False,
     ):
         """Initializes the launcher.
 
@@ -122,6 +123,7 @@ def __init__(
         self._snapshot = snapshot
         self._step = step
         self._orchestrator_run_id = orchestrator_run_id
+        self._dynamic = dynamic
 
         if not snapshot.stack:
             raise RuntimeError(
@@ -306,7 +308,8 @@ def launch(self) -> StepRunResponse:
                 stack=self._stack,
             )
             step_run_request = request_factory.create_request(
-                invocation_id=self._invocation_id
+                invocation_id=self._invocation_id,
+                dynamic_config=self._step if self._dynamic else None,
             )
             step_run_request.logs = logs_model
 

src/zenml/orchestrators/step_run_utils.py

@@ -75,7 +75,9 @@ def has_caching_enabled(self, invocation_id: str) -> bool:
             is_enabled_on_pipeline=self.snapshot.pipeline_configuration.enable_cache,
         )
 
-    def create_request(self, invocation_id: str) -> StepRunRequest:
+    def create_request(
+        self, invocation_id: str, dynamic_config: Optional[Step] = None
+    ) -> StepRunRequest:
         """Create a step run request.
 
         This will only create a request with basic information and will not yet
@@ -95,6 +97,7 @@ def create_request(self, invocation_id: str) -> StepRunRequest:
             status=ExecutionStatus.RUNNING,
             start_time=utc_now(),
             project=Client().active_project.id,
+            dynamic_config=dynamic_config,
         )
 
     def populate_request(
@@ -110,7 +113,10 @@ def populate_request(
                 input resolution. This will be updated in-place with newly
                 fetched step runs.
         """
-        step = self.snapshot.step_configurations[request.name]
+        step = (
+            request.dynamic_config
+            or self.snapshot.step_configurations[request.name]
+        )
 
         input_artifacts = input_utils.resolve_step_inputs(
             step=step,
@@ -133,7 +139,9 @@ def populate_request(
         (
             docstring,
             source_code,
-        ) = self._get_docstring_and_source_code(invocation_id=request.name)
+        ) = self._get_docstring_and_source_code(
+            invocation_id=request.name, step=step
+        )
 
         request.docstring = docstring
         request.source_code = source_code
@@ -174,7 +182,7 @@ def populate_request(
                     request.docstring = cached_step_run.docstring
 
     def _get_docstring_and_source_code(
-        self, invocation_id: str
+        self, invocation_id: str, step: "Step"
     ) -> Tuple[Optional[str], Optional[str]]:
         """Get the docstring and source code for the step.
 
@@ -185,8 +193,6 @@ def _get_docstring_and_source_code(
         Returns:
             The docstring and source code of the step.
         """
-        step = self.snapshot.step_configurations[invocation_id]
-
         try:
             return self._get_docstring_and_source_code_from_step_instance(
                 step=step

src/zenml/pipelines/dynamic/runner.py

@@ -13,12 +13,11 @@
     Tuple,
     Union,
 )
-from zenml.client import Client
-from zenml.config.compiler import Compiler
-
 from uuid import UUID
+
 from zenml import ExternalArtifact
 from zenml.client import Client
+from zenml.config.compiler import Compiler
 from zenml.config.step_configurations import Step
 from zenml.exceptions import RunStoppedException
 from zenml.logger import get_logger
@@ -27,7 +26,6 @@
     ArtifactVersionResponse,
     PipelineRunResponse,
     PipelineSnapshotResponse,
-    PipelineSnapshotUpdate,
 )
 from zenml.models.v2.core.step_run import StepRunResponse
 from zenml.orchestrators.publish_utils import (
@@ -169,18 +167,19 @@ def run_step_sync(
         compiled_step, invocation_id = _compile_step(
             self.pipeline, step, id, upstream_steps, inputs
         )
-        updated_snapshot = Client().zen_store.update_snapshot(
-            self._snapshot.id,
-            snapshot_update=PipelineSnapshotUpdate(
-                add_steps={invocation_id: compiled_step}
-            ),
-        )
-        step_config = updated_snapshot.step_configurations[invocation_id]
+        # updated_snapshot = Client().zen_store.update_snapshot(
+        #     self._snapshot.id,
+        #     snapshot_update=PipelineSnapshotUpdate(
+        #         add_steps={invocation_id: compiled_step}
+        #     ),
+        # )
+        # step_config = updated_snapshot.step_configurations[invocation_id]
         step_run = _run_step_sync(
-            snapshot=updated_snapshot,
-            step=step_config,
+            snapshot=self._snapshot,
+            step=compiled_step,
             orchestrator_run_id=self._run.orchestrator_run_id,
-            retry=_should_retry_locally(step_config),
+            retry=_should_retry_locally(compiled_step),
+            dynamic=True,
         )
         return _load_step_result(step_run.id)
 
@@ -198,31 +197,20 @@ def run_step_in_thread(
         compiled_step, invocation_id = _compile_step(
             self.pipeline, step, id, upstream_steps, inputs
         )
-        updated_snapshot = Client().zen_store.update_snapshot(
-            self._snapshot.id,
-            snapshot_update=PipelineSnapshotUpdate(
-                add_steps={invocation_id: compiled_step}
-            ),
-        )
-        step_config = updated_snapshot.step_configurations[invocation_id]
 
         def _run() -> StepRunResult:
             step_run = _run_step_sync(
-                snapshot=updated_snapshot,
-                step=step_config,
+                snapshot=self._snapshot,
+                step=compiled_step,
                 orchestrator_run_id=self._run.orchestrator_run_id,
-                retry=_should_retry_locally(step_config),
+                retry=_should_retry_locally(compiled_step),
+                dynamic=True,
             )
             return _load_step_result(step_run.id)
 
         ctx = contextvars.copy_context()
-        future = self._executor.submit(
-            ctx.run, _run
-        )
-        return StepRunResultFuture(
-            wrapped=future, invocation_id=invocation_id
-        )
-
+        future = self._executor.submit(ctx.run, _run)
+        return StepRunResultFuture(wrapped=future, invocation_id=invocation_id)
 
 
 def _prepare_step_run(
@@ -318,12 +306,14 @@ def _run_step_sync(
     step: "Step",
     orchestrator_run_id: str,
     retry: bool = False,
+    dynamic: bool = False,
 ) -> StepRunResponse:
     def _launch_step() -> StepRunResponse:
         launcher = StepLauncher(
             snapshot=snapshot,
             step=step,
             orchestrator_run_id=orchestrator_run_id,
+            dynamic=dynamic,
         )
         return launcher.launch()
 
@@ -362,7 +352,7 @@ def _launch_step() -> StepRunResponse:
                     raise
             else:
                 break
-    
+
     return step_run
 
 
@@ -377,7 +367,7 @@ def _convert_output_artifact(
             step_name=step_run.name,
             **artifact.model_dump(),
         )
-        
+
     output_artifacts = step_run.regular_outputs
     if len(output_artifacts) == 0:
         return None
@@ -399,16 +389,19 @@ def _should_retry_locally(step: "Step") -> bool:
         return True
     else:
         # Running out of process with the orchestrator
-        return not Client().active_stack.orchestrator.config.handles_step_retries
+        return (
+            not Client().active_stack.orchestrator.config.handles_step_retries
+        )
+
 
 def _runs_in_process(step: "Step") -> bool:
     if step.config.step_operator:
         return False
-    
+
     if not Client().active_stack.orchestrator.supports_dynamic_out_of_process_steps:
         return False
 
     if step.config.in_process is False:
         return False
-    
-    return True
+
+    return True

src/zenml/zen_stores/schemas/pipeline_snapshot_schemas.py

@@ -586,6 +586,7 @@ class StepConfigurationSchema(BaseSchema, table=True):
     __table_args__ = (
         UniqueConstraint(
             "snapshot_id",
+            "step_run_id",
             "name",
             name="unique_step_name_for_snapshot",
         ),
@@ -608,5 +609,13 @@ class StepConfigurationSchema(BaseSchema, table=True):
         source_column="snapshot_id",
         target_column="id",
         ondelete="CASCADE",
-        nullable=False,
+        nullable=True,
+    )
+    step_run_id: UUID = build_foreign_key_field(
+        source=__tablename__,
+        target="step_run",
+        source_column="step_run_id",
+        target_column="id",
+        ondelete="CASCADE",
+        nullable=True,
     )

src/zenml/zen_stores/schemas/step_run_schemas.py

@@ -216,6 +216,9 @@ class StepRunSchema(NamedSchema, RunMetadataInterface, table=True):
             ),
         )
     )
+    dynamic_step_configuration_schema: Optional["StepConfigurationSchema"] = (
+        Relationship()
+    )
 
     model_config = ConfigDict(protected_namespaces=())  # type: ignore[assignment]
 
@@ -344,7 +347,10 @@ def get_step_configuration(self) -> Step:
         step = None
 
         if self.snapshot is not None:
-            if self.step_configuration_schema:
+            if (
+                self.dynamic_step_configuration_schema
+                or self.step_configuration_schema
+            ):
                 pipeline_configuration = (
                     PipelineConfiguration.model_validate_json(
                         self.snapshot.pipeline_configuration
@@ -354,8 +360,12 @@ def get_step_configuration(self) -> Step:
                     start_time=self.pipeline_run.start_time,
                     inplace=True,
                 )
+                config_schema = (
+                    self.dynamic_step_configuration_schema
+                    or self.step_configuration_schema
+                )
                 step = Step.from_dict(
-                    json.loads(self.step_configuration_schema.config),
+                    json.loads(config_schema.config),
                     pipeline_configuration=pipeline_configuration,
                 )
         if not step and self.step_configuration:

src/zenml/zen_stores/sql_zen_store.py

@@ -5167,34 +5167,6 @@ def update_snapshot(
                 session=session,
             )
 
-            if snapshot_update.add_steps:
-                # TODO: this doesn't work for scheduled runs that reuse a
-                # snapshot. For that, we'd have to separate config from the
-                # snapshot, which is something we want to do anyway.
-                if not snapshot.is_dynamic:
-                    raise ValueError(
-                        "Cannot dynamically update steps of a static snapshot."
-                    )
-
-                # TODO: race conditions
-                current_index = len(snapshot.step_configurations)
-                for index, (step_name, step_configuration) in enumerate(
-                    snapshot_update.add_steps.items()
-                ):
-                    step_configuration_schema = StepConfigurationSchema(
-                        index=current_index + index,
-                        name=step_name,
-                        # Don't include the merged config in the step
-                        # configurations, we reconstruct it in the `to_model` method
-                        # using the pipeline configuration.
-                        config=step_configuration.model_dump_json(
-                            exclude={"config"}
-                        ),
-                        snapshot_id=snapshot.id,
-                    )
-                    session.add(step_configuration_schema)
-                session.commit()
-
             session.refresh(snapshot)
             return snapshot.to_model(
                 include_metadata=True, include_resources=True
@@ -9677,7 +9649,10 @@ def create_run_step(self, step_run: StepRunRequest) -> StepRunResponse:
                 session=session,
                 reference_type="original step run",
             )
-            step_config = run.get_step_configuration(step_name=step_run.name)
+            step_config = (
+                step_run.dynamic_config
+                or run.get_step_configuration(step_name=step_run.name)
+            )
 
             # Release the read locks of the previous two queries before we
             # try to acquire more exclusive locks
@@ -9929,6 +9904,20 @@ def create_run_step(self, step_run: StepRunRequest) -> StepRunResponse:
                     pipeline_run_id=step_run.pipeline_run_id, session=session
                 )
 
+            if step_run.dynamic_config:
+                step_configuration_schema = StepConfigurationSchema(
+                    index=0,
+                    name=step_run.name,
+                    # Don't include the merged config in the step
+                    # configurations, we reconstruct it in the `to_model` method
+                    # using the pipeline configuration.
+                    config=step_run.dynamic_config.model_dump_json(
+                        exclude={"config"}
+                    ),
+                    step_run_id=step_schema.id,
+                )
+                session.add(step_configuration_schema)
+
             session.commit()
             session.refresh(
                 step_schema, ["input_artifacts", "output_artifacts"]