[dagster-airlift][jobs 5/n] Definitions.execute_job_in_process

Author: dpeng817

python_modules/dagster/dagster/_core/definitions/definitions_class.py

@@ -44,6 +44,8 @@
 from dagster._utils.warnings import disable_dagster_warnings
 
 if TYPE_CHECKING:
+    from dagster._core.definitions.run_config import RunConfig
+    from dagster._core.execution.execute_in_process_result import ExecuteInProcessResult
     from dagster._core.storage.asset_value_loader import AssetValueLoader
 
 
@@ -795,3 +797,75 @@ def asset2(_): ...
             *[d for d in self.assets or [] if not isinstance(d, (AssetsDefinition, AssetSpec))],
         ]
         return replace(self, assets=assets)
+
+    def execute_job_in_process(
+        self,
+        job_name: str,
+        run_config: Optional[Union[Mapping[str, Any], "RunConfig"]] = None,
+        instance: Optional["DagsterInstance"] = None,
+        partition_key: Optional[str] = None,
+        raise_on_error: bool = True,
+        op_selection: Optional[Sequence[str]] = None,
+        asset_selection: Optional[Sequence[AssetKey]] = None,
+        run_id: Optional[str] = None,
+        input_values: Optional[Mapping[str, object]] = None,
+        tags: Optional[Mapping[str, str]] = None,
+        resources: Optional[Mapping[str, object]] = None,
+    ) -> "ExecuteInProcessResult":
+        from dagster._core.definitions.job_base import RepoBackedJob
+        from dagster._core.execution.execute_in_process import (
+            core_execute_in_process,
+            merge_run_tags,
+            type_check_and_normalize_args,
+        )
+
+        run_config, op_selection, asset_selection, resource_defs, partition_key, input_values = (
+            type_check_and_normalize_args(
+                run_config=run_config,
+                partition_key=partition_key,
+                op_selection=op_selection,
+                asset_selection=asset_selection,
+                input_values=input_values,
+                resources=resources,
+            )
+        )
+
+        job_def = check.not_none(get_job_from_defs(job_name, self)).as_ephemeral_job(
+            resource_defs=resource_defs,
+            input_values=check.opt_mapping_param(
+                input_values, "input_values", key_type=str, value_type=object
+            ),
+            op_selection=op_selection,
+            asset_selection=asset_selection,
+        )
+        new_job_list = [job for job in (self.jobs or []) if job.name != job_name] + [job_def]
+        new_defs_obj = replace(self, jobs=new_job_list)
+        resolved_repo = new_defs_obj.get_repository_def()
+        wrapped_job = RepoBackedJob(job_name=job_name, repository_def=resolved_repo)
+        return core_execute_in_process(
+            job=wrapped_job,
+            run_config=run_config,
+            instance=instance,
+            output_capturing_enabled=True,
+            raise_on_error=raise_on_error,
+            run_tags=merge_run_tags(
+                job_def=job_def,
+                partition_key=partition_key,
+                tags=tags,
+                asset_selection=asset_selection,
+                instance=instance,
+                run_config=run_config,
+            ),
+            run_id=run_id,
+            asset_selection=frozenset(asset_selection),
+        )
+
+
+def get_job_from_defs(
+    name: str, defs: Definitions
+) -> Optional[Union[JobDefinition, UnresolvedAssetJobDefinition]]:
+    """Get the job from the definitions by its name."""
+    return next(
+        iter(job for job in (defs.jobs or []) if job.name == name),
+        None,
+    )

python_modules/dagster/dagster/_core/definitions/job_base.py

@@ -95,3 +95,49 @@ def asset_selection(self) -> Optional[AbstractSet[AssetKey]]:
     @property
     def asset_check_selection(self) -> Optional[AbstractSet[AssetCheckKey]]:
         return self._job_def.asset_check_selection
+
+
+class RepoBackedJob(IJob):
+    def __init__(
+        self,
+        job_name: str,
+        repository_def: "RepositoryDefinition",
+    ):
+        self._job_name = job_name
+        self._repository_def = repository_def
+
+    def get_definition(self) -> "JobDefinition":
+        return self._repository_def.get_job(self._job_name)
+
+    def get_repository_definition(self) -> Optional["RepositoryDefinition"]:
+        return self._repository_def
+
+    @property
+    def op_selection(self) -> Optional[AbstractSet[str]]:
+        return self.get_definition().op_selection
+
+    @property
+    def asset_selection(self) -> Optional[AbstractSet[AssetKey]]:
+        return self.get_definition().asset_selection
+
+    @property
+    def asset_check_selection(self) -> Optional[AbstractSet[AssetCheckKey]]:
+        return self.get_definition().asset_check_selection
+
+    def get_subset(
+        self,
+        *,
+        op_selection: Optional[Iterable[str]] = None,
+        asset_selection: Optional[AbstractSet[AssetKey]] = None,
+        asset_check_selection: Optional[AbstractSet[AssetCheckKey]] = None,
+    ) -> "RepoBackedJob":
+        op_selection = set(op_selection) if op_selection else None
+        return RepoBackedJob(
+            self._job_name,
+            self._repository_def.get_subset(
+                job_name=self._job_name,
+                op_selection=op_selection,
+                asset_selection=asset_selection,
+                asset_check_selection=asset_check_selection,
+            ),
+        )

python_modules/dagster/dagster/_core/definitions/job_definition.py

