[2/n][structured config] Enable struct config resources, IO managers to depend on other resources #11645

benpankow · 2023-01-11T20:00:16Z

Summary

Allows resources to depend on other resources:

class AWSCredentials(Resource):
    username: str
    password: str

class S3Resource(Resource):
    credentials: AWSCredentials
    region: str
    bucket: str

class EC2Resource(Resource):
    credentials: AWSCredentials

credentials = AWSCredentials(
    username=EnvVar("AWS_USERNAME"),
    password=EnvVar("AWS_PASSWORD"),
)

s3 = S3Resource(credentials=credentials)
ec2 = EC2Resource(credentials=credentials)

defs = Definitions(
    assets=[...],
    resources={
        "aws_credentials": credentials,
        "s3": s3,
        "ec2": snowflake,
    }
)

When using raw-Python outputs, these must be categorized as resources using the ResourceDependency annotation:

@resource(config_schema={"username": str, "password": str})
def postgres_engine(context):
    return sqlalchemy.create_engine(...)

class DBResource:
    engine: ResourceOutput[sqlalchemy.engine.Engine]

engine = postgres_engine.configured({"username": "foo", "password": "bar"})

# engine input is typed to take any of
# sqlalchemy.engine.Engine, Resource[sqlalchemy.engine.Engine],
# PartialResource[sqlalchemy.engine.Engine],
# or ResourceDefintion (for old, function-based resources)
db = DBResource(engine=engine)

Implementation

The goal of this PR is to allow for the following pattern:

foo = FooResource.configure_at_runtime()
bar = BarResource(inner=foo)

defs = Definitions(
    ...
    resources={
        "my_foo": foo,
        "my_bar": bar,
    }
)

Here, bar has a pointer to a ResourceDefinition for foo, but has no knowledge that foo is tied to the resource key "my_foo" (this is only specified at Definitions time). However, any pipeline which depends on "my_bar" as a required resource key must also depend on "my_foo", so bar.required_resource_keys must return {"my_foo"}.

Getting this information (that foo has "my_foo" as a resource key) to required_resource_keys is tricky.

This PR adds a new wrapper ResourceDefinition called ResourceWithKeyMapping which pairs the wrapped resource with mapping that takes ResourceDefinition IDs to resource keys (e.g. allowing us to map foo->"my_foo"). It catches any calls to required_resource_keys and instead forwards them to _resolve_required_resource_keys, which takes the mapping.

See #12232 for more context.

Test Plan

Added unit tests for I/O managers, resources.

Test Plan

vercel · 2023-01-11T20:00:20Z

The latest updates on your projects. Learn more about Vercel for Git ↗︎

2 Ignored Deployments

Name	Status	Preview	Comments	Updated
dagit-storybook	⬜️ Ignored (Inspect)			Feb 15, 2023 at 9:13PM (UTC)
dagster	⬜️ Ignored (Inspect)	Visit Preview		Feb 15, 2023 at 9:13PM (UTC)

benpankow · 2023-01-11T20:00:28Z

Current dependencies on/for this PR:

master
- PR [1/n][structured config] Add ability to runtime-configure struct-config resources #11773
  - PR [2/n][structured config] Enable struct config resources, IO managers to depend on other resources #11645 👈

This comment was auto-generated by Graphite.

python_modules/dagster/dagster/_config/structured_config.py

schrockn

I'm definitely scared of this one. It is fiendishly clever, but I'm concerned about maintainability over time.

What alternatives did we consider? Is there an approach available to use where we use a type annotation for resource-to-resource dependencies.

class MyResource:
   inner_resource: DependsOnResource[resource_key, InnerResource]

Stepping back I think this is type of situation where we should explore imposing more burden on users here in order to simplify the implementation and make the behavior more predictable.

schrockn · 2023-01-20T18:44:40Z

python_modules/dagster/dagster/_config/structured_config.py

+@dataclass_transform()
+class BaseResourceMeta(pydantic.main.ModelMetaclass):
+    def __new__(self, name, bases, namespaces, **kwargs):
+        annotations = namespaces.get("__annotations__", {})
+        for base in bases:
+            if hasattr(base, "__annotations__"):
+                annotations.update(base.__annotations__)
+        for field in annotations:
+            if not field.startswith("__"):
+                if get_origin(annotations[field]) == _ResourceDep:
+                    arg = get_args(annotations[field])[0]
+                    annotations[field] = Union[_PartialResource[arg], _Resource[arg]]
+                elif _safe_is_subclass(annotations[field], _Resource):
+                    base = annotations[field]
+                    annotations[field] = Union[_PartialResource[base], base]
+
+        namespaces["__annotations__"] = annotations
+        return super().__new__(self, name, bases, namespaces, **kwargs)


