[TRIS-722] Custom Tagging

[RikishK]

Options to add custom tags are in the order:

environment variable dictionary,
step functions cli create flag,
batch decorator variable dictionary

This PR builds ontop of this PR: #1628 - this original PR did not work when testing but I would like to credit it here as some of the changes are still used.

[RikishK]

Testing

Setup:

Environment variable:

{ name: 'METAFLOW_BATCH_EMIT_TAGS', value: 'True' },
{ name: 'METAFLOW_BATCH_DEFAULT_TAGS', value: '{"default_greeting": "default_world"}' }

Step functions create command:

step-functions create --aws-tags "hello=world"

Flow:

from metaflow import (
    FlowSpec,
    step,
    batch,
    retry,
    schedule,
    project,
    conda,
    current
)
import logging
import logging_setup

logging_setup.configure()


custom_step_tags = {
    "goodbye": "world",
    "hello": "universe",
   #"inv@lid": "t@g"
    }

@project(name="ml_apac_trisolaris")
@schedule(hourly=True)
class CanarySimpleAdhoc(FlowSpec):
    @step
    def start(self):
        self.logger = logging.getLogger(self.__class__.__name__)
        self.logger.info(f"Canary Hello")
        self.next(self.hello)

    @batch(cpu=1, memory=500, tags=custom_step_tags)
    @retry
    @step
    def hello(self):
        self.logger.info("Canary Hello World in Batch!")
        self.next(self.end)

    @step
    def end(self):
        self.logger.info("HelloAWS is finished.")


if __name__ == "__main__":
    CanarySimpleAdhoc()

Output/Result:

Start step batch tags:

Hello step batch tags:

Additional notes:
Tag validation was tested with "inv@lid": "t@g" which caused an error as expected.

[saikonen]

notes for testing, the config variable

METAFLOW_BATCH_EMIT_TAGS=True 

needs to be set in order for tags to have any effect (even validation).

For a better UX, there could be a check somewhere in the path that if aws tags are being set, but emit tags is set to false, we print a warning on this.

Validation could also maybe be run despite the emit_tags feature flag, in order to catch issues.

[saikonen]

we need to add similar logic that is present in f.ex. the kubernetes_decorator.py#runtime_step_cli which json.dumps some attributes to the cli_args.commands, namely the new aws_tags in this case.

otherwise the passed in value will be not handled correctly in the CLI

[saikonen]

for example replacing the cli_args.commands.update(self.attributes) with the following should work:

# dictionaries need to be json dumped for the CLI
            json_types = ["aws_tags"]
            for k, v in self.attributes.items():
                if k in json_types:
                    cli_args.command[k] = json.dumps(v)
                else:
                    cli_args.command_options[k] = v

[saikonen]

also for the batch_cli.py side, the --aws_tags option needs to be

@click.option("--aws-tags", multiple=False, type=JSONTypeClass(), default=None, help="AWS tags.")

in order to convert types correctly. note the multiple=False also, was there a reason to have the plural option act as multiple?

[saikonen]

I didn't test this part thoroughly, but there is one significant caveat with running validations in the batch setup phase, namely: when validation fails, the metadata for the run/step/task have already been registered in this case, leading to a failed task with no logs.

An alternative approach with a better UX would be to perform batch tag validation inside the batch_decorator.py#step_init, which will happen before job submission. This allows for erroring out fairly quickly due to malformed tags.

Is there any issue with moving the validation and the BATCH_DEFAULT_TAGS usage into the decorator completely?

[saikonen]

also can you verify after the changes that the following flows run as expected? both with batch, and step-functions

expected to fail

from metaflow import step, FlowSpec, batch

class BatchTags(FlowSpec):
    @batch(aws_tags={"inv@lid": "t@g"})
    @step
    def start(self):
        print("Starting 👋")
        self.next(self.end)

    @step
    def end(self):
        print("Done! 🏁")


if __name__ == "__main__":
    BatchTags()

expected to run

from metaflow import (
    FlowSpec,
    step,
    batch,
)

custom_step_tags = {
    "goodbye": "world",
    "hello": "universe",
    }

class BatchTags2(FlowSpec):
    @step
    def start(self):
        self.next(self.hello)

    @batch(aws_tags=custom_step_tags)
    @step
    def hello(self):
        self.next(self.end)

    @step
    def end(self):
        pass

if __name__ == "__main__":
    BatchTags2()