dask flyte task plugin

gistfile1.py

class DaskFunctionTask(PythonFunctionTask[DaskConfig]):
    def pre_execute(self, user_params: ExecutionParameters):
       self.cluster = KubeCluster(...)
       self.client = Client(cluster)
        return (
            user_params.builder()
            .add_attr("DASK_CLIENT", self.client)
            .add_attr("DASK_CLUSTER", self.cluster)
            .build()
        )
        
    def execute(self, **kwargs):
        try:
            return super().execute(**kwargs)
        except Exception as e:
            self.cleanup()
            raise e
    
    def post_execute(self, user_params: ExecutionParameters, rval: Any):
        self.cleanup()
        return rval

    def cleanup(self):
        self.client.close()
        self.cluster.close()