[6/n subset refactor] Serialize AssetGraphSubset and AssetBackfillData with whitelist_for_serdes

[clairelin135]

This PR makes AssetGraphSubset and AssetBackfillData serializable via whitelist_for_serdes.

This involves the following changes:

• converts AssetGraphSubset and AssetBackfillData into NamedTuples
  • adds a custom serializer to convert PartitionKeysTimeWindowPartitionsSubset to TimeWindowsPartitionsSubset
• removes the asset graph property from AssetGraphSubset as it is not serializable. This has a cascading effect:
  • callsites where we previously called asset_graph_subset.asset_graph must instead have an asset graph passed in
  • previously we could build empty partitions subsets when needed within AssetGraphSubset (i.e. within __or__). This logic now must be updated to handle cases where a partitions subset is currently None
  • now AssetGraphSubset or, and, and sub (|, &, -) operations cannot operate directly against sets of AssetKeyPartitionKeys, since the asset graph is required to build subsets from these AssetKeyPartitionKeys

[clairelin135]

Current dependencies on/for this PR:

• master
  • PR [1/n subset refactor] Split TimeWindowPartitionsSubset #17684
    • PR [2/n subset refactor] Make TimeWindowPartitionsDefinition serializable #17660
      • PR [3/n subset refactor] Add whitelist_for_serdes to TimeWindowPartitionsSubset #17702
        • PR [4/n subset refactor] Add whitelist_for_serdes to DefaultPartitionsSubset #17703
          • PR [5/n subset refactor] [serdes] Enable serializing mappings with non-scalar keys #18057
            • PR [6/n subset refactor] Serialize AssetGraphSubset and AssetBackfillData with whitelist_for_serdes #17844 👈
              • PR [7/n subset refactor] Use new asset backfill data serialization format #17929

This stack of pull requests is managed by Graphite.

[clairelin135]

Unfortunately json cannot serialize dictionaries that aren't keyed by a primitive or a string, so we key by the serialized asset key instead

[sryza]

Could it make sense to handle that in the custom serializer so it doesn't need to leak into the data model for the class?

[sryza]

Also, don't want to increase scope too much, but it could make sense to try to address this at the serialization layer if it's not too big a lift.

[clairelin135]

Also, don't want to increase scope too much, but it could make sense to try to address this at the serialization layer if it's not too big a lift.

I did take a stab at implementing it this way.

• The properly serialized asset key pack_value(asset_key....) value is a dictionary, which cannot be used as a key to a dictionary
• We could add custom logic in the serialization layer to convert an asset key to a string (i.e. asset_key.to_user_string() but this seems like a pain to deal with if we did decide to add additional fields to the AssetKey class

[clairelin135]

Could it make sense to handle that in the custom serializer so it doesn't need to leak into the data model for the class?

My initial version actually used this approach, it looked like this:

dagster/python_modules/dagster/dagster/_core/definitions/asset_graph_subset.py

Lines 42 to 69 in 79652f9

	class PartitionsSubsetByAssetKeySerializer(FieldSerializer):
	"""Packs and unpacks a mapping from AssetKey to PartitionsSubset.
	In JSON, a key must be a str, int, float, bool, or None. This serializer packs the AssetKey
	into a str, and unpacks it back into an AssetKey.
	It also converts PartitionKeysTimeWindowPartitionsSubset into serializable TimeWindowPartitionsSubsets.
	"""
	def pack(self, mapping: Mapping[AssetKey, Any], **_kwargs) -> Mapping[str, Any]:
	return {
	serialize_value(key): serialize_value(
	value.to_time_window_partitions_subset()
	if isinstance(value, PartitionKeysTimeWindowPartitionsSubset)
	else value
	)
	for key, value in mapping.items()
	}
	def unpack(
	self,
	mapping: Mapping[str, Any],
	**_kwargs,
	) -> Mapping[AssetKey, Any]:
	return {
	deserialize_value(key, AssetKey): deserialize_value(value, TimeWindowPartitionsSubset)
	for key, value in mapping.items()
	}

I moved away from this implementation because it felt like a duplicate of the existing logic, with the exception of converting the asset key to a string, so maybe it would be better to reduce code surface area by just converting the asset keys to serialized form.

I feel mixed on this though -- I can see either making sense. It certainly is cleaner to not have to convert to/from serialized asset keys.

[sryza]

but this seems like a pain to deal with if we did decide to add additional fields to the AssetKey class

I think that we can be confident that we are not going to do this. I suspect a lot of other places would break as well.

[clairelin135]

passed in asset graph because we'd like to get an empty subset instead of None if possible

the __oper__ callsite does not have access to the asset graph, so this is an optional param

[sryza]

Would it make sense to put the custom serializer on PartitionKeysTimeWindowPartitionsSubset itself?

[clairelin135]

It would be nice to not have to convert these subset objects everywhere they're used, but I think on principle objects that don't cross a serialization boundary should not be decorated with @whitelist_for_serdes.

Ideally we have some parallel entity that means "convert to Y, then serialize/deserialize Y". Implementing something like that adds additional scope I'd prefer to avoid at this time

[sryza]

Could it make sense to handle that in the custom serializer so it doesn't need to leak into the data model for the class?

[clairelin135]

@sryza This PR has been updated to directly serialize partitions_subsets_by_asset_key following the changes in part 5 that enable this in the serdes layer

[clairelin135]

@sryza I've refactored this serializer to be generalizable to other named tuples that contain partitions subsets mappings. Unfortunately serializers are instantiated on demand rather than at definition time, so we can't easily provide field names to apply this custom logic.

Instead, I've added logic to detect when partitions subsets exist and to convert those accordingly

[sryza]

That makes sense

[sryza]

This could use short docstring.

[sryza]

A few small remaining comments. Otherwise LGTM!

[sryza]

Could this just be

@whitelist_for_serdes(...)
class AssetGraphSubset(NamedTuple):
    partitions_subsets_by_asset_key: Mapping[AssetKey, PartitionsSubset]
   non_partitioned_asset_keys: AbstractSet[AssetKey]

with no __new__?

[sryza]

That makes sense

[sryza]

This could use short docstring.