Handle @kubernetes/@batch specific attributes and local runs better

[tuulos]

Both @kubernetes and @batch have attributes that can't be set via @resources (e.g. tmpfs). To set these attributes, you need to use the decorator, but using the decorator forces the step to execute remotely.

There should be a way to define decorator-specific attributes while allowing local runs

[tuulos]

see this thread for an example https://outerbounds-community.slack.com/archives/C02116BBNTU/p1684234117421379

[tanzeyy]

Looks like that there is a config named ENABLE_STEP_DECORATOR trying to achive the goal but I haven't found any code associated with it: https://github.com/Netflix/metaflow/blob/master/metaflow/metaflow_config.py#L358

[roofurmston]

@tanzeyy I believe that is used as part of the metaflow extensions framework. IIRC you can use it to specify which decorators you want to keep as part of your Metaflow extension.

I think this is tangential to this issue, in which you would like to keep the decorators are part of your Metaflow install, but for some of them not to be "active" during local runs.

[roofurmston]

We have this issue. We are running flows remotely on Kubernetes/Argo Workflows.

We would like to specify certain pieces of configuration, such as the amount ephemeral storage, which is only available in the kubernetes decorator. However, it is also a requirement for us to be able to develop flows locally, which is not possible with the kubernetes decorator.

[roofurmston]

I've had a look through the codebase and, at least for the resources/kubernetes decorator situation, I have a rough proposal. I was thinking that the following could work:

• Add a KubernetesSpecification step decorator. This is used for specifying the Kubernetes configuration for a step, but doesn't try to run any code remotely. It would mainly consist of the defaults class variable and the __init__ method, both of which would match those of the kubernetes decorator.

• The step_init method of the kubernetes decorator would be updated so as to update its attribute values to match any specified in the kubernetes_specification. This would be analogous to how the kubernetes decorator currently updates its attributes using the resources decorator.

I believe the kubernetes_specification decorator would work in the same manner as the kubernetes decorator when it comes to Argo Workflows, Airflow & the --with kubernetes flag. The only difference would be that the vanilla run command would now run locally, with no steps running on Kubernetes.

On face value, it seems fairly easy to implement. I also believe this should be backward compatible.

A couple of open questions that would need to be resolved:

• How best to write the additional kubernetes_specification decorator so that there is no code duplication between this & the kubernetes decorator.

• Work out the logic of how to select which attribute value to use when it is specified in both decorators. Initial suggestion would be to use any non-default value and throw an error if both decorators specify a non-default value.

What do you think? Happy to flesh it out in a bit more detail, if you want. Or consider alternatives if you think this approach is no good.

[savingoyal]

One workaround that wouldn't require any changes to metaflow would be to create a pure python decorator whose behavior can be controlled using environment variables. Here is an approach - https://github.com/jacopotagliabue/you-dont-need-a-bigger-boat/blob/main/remote_flow/metaflow/custom_decorators.py#L29

[tanzeyy]

Another approach to disable decos like Kubernetes or batch when running locally could be adding a command line argument named --without or --disable. For example,

python flow.py run --disable=kubernetes

To achieve this, just modify the click option on start and run, and remove the decos passing through the cmd args. I'm happy to provide a demo if needed.

As for the Python decorator-based method mentioned by @savingoyal :

One workaround that wouldn't require any changes to metaflow would be to create a pure python decorator whose behavior can be controlled using environment variables. Here is an approach - https://github.com/jacopotagliabue/you-dont-need-a-bigger-boat/blob/main/remote_flow/metaflow/custom_decorators.py#L29

defacto such a method is already been used by us (we add a flow decorator named disable), but using a command line would be more elegant since we don't need to modify the source code.

[roofurmston]

@savingoyal It is a good point about using a no-op approach!

I originally had such an approach in mind for ourselves internally. We are using Metaflow extensions and it is easy for us to add this configuration in one place.

I tend to agree with @tanzeyy though. If we were using the native Metaflow I'd rather this functionality were part of the library.

Would also like to keep our use of extensions to a minimum. (Knowing it is not officially supported etc....) So in that respect too, it would be nice if it was part of the core library.

If it is of interest, here is how we set it up internally. We followed something very similar to your suggestion, but decorated the class itself!

from functools import wraps
from typing import Tuple

from metaflow.metaflow_config import (
    LOCAL_DEVELOPMENT
)

from metaflow.plugins.kubernetes.kubernetes_decorator import KubernetesDecorator as MetaflowKubernetesDecorator


def non_op_decorator(func):
    """Add non-operation decorator to the given function."""

    def no_op():
        pass

    @wraps(func)
    def wrapper(*args, **kwargs):
        if LOCAL_DEVELOPMENT:
            return no_op()
        return func(*args, **kwargs)
    return wrapper


def non_op_methods(cls, to_exclude: Tuple[str] = ('__init__', )):
    """
    Add no-op decorator to all methods, other than those specified in to_exclude, of the given class.

    param: cls: The class to which the no-op decorator is to be attached to member functions.
    param: to_exclude: A tuple of function names to which the decorator is not to be attached.
    """
    for k, v in vars(cls).items():
        if callable(v) and k not in to_exclude:
            setattr(cls, k, non_op_decorator(v))
    return cls


KubernetesDecorator = non_op_methods(MetaflowKubernetesDecorator)

Feel like it wouldn't be difficult to add something similar to the main library. Of course, in the main library one could simply decorate the class directly.

[elgehelge]

I think being able to easily switch between remote and local execution is extremely important for a great developer experience.

I really like the idea (that @tanzeyy came up with) of just having a flag like --with local or --without kubernetes 👍