Made some progress here if anyone is interested: https://prefect-community.slack.com/archives/CL09KU1K7/p1691527886984439 basically I had to split out the data context from the class and into its own dataclass with data related methods , this helped clean it up alot and made more idiomatic