Can you explain, in prose, what this witchcraft is doing?

When a user defines a class extending Resource, Pydantic will use the type annotations to build a schema of valid values for each field.

For example, if the user defines the resources:

class MyCredentials(Resource): username: str password: str class MyResource(Resource): creds: MyCredentials

Then the schema requires creds to be an instance MyCredentials when MyResource is constructed. However, we would like to enable a user to specify a runtime-configured resource, which is of the type Partial[MyCredentials].

To do so, we want to redefine the user's type, behind the scenes, to:

class MyResource(Resource): creds: MyCredentials | PartialResource[MyCredentials]

This is what this metaclass does - when the user defines a class extending Resource, it iterates over all the fields and remaps any Resource-annotated field with a union of Resource and PartialResource.

(I've added some more comments to describe this inline)

python_modules/dagster/dagster_tests/core_tests/resource_tests/test_struct_config_resources.py

) ## Summary Alternate approach stacked on #11645 which removes the `top_level_key` shenanigans. The goal of this PR is to allow for the following pattern: ```python foo = FooResource.configure_at_runtime() bar = BarResource(inner=foo) defs = Definitions( ... resources={ "my_foo": foo, "my_bar": bar, } ) ``` Here, `bar` has a pointer to a `ResourceDefinition` for `foo`, but has no knowledge that `foo` is tied to the resource key `"my_foo"` (this is only specified at `Definitions` time). However, any pipeline which depends on `"my_bar"` as a required resource key must also depend on `"my_foo"`, so `bar.required_resource_keys` must return `{"my_foo"}`. Getting this information (that `foo` has `"my_foo"` as a resource key) to `required_resource_keys` is tricky. This PR adds a new wrapper `ResourceDefinition` called `ResourceWithKeyMapping` which pairs the wrapped resource with mapping that takes `ResourceDefinition` IDs to resource keys (e.g. allowing us to map `foo`->`"my_foo"`). It catches any calls to `required_resource_keys` and instead forwards them to `_resolve_required_resource_keys`, which takes the mapping. ## Test Plan Existing unit tests.

benpankow · 2023-02-10T19:04:39Z

Merged #12232 into this PR, avoiding the top_level_keys approach.

queue

schrockn

This looks mostly good, except for the string mangling and jamming dynamically named properties on the resources object. Let's do something better than that.

python_modules/dagster/dagster/_config/structured_config/__init__.py

schrockn · 2023-02-10T19:41:35Z

python_modules/dagster/dagster/_config/structured_config/__init__.py

+def _get_resource_by_id(context: InitResourceContext, resource_id: ResourceId) -> Any:
+    """
+    Retrieves a resource's value from the resource context object from the ResourceDefinition's ID.
+    This is used to retrieve the value of nested resources which are not fully configured, since the
+    resource key is not known by the nested resource's parent.
+    """
+    return getattr(context.resources, f"id_{resource_id}", None)


we should model this with a data structure somewhere, not mangle strings and monkeypatch

Attached to a new InitResourceContext subclass, InitResourceContextWithKeyMapping, which includes the id-to-resource mapping.

queue

schrockn

One last turn.

We definitely don't want to slow down the core tests case suite this much with those mypy and pyright tests.

python_modules/dagster/dagster/_config/structured_config/__init__.py

schrockn · 2023-02-11T10:15:22Z

python_modules/dagster/dagster_tests/core_tests/resource_tests/test_struct_config_resources.py