@@ -666,34 +666,61 @@ def execute_in_process(
             :py:class:`~dagster.ExecuteInProcessResult`
 
         """
-        from dagster._core.definitions.executor_definition import execute_in_process_executor
-        from dagster._core.definitions.run_config import convert_config_input
-        from dagster._core.execution.build_resources import wrap_resources_for_execution
-        from dagster._core.execution.execute_in_process import core_execute_in_process
-
-        run_config = check.opt_mapping_param(convert_config_input(run_config), "run_config")
-        op_selection = check.opt_sequence_param(op_selection, "op_selection", str)
-        asset_selection = check.opt_sequence_param(asset_selection, "asset_selection", AssetKey)
-        resources = check.opt_mapping_param(resources, "resources", key_type=str)
+        from dagster._core.definitions.job_base import InMemoryJob
+        from dagster._core.execution.execute_in_process import (
+            core_execute_in_process,
+            merge_run_tags,
+            type_check_and_normalize_args,
+        )
 
-        resource_defs = wrap_resources_for_execution(resources)
+        run_config, op_selection, asset_selection, resource_defs, partition_key, input_values = (
+            type_check_and_normalize_args(
+                run_config=run_config,
+                partition_key=partition_key,
+                op_selection=op_selection,
+                asset_selection=asset_selection,
+                input_values=input_values,
+                resources=resources,
+            )
+        )
 
-        check.invariant(
-            not (op_selection and asset_selection),
-            "op_selection and asset_selection cannot both be provided as args to"
-            " execute_in_process",
+        ephemeral_job = self.as_ephemeral_job(
+            resource_defs=resource_defs,
+            input_values=input_values,
+            op_selection=op_selection,
+            asset_selection=asset_selection,
         )
 
-        partition_key = check.opt_str_param(partition_key, "partition_key")
-        input_values = check.opt_mapping_param(input_values, "input_values")
+        wrapped_job = InMemoryJob(job_def=ephemeral_job)
+        return core_execute_in_process(
+            job=wrapped_job,
+            run_config=run_config,
+            instance=instance,
+            output_capturing_enabled=True,
+            raise_on_error=raise_on_error,
+            run_tags=merge_run_tags(
+                job_def=self,
+                partition_key=partition_key,
+                tags=tags,
+                asset_selection=asset_selection,
+                instance=instance,
+                run_config=run_config,
+            ),
+            run_id=run_id,
+            asset_selection=frozenset(asset_selection),
+        )
 
-        # Combine provided input values at execute_in_process with input values
-        # provided to the definition. Input values provided at
-        # execute_in_process will override those provided on the definition.
-        input_values = merge_dicts(self.input_values, input_values)
+    def as_ephemeral_job(
+        self,
+        resource_defs: Mapping[str, ResourceDefinition],
+        input_values: Mapping[str, object],
+        op_selection: Optional[Sequence[str]] = None,
+        asset_selection: Optional[Sequence[AssetKey]] = None,
+    ) -> "JobDefinition":
+        from dagster._core.definitions.executor_definition import execute_in_process_executor
 
         bound_resource_defs = dict(self.resource_defs)
-        ephemeral_job = JobDefinition.dagster_internal_init(
+        return JobDefinition.dagster_internal_init(
             name=self._name,
             graph_def=self._graph_def,
             resource_defs={**_swap_default_io_man(bound_resource_defs, self), **resource_defs},
@@ -705,54 +732,17 @@ def execute_in_process(
             run_tags=self._run_tags,
             op_retry_policy=self._op_retry_policy,
             asset_layer=self.asset_layer,
-            input_values=input_values,
+            input_values=merge_dicts(self.input_values, input_values),
             description=self.description,
             partitions_def=self.partitions_def,
             metadata=self.metadata,
             _subset_selection_data=None,  # this is added below
             _was_explicitly_provided_resources=True,
-        )
-
-        ephemeral_job = ephemeral_job.get_subset(
+        ).get_subset(
             op_selection=op_selection,
             asset_selection=frozenset(asset_selection) if asset_selection else None,
         )
 
-        merged_run_tags = merge_dicts(self.run_tags, tags or {})
-        if partition_key:
-            ephemeral_job.validate_partition_key(
-                partition_key,
-                selected_asset_keys=asset_selection,
-                dynamic_partitions_store=instance,
-            )
-            tags_for_partition_key = ephemeral_job.get_tags_for_partition_key(
-                partition_key,
-                selected_asset_keys=asset_selection,
-            )
-
-            if not run_config and self.partitioned_config:
-                run_config = self.partitioned_config.get_run_config_for_partition_key(partition_key)
-
-            if self.partitioned_config:
-                merged_run_tags.update(
-                    self.partitioned_config.get_tags_for_partition_key(
-                        partition_key, job_name=self.name
-                    )
-                )
-            else:
-                merged_run_tags.update(tags_for_partition_key)
-
-        return core_execute_in_process(
-            ephemeral_job=ephemeral_job,
-            run_config=run_config,
-            instance=instance,
-            output_capturing_enabled=True,
-            raise_on_error=raise_on_error,
-            run_tags=merged_run_tags,
-            run_id=run_id,
-            asset_selection=frozenset(asset_selection),
-        )
-
     def _get_partitions_def(
         self, selected_asset_keys: Optional[Iterable[AssetKey]]
     ) -> PartitionsDefinition: