pin s3 worker count for kubernetes tasks

[savingoyal]

No description provided.

[savingoyal]

need to set the right vars in argo workflows and airflow too

[savingoyal]

would we need a similar change in argo workflows for pods and jobsets?

[savingoyal]

would we need a change in step functions as well?

[saikonen]

step-functions calls the batch implementation so this should apply there as well. will verify

[saikonen]

confirmed, step-functions is setting the env correctly through the changes in Batch

[saikonen]

noting something extra that needs to be handled here:

from metaflow import step, FlowSpec, resources
import os

class S3WorkerFlow(FlowSpec):
    @resources(cpu=1)
    @step
    def start(self):
        val = int(os.environ.get("METAFLOW_S3_WORKER_COUNT"))
        assert val == 1, "Worker count should be 1!"
        self.next(self.bigger)

    @resources(cpu=4)
    @step
    def bigger(self):
        val = int(os.environ.get("METAFLOW_S3_WORKER_COUNT"))
        assert val == 2, "Worker count should be 2!"
        self.next(self.end)

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


if __name__ == "__main__":
    S3WorkerFlow()

running this --with kubernetes or argo workflows fails, as resource decorator values are cast to a str-float, which fails to convert back to an int

[saikonen]

running this --with kubernetes or argo workflows fails, as resource decorator values are cast to a str-float, which fails to convert back to an int

The cause for this is the different defaults in @batch compared to @kubernetes for the cpu value, which leads to a string-float cpu value with resources&kubernetes, but a string-int cpu value for resources&batch