From 7de0f2c3b211140ba496b093b76221b15b877edf Mon Sep 17 00:00:00 2001 From: Mariam Soliman Date: Wed, 15 Jul 2026 06:37:07 +0300 Subject: [PATCH] Log which ops require each resource during initialization The RESOURCE_INIT_STARTED log line previously listed only resource keys: Starting initialization of resources [io_manager, resource_a, resource_b]. On jobs with many resources this made it impossible to tell which op pulled in a given resource. It now attributes each resource to the ops that require it: Starting initialization of resources: io_manager - required by the I/O layer resource_a - required by op instances [op_one] resource_b - required by op instances [op_two] get_resource_origins_mapping builds resource_key -> {op instance names} from op and hook required_resource_keys, keyed by full node handle so nested and aliased ops are unambiguous, and propagates through resource-to-resource dependencies. job_def is threaded to the resource init generator to supply it; when absent (e.g. build_resources) the original message is preserved. The dagstermill resource manager is updated to match the new signature. Resolves #2307. --- .../dagster/dagster/_core/events/__init__.py | 25 ++- .../_core/execution/context_creation_job.py | 1 + .../dagster/_core/execution/resources_init.py | 43 ++++- .../dagster_tests/core_tests/test_job_init.py | 168 ++++++++++++++++++ .../dagstermill/dagstermill/manager.py | 2 + 5 files changed, 235 insertions(+), 4 deletions(-) diff --git a/python_modules/dagster/dagster/_core/events/__init__.py b/python_modules/dagster/dagster/_core/events/__init__.py index 563a81c467d98..f94ccde771358 100644 --- a/python_modules/dagster/dagster/_core/events/__init__.py +++ b/python_modules/dagster/dagster/_core/events/__init__.py @@ -410,6 +410,26 @@ def log_resource_event(log_manager: DagsterLogManager, event: "DagsterEvent") -> log_manager.log_dagster_event(level=log_level, msg=event.message or "", dagster_event=event) +def _resource_init_start_message( + resource_keys: AbstractSet[str], + resource_key_to_op_names: Mapping[str, AbstractSet[str]] | None, +) -> str: + if not resource_key_to_op_names: + return "Starting initialization of resources [{}].".format(", ".join(sorted(resource_keys))) + + # Iterate resource_keys (scoped to this step) rather than the job-wide mapping, so a step + # only lists its own ops. Keys not in the mapping are required by an io/input manager or + # type loader rather than by an op directly. + lines = ["Starting initialization of resources:"] + for resource_key in sorted(resource_keys): + op_names = sorted(resource_key_to_op_names.get(resource_key, set())) + if op_names: + lines.append(f"{resource_key} - required by op instances [{', '.join(op_names)}]") + else: + lines.append(f"{resource_key} - required by the I/O layer") + return "\n".join(lines) + + class DagsterEventSerializer(NamedTupleSerializer["DagsterEvent"]): def before_unpack(self, context, unpacked_dict: Any) -> dict[str, Any]: event_type_value, event_specific_data = _handle_back_compat( @@ -1337,15 +1357,14 @@ def resource_init_start( execution_plan: "ExecutionPlan", log_manager: DagsterLogManager, resource_keys: AbstractSet[str], + resource_key_to_op_names: Mapping[str, AbstractSet[str]] | None = None, ) -> "DagsterEvent": return DagsterEvent.from_resource( DagsterEventType.RESOURCE_INIT_STARTED, job_name=job_name, execution_plan=execution_plan, log_manager=log_manager, - message="Starting initialization of resources [{}].".format( - ", ".join(sorted(resource_keys)) - ), + message=_resource_init_start_message(resource_keys, resource_key_to_op_names), event_specific_data=EngineEventData(metadata={}, marker_start="resources"), ) diff --git a/python_modules/dagster/dagster/_core/execution/context_creation_job.py b/python_modules/dagster/dagster/_core/execution/context_creation_job.py index 900d3bcaf402b..ee1335dbd02de 100644 --- a/python_modules/dagster/dagster/_core/execution/context_creation_job.py +++ b/python_modules/dagster/dagster/_core/execution/context_creation_job.py @@ -238,6 +238,7 @@ def execution_context_event_generator( instance=instance, emit_persistent_events=True, event_loop=event_loop, + job_def=job_def, ) yield from resources_manager.generate_setup_events() scoped_resources_builder = check.inst( diff --git a/python_modules/dagster/dagster/_core/execution/resources_init.py b/python_modules/dagster/dagster/_core/execution/resources_init.py index 309c9f2c5e060..23c06c3209076 100644 --- a/python_modules/dagster/dagster/_core/execution/resources_init.py +++ b/python_modules/dagster/dagster/_core/execution/resources_init.py @@ -1,6 +1,6 @@ import inspect from asyncio import AbstractEventLoop -from collections import deque +from collections import defaultdict, deque from collections.abc import Generator, Mapping from contextlib import ContextDecorator from typing import AbstractSet, Any, Callable, cast # noqa: UP035 @@ -8,6 +8,7 @@ from dagster_shared.utils.timing import format_duration import dagster._check as check +from dagster._core.definitions.dependency import OpNode from dagster._core.definitions.job_definition import JobDefinition from dagster._core.definitions.resource_definition import ( ResourceDefinition, @@ -49,6 +50,7 @@ def resource_initialization_manager( instance: DagsterInstance | None, emit_persistent_events: bool | None, event_loop: AbstractEventLoop | None, + job_def: JobDefinition | None = None, ): generator = resource_initialization_event_generator( resource_defs=resource_defs, @@ -60,6 +62,7 @@ def resource_initialization_manager( instance=instance, emit_persistent_events=emit_persistent_events, event_loop=event_loop, + job_def=job_def, ) return EventGenerationManager(generator, ScopedResourcesBuilder) @@ -116,6 +119,39 @@ def _get_deps_helper(resource_key): return reqd_resources +def get_resource_origins_mapping(job_def: JobDefinition) -> Mapping[str, AbstractSet[str]]: + """Maps each resource key to the op instances that require it, keyed by full node handle so + that aliased and nested ops are unambiguous. + + A resource required by another resource (e.g. spark, pulled in by a pyspark resource) is + attributed to the ops that required the outer resource. Resources no op requires directly, + such as the default io_manager, are absent from the mapping. + """ + origins: dict[str, set[str]] = defaultdict(set) + + for handle in job_def.graph.iterate_node_handles(): + node = job_def.get_node(handle) + if not isinstance(node, OpNode): + continue + + op_name = str(handle) + + for resource_key in node.definition.required_resource_keys: + origins[resource_key].add(op_name) + + for hook_def in job_def.get_all_hooks_for_handle(handle): + for resource_key in hook_def.required_resource_keys: + origins[resource_key].add(op_name) + + resource_dependencies = resolve_resource_dependencies(job_def.resource_defs) + for resource_key, op_names in list(origins.items()): + for dep_key in get_dependencies(resource_key, resource_dependencies): + if dep_key != resource_key: + origins[dep_key].update(op_names) + + return dict(origins) + + def _core_resource_initialization_event_generator( resource_defs: Mapping[str, ResourceDefinition], resource_configs: Mapping[str, ResourceConfig], @@ -127,6 +163,7 @@ def _core_resource_initialization_event_generator( instance: DagsterInstance | None, emit_persistent_events: bool | None, event_loop, + job_def: JobDefinition | None = None, ): job_name = "" # Must be initialized to a string to satisfy typechecker contains_generator = False @@ -147,6 +184,7 @@ def _core_resource_initialization_event_generator( cast("ExecutionPlan", execution_plan), resource_log_manager, resource_keys_to_init, + get_resource_origins_mapping(job_def) if job_def else None, ) resource_dependencies = resolve_resource_dependencies(resource_defs) @@ -225,6 +263,7 @@ def resource_initialization_event_generator( instance: DagsterInstance | None, emit_persistent_events: bool | None, event_loop: AbstractEventLoop | None, + job_def: JobDefinition | None = None, ): check.inst_param(log_manager, "log_manager", DagsterLogManager) resource_keys_to_init = check.opt_set_param( @@ -233,6 +272,7 @@ def resource_initialization_event_generator( check.opt_inst_param(execution_plan, "execution_plan", ExecutionPlan) check.opt_inst_param(dagster_run, "dagster_run", DagsterRun) check.opt_inst_param(instance, "instance", DagsterInstance) + check.opt_inst_param(job_def, "job_def", JobDefinition) if execution_plan and execution_plan.step_handle_for_single_step_plans(): step = execution_plan.get_step( @@ -260,6 +300,7 @@ def resource_initialization_event_generator( instance=instance, emit_persistent_events=emit_persistent_events, event_loop=event_loop, + job_def=job_def, ) except GeneratorExit: # Shouldn't happen, but avoid runtime-exception in case this generator gets GC-ed diff --git a/python_modules/dagster/dagster_tests/core_tests/test_job_init.py b/python_modules/dagster/dagster_tests/core_tests/test_job_init.py index e9182a9012ba6..4682ad2bf9592 100644 --- a/python_modules/dagster/dagster_tests/core_tests/test_job_init.py +++ b/python_modules/dagster/dagster_tests/core_tests/test_job_init.py @@ -2,9 +2,13 @@ import pytest from dagster import DagsterInstance from dagster._core.definitions.job_base import InMemoryJob +from dagster._core.events import DagsterEvent, DagsterEventType +from dagster._core.execution import resources_init from dagster._core.execution.api import create_execution_plan from dagster._core.execution.context_creation_job import PlanExecutionContextManager from dagster._core.execution.resources_init import ( + get_required_resource_keys_to_init, + get_resource_origins_mapping, resource_initialization_event_generator, resource_initialization_manager, single_resource_event_generator, @@ -117,3 +121,167 @@ def test_clean_event_generator_exit(): @dg.op def fake_op(_): pass + + +def test_get_resource_origins_mapping(): + @dg.resource + def spark(_): + return "spark" + + @dg.resource(required_resource_keys={"spark"}) + def pyspark(_): + return "pyspark" + + @dg.resource + def required_by_nobody(_): + return "nobody" + + @dg.op(required_resource_keys={"pyspark"}) + def uses_pyspark(_): + pass + + @dg.op(required_resource_keys={"spark"}) + def uses_spark_directly(_): + pass + + @dg.graph + def nested(): + uses_spark_directly() + + @dg.job( + resource_defs={ + "spark": spark, + "pyspark": pyspark, + "required_by_nobody": required_by_nobody, + } + ) + def origins_job(): + uses_pyspark() + nested() + + mapping = get_resource_origins_mapping(origins_job) + + assert mapping["pyspark"] == {"uses_pyspark"} + + # spark is required transitively (via pyspark) and directly by a nested op, reported by handle + assert mapping["spark"] == {"uses_pyspark", "nested.uses_spark_directly"} + + # resources no op requires are absent from the mapping + assert "required_by_nobody" not in mapping + assert "io_manager" not in mapping + + +def gen_two_op_resource_job(): + """A job where each op requires a different resource.""" + + @dg.resource + def resource_a(_): + return "a" + + @dg.resource + def resource_b(_): + return "b" + + @dg.op(required_resource_keys={"resource_a"}) + def op_one(_): + pass + + @dg.op(required_resource_keys={"resource_b"}) + def op_two(_): + pass + + @dg.job(resource_defs={"resource_a": resource_a, "resource_b": resource_b}) + def two_op_resource_job(): + op_one() + op_two() + + return two_op_resource_job + + +def _resource_init_started_message(job_def, *, pass_job_def: bool) -> str: + """Drives the real resource-init generator and returns the RESOURCE_INIT_STARTED message.""" + instance = DagsterInstance.ephemeral() + execution_plan = create_execution_plan(job_def) + run = instance.create_run_for_job(job_def=job_def, execution_plan=execution_plan) + log_manager = DagsterLogManager.create(loggers=[], dagster_run=run) + resolved_run_config = ResolvedRunConfig.build(job_def) + + events = list( + resource_initialization_event_generator( + resource_defs=job_def.resource_defs, + resource_configs=resolved_run_config.resources, + log_manager=log_manager, + execution_plan=execution_plan, + dagster_run=run, + resource_keys_to_init=get_required_resource_keys_to_init(execution_plan, job_def), + instance=instance, + emit_persistent_events=True, + event_loop=None, + job_def=job_def if pass_job_def else None, + ) + ) + + started = [ + event + for event in events + if isinstance(event, DagsterEvent) + and event.event_type_value == DagsterEventType.RESOURCE_INIT_STARTED.value + ] + assert len(started) == 1, f"expected exactly one RESOURCE_INIT_STARTED, got {len(started)}" + return started[0].message + + +def test_resource_init_message_reports_op_origins(): + """With job_def available, the message says which ops require each resource (issue #2307).""" + message = _resource_init_started_message(gen_two_op_resource_job(), pass_job_def=True) + + assert message == ( + "Starting initialization of resources:\n" + "io_manager - required by the I/O layer\n" + "resource_a - required by op instances [op_one]\n" + "resource_b - required by op instances [op_two]" + ) + + +def test_resource_init_message_falls_back_without_job_def(): + """Without job_def (e.g. build_resources), the message is byte-for-byte the original.""" + message = _resource_init_started_message(gen_two_op_resource_job(), pass_job_def=False) + + assert message == "Starting initialization of resources [io_manager, resource_a, resource_b]." + + +def test_job_def_reaches_innermost_resource_generator(monkeypatch): + """job_def must be forwarded from resource_initialization_manager down to the innermost + generator. A param that is accepted but not forwarded still typechecks, so assert it arrives. + """ + job_def = gen_two_op_resource_job() + instance = DagsterInstance.ephemeral() + execution_plan = create_execution_plan(job_def) + run = instance.create_run_for_job(job_def=job_def, execution_plan=execution_plan) + log_manager = DagsterLogManager.create(loggers=[], dagster_run=run) + resolved_run_config = ResolvedRunConfig.build(job_def) + + captured = {} + real_core_generator = resources_init._core_resource_initialization_event_generator # noqa: SLF001 + + def spy(**kwargs): + captured["job_def"] = kwargs.get("job_def", "<>") + return real_core_generator(**kwargs) + + monkeypatch.setattr(resources_init, "_core_resource_initialization_event_generator", spy) + + manager = resource_initialization_manager( + resource_defs=job_def.resource_defs, + resource_configs=resolved_run_config.resources, + log_manager=log_manager, + execution_plan=execution_plan, + dagster_run=run, + resource_keys_to_init=get_required_resource_keys_to_init(execution_plan, job_def), + instance=instance, + emit_persistent_events=True, + event_loop=None, + job_def=job_def, + ) + list(manager.generate_setup_events()) + + assert captured["job_def"] is job_def diff --git a/python_modules/libraries/dagstermill/dagstermill/manager.py b/python_modules/libraries/dagstermill/dagstermill/manager.py index 188582d285599..e4257f864b5e4 100644 --- a/python_modules/libraries/dagstermill/dagstermill/manager.py +++ b/python_modules/libraries/dagstermill/dagstermill/manager.py @@ -86,6 +86,7 @@ def _setup_resources( instance: DagsterInstance | None, emit_persistent_events: bool | None, event_loop: AbstractEventLoop | None, + job_def: JobDefinition | None = None, ): """Drop-in replacement for `dagster._core.execution.resources_init.resource_initialization_manager`. It uses a @@ -101,6 +102,7 @@ def _setup_resources( instance=instance, emit_persistent_events=emit_persistent_events, event_loop=event_loop, + job_def=job_def, ) self.resource_manager = DagstermillResourceEventGenerationManager( generator, ScopedResourcesBuilder