+def test_type_signatures_config_at_launch():
+    with tempfile.TemporaryDirectory() as tempdir:
+        filename = os.path.join(tempdir, "test.py")
+
+        with open(filename, "w") as f:
+            f.write(
+                """
+from dagster._config.structured_config import Resource
+
+class MyResource(Resource):
+    a_string: str
+
+reveal_type(MyResource.configure_at_launch())
+"""
+            )
+
+        pyright_out = get_pyright_reveal_type_output(filename)
+        mypy_out = get_mypy_type_output(filename)
+
+        # Ensure partial resource is correctly parameterized
+        assert pyright_out[0] == "PartialResource[MyResource]"
+        assert mypy_out[0].endswith("PartialResource[test.MyResource]")
+
+
+def test_type_signatures_constructor_resource_dependency():
+    with tempfile.TemporaryDirectory() as tempdir:
+        filename = os.path.join(tempdir, "test.py")
+
+        with open(filename, "w") as f:
+            f.write(
+                """
+from dagster._config.structured_config import Resource, ResourceDependency
+
+class StringDependentResource(Resource):
+    a_string: ResourceDependency[str]
+
+reveal_type(StringDependentResource.__init__)
+
+my_str_resource = StringDependentResource(a_string="foo")


These tests are supppper slow. I think we should disable them by default locally and additionally, run them in their own test suite to keep the core_test time reasonable.

Gated behind -m 'typesignature' mark. They'll run separately in CI now & only run locally with that flag.

schrockn · 2023-02-11T10:26:45Z

python_modules/dagster/dagster/_config/structured_config/typing_utils.py

+import pydantic
+from typing_extensions import dataclass_transform, get_origin
+
+from .utils import safe_is_subclass

 if TYPE_CHECKING:


Worthy of comment/explanation as to why this works while still having to do all the crazy shenanigans around temp classes.

schrockn · 2023-02-11T10:33:21Z

Also I think you should do something like #12273 to hide all the temp class stuff in a single scope. Feels a little java-y, but those top-level symbols aren't great.

queue

… depend on other resources

schrockn

ok lfg

python_modules/dagster/dagster/_config/structured_config/__init__.py

schrockn · 2023-02-14T01:38:06Z

python_modules/dagster/dagster/_config/structured_config/__init__.py

@@ -250,6 +506,23 @@ def resource_fn(context: InitResourceContext):


 ResourceOrPartial: TypeAlias = Union[Resource[ResValue], PartialResource[ResValue]]
+ResourceOrPartialOrBase: TypeAlias = Union[


or base? base what?

Changed to OrValue, since this can also hold the raw resource value (in the case of a bare Python object)

benpankow force-pushed the benpankow/wip-partial-resources branch from 563d10e to 3051156 Compare January 18, 2023 00:28

benpankow commented Jan 18, 2023

View reviewed changes

python_modules/dagster/dagster/_config/structured_config.py Outdated Show resolved Hide resolved

benpankow force-pushed the benpankow/wip-partial-resources branch from 3051156 to 6fbf9dc Compare January 18, 2023 18:46

benpankow changed the base branch from master to benpankow/wip-partial-resources-0.5 January 18, 2023 19:05

benpankow force-pushed the benpankow/wip-partial-resources branch from 6fbf9dc to 096796c Compare January 18, 2023 19:05

benpankow mentioned this pull request Jan 18, 2023

[1/n][structured config] Add ability to runtime-configure struct-config resources #11773

Merged

benpankow requested review from schrockn and smackesey January 18, 2023 19:11

benpankow changed the title ~~[structured config][wip] resource-resource deps w/ partial resources~~ [structured config] Enable struct config resources, IO managers to depend on other resources Jan 18, 2023

benpankow marked this pull request as ready for review January 18, 2023 19:13

benpankow force-pushed the benpankow/wip-partial-resources branch from 096796c to 6348623 Compare January 18, 2023 19:14

benpankow force-pushed the benpankow/wip-partial-resources-0.5 branch from 672d984 to 50e613c Compare January 18, 2023 23:20

benpankow force-pushed the benpankow/wip-partial-resources branch from 6348623 to a8fa3c1 Compare January 18, 2023 23:20

This was referenced Jan 18, 2023

[structured config] Migrate resources from project-fully-featured to struct config resources #11785

Merged

[structured config] Enable nesting of non-structured resources in structured resources #11791

Closed

benpankow force-pushed the benpankow/wip-partial-resources-0.5 branch from 50e613c to c2c6500 Compare January 20, 2023 18:04

benpankow force-pushed the benpankow/wip-partial-resources branch from 04934d1 to a95f683 Compare January 20, 2023 18:05

vercel bot deployed to Preview – dagster January 20, 2023 18:06 View deployment

vercel bot deployed to Preview – dagster January 20, 2023 18:28 View deployment

schrockn previously requested changes Jan 20, 2023

View reviewed changes

benpankow force-pushed the benpankow/wip-partial-resources branch from e13da59 to e8fac79 Compare January 20, 2023 23:07

vercel bot deployed to Preview – dagster January 20, 2023 23:09 View deployment

vercel bot deployed to Preview – dagster January 20, 2023 23:15 View deployment

benpankow force-pushed the benpankow/wip-partial-resources-0.5 branch from 9e595b2 to c8fcefe Compare January 24, 2023 00:59

benpankow force-pushed the benpankow/wip-partial-resources branch from d34ab46 to 886dfd9 Compare January 24, 2023 00:59

vercel bot deployed to Preview – dagster January 24, 2023 01:01 View deployment

benpankow force-pushed the benpankow/wip-partial-resources-0.5 branch from c8fcefe to ebd0178 Compare January 26, 2023 23:46

benpankow force-pushed the benpankow/wip-partial-resources branch from 1638875 to 1b4d8e3 Compare January 26, 2023 23:46

benpankow force-pushed the benpankow/wip-partial-resources branch from 0987feb to 74a648c Compare February 10, 2023 17:52

benpankow requested a review from schrockn February 10, 2023 19:04

schrockn previously requested changes Feb 10, 2023

View reviewed changes

benpankow requested a review from schrockn February 10, 2023 23:52

schrockn previously requested changes Feb 11, 2023

View reviewed changes

benpankow force-pushed the benpankow/wip-partial-resources branch from beaf683 to 46be9e3 Compare February 13, 2023 18:58

benpankow requested a review from schrockn February 13, 2023 19:00

benpankow force-pushed the benpankow/wip-partial-resources branch from 4c7dfda to 5fc053b Compare February 14, 2023 21:32

benpankow and others added 9 commits February 14, 2023 13:33

[structured config] Allow structured config resources, IO managers to…

32a1236

… depend on other resources

pyright

ca7a68a

pyright 2

9a43eeb

respond to review comments:

e722558

update err

9202ca7

code

e058212

respond to review comments

49cb1ad

no longer wrap normal resources with mapping

7f36b6d

run new tests in tox

45d2669

benpankow force-pushed the benpankow/wip-partial-resources branch from 5ef268e to 45d2669 Compare February 14, 2023 21:33

add mypy extra again

78a0baa

schrockn approved these changes Feb 15, 2023

View reviewed changes

respond to last set of review comments

e2d9c4d

benpankow merged commit 0d62593 into master Feb 15, 2023

benpankow deleted the benpankow/wip-partial-resources branch February 15, 2023 21:33

benpankow mentioned this pull request Feb 15, 2023

[pythonic config] Adjust public type signatures to take RunConfig #12375

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[2/n][structured config] Enable struct config resources, IO managers to depend on other resources #11645

[2/n][structured config] Enable struct config resources, IO managers to depend on other resources #11645

benpankow commented Jan 11, 2023 •

edited

Loading

vercel bot commented Jan 11, 2023 •

edited

Loading

benpankow commented Jan 11, 2023 •

edited

Loading

schrockn left a comment

schrockn Jan 20, 2023

benpankow Jan 20, 2023

benpankow commented Feb 10, 2023

schrockn left a comment

schrockn Feb 10, 2023

benpankow Feb 10, 2023

schrockn left a comment

schrockn Feb 11, 2023

benpankow Feb 13, 2023

schrockn Feb 11, 2023

benpankow Feb 13, 2023

schrockn commented Feb 11, 2023

schrockn left a comment

schrockn Feb 14, 2023

benpankow Feb 15, 2023

		@@ -250,6 +506,23 @@ def resource_fn(context: InitResourceContext):


		ResourceOrPartial: TypeAlias = Union[Resource[ResValue], PartialResource[ResValue]]
		ResourceOrPartialOrBase: TypeAlias = Union[

[2/n][structured config] Enable struct config resources, IO managers to depend on other resources #11645

[2/n][structured config] Enable struct config resources, IO managers to depend on other resources #11645

Conversation

benpankow commented Jan 11, 2023 • edited Loading

Summary

Implementation

Test Plan

Test Plan

vercel bot commented Jan 11, 2023 • edited Loading

benpankow commented Jan 11, 2023 • edited Loading

schrockn left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

benpankow commented Feb 10, 2023

schrockn left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

schrockn left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

schrockn commented Feb 11, 2023

schrockn left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

benpankow commented Jan 11, 2023 •

edited

Loading

vercel bot commented Jan 11, 2023 •

edited

Loading

benpankow commented Jan 11, 2023 •

edited

Loading