Extend flytekit version hash calculation to be pluggable

[ddl-rliu]

Tracking issue

Closes flyteorg/flyte#5364

Why are the changes needed?

Extends https://github.com/flyteorg/flytekit/pull/2039/files – that PR gives the minimum API surface for configuring the API via FlytekitPlugin.

This PR increases how pyflyte is pluggable by external libraries, specifically for controlling the additional context used to generate the version string hash.

What changes were proposed in this pull request?

This can be used with a plugin like:

class MyPlugin(FlytekitPlugin):
    @staticmethod
    def get_additional_context(entity: Union[PythonAutoContainerTask, WorkflowBase]) -> List[str]:
        """Get additional context to be used for calculating the version hash."""
        if isinstance(entity, PythonTask):
            return [str(entity.task_config)]
        if isinstance(entity, WorkflowBase):
            task_configs = []
            for n in entity.nodes:
                task_configs.extend(DominoJobPlugin.get_additional_context(n.flyte_entity))
            return task_configs
        return []

This will add the task config to the calculation of the version hash.

How was this patch tested?

Setup process

Screenshots

Check all the applicable boxes

• I updated the documentation accordingly.
• All new and existing tests passed.
• All commits are signed-off.

Related PRs

Docs link

[codecov]

Codecov Report

Attention: Patch coverage is 0% with 2 lines in your changes are missing coverage. Please review.

Project coverage is 72.09%. Comparing base (e6e08f9) to head (a744848).
Report is 15 commits behind head on master.

❗ Current head a744848 differs from pull request most recent head d4cda94

Please upload reports for the commit d4cda94 to get more accurate results.

Files	Patch %	Lines
flytekit/remote/remote.py	0.00%	2 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##           master    #2428       +/-   ##
===========================================
+ Coverage   42.77%   72.09%   +29.31%     
===========================================
  Files         185      181        -4     
  Lines       18677    18397      -280     
  Branches     2665     3601      +936     
===========================================
+ Hits         7990    13264     +5274     
+ Misses      10599     4508     -6091     
- Partials       88      625      +537     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[ddl-ebrown]

Probably not the ideal place to have an import.

It also seems like the additional_context (which could probably be named a bit more descriptively?) value should be passed into the register_script function, shouldn't it? The pattern of pulling in special state like this makes this code a bit harder to test / reason about I think.

[ddl-rliu]

I use lazy import here to solve the circular import issue (remote imports plugin, plugin imports remote)

[ddl-ebrown]

Yeah, I've used this tactic to reduce startup time in Ruby -- but:

• It looks like there's already a lazy_module helper in Flytekit
• if this is a tactic for dealing with a potential circular dependency, that can be the sign of a flawed design
• Remote now takes a dependency on this plugin helper method -- it's cleaner if you can invert the behavior so that this function doesn't need to know anything about finding special mutation functions in plugins

[ddl-rliu]

I could rename additional_context to version_hash_additional_context? I wouldn't rename it to task_configs because it does not necessarily contain the task configs, that only happens for a particular use case.

Moving it to a parameter of register_script seems sensible, but for the moment I'm thinking of keeping the PR as very easy to review/not changing too many methods. Will give it a try at a later point.

edit: will rename to version_hash_additional_context

[ddl-ebrown]

Would it make more sense to have pluggable versioning instead? i.e. the plugin defines a custom version function that gets passed md5_bytes, serialization_settings, default_inputs, *_get_image_names(entity) if it exists? Would such a function need a dual?

[ddl-rliu]

My feeling is that core logic i.e. the default version hash logic, should remain in its section, rather than moved to the plugin methods. This mostly follows the existing pattern I noticed in FlytekitPlugin, where core logic is not often moved to its methods – however there are some exceptions, as I suppose with FlytekitPlugin.get_remote

[ddl-ebrown]

@pingsutw this is a similar "extension" as #2661

Should we try to get these 2 PRs lined up for the new release of flytekit @ddl-rliu ?

[thomasjpfan]

I prefer not to have a new extension point in FlytekitPluginProtocol, if there is an simpler alternative.

For your use case, can your external library provide a FlyteRemote subclass and override _version_from_hash? This FlyteRemote subclass can then be returned by your custom FlytekitPluginProtocol.get_remote.

[ddl-rliu]

@thomasjpfan Thanks, it's an interesting idea, and I never thought about that possibility! With a modest refactor of _version_from_hash, particularly to accept entity as a parameter, I think it will work exactly like you say. Proof of concept refactor: https://github.com/flyteorg/flytekit/pull/2688/files

To explain more – the parent issue deals with adding entity.task_config to the version hash calculation, therefore, the refactor would expose entity via a parameter to _version_from_hash.