Package: mlflowhelper
mlflowhelperA set of tools for working withmlflow(seehttps://mlflow.org)Featuresmanaged artifact logging andloadingautomaticartifact logging and cleanupno overwriting fileswhen running scripts in parallelloadingartifactcentral configurationof logging and loading behaviorlogallfunction parameters and locals with a simple call tomlflowhelper.log_vars()Documentationpip install mlflowhelperManaged artifact logging and loadingGeneral functionalityfrommatplotlibimportpyplotaspltimportmlflowhelperwithmlflowhelper.start_run():withmlflowhelper.managed_artifact("plot.png")asartifact:fig=plt.figure()plt.plot([1,2,3],[1,2,3])fig.savefig(artifact.get_path())This code snippet automatically logs the created artifact (plot.png).
At the same time if will create the artifact in a temporary folder so that you don't have to worry about
overwriting it when running your scripts in parallel.
By default, this also cleans up the artifact and the temporary folder after logging.You can also manage artifacts on a directory level:frommatplotlibimportpyplotaspltimportmlflowhelperwithmlflowhelper.start_run():withmlflowhelper.managed_artifact_dir("plots")asartifact_dir:# plot 1fig=plt.figure()plt.plot([1,2,3],[1,2,3])fig.savefig(artifact_dir.get_path("plot1.png"))# plot 2fig=plt.figure()plt.plot([1,2,3],[1,2,3])fig.savefig(artifact_dir.get_path("plot2.png"))Artifact loadingYou may want to run experiments but reuse some precomputed artifact from a different run (such
as preprocessed data, trained models, etc.). This can be done as follows:importmlflowhelperimportpandasaspdwithmlflowhelper.start_run():mlflowhelper.set_load(run_id="e1363f760b1e4ab3a9e93f856f2e9341",stages=["load_data"])# activate loading from previous runwithmlflowhelper.managed_artifact_dir("data.csv",stage="load_data")asartifact:ifartifact.loaded:# load artifactdata=pd.read_csv(artifact.get_path())else:# create and save artifactdata=pd.read_csv("/shared/dir/data.csv").sample(frac=1)data.to_csv(artifact.get_path())Similarly, this works for directories of course:importmlflowhelperimportpandasaspdmlflowhelper.set_load(run_id="e1363f760b1e4ab3a9e93f856f2e9341",stages=["load_data"])# activate loading from previous runwithmlflowhelper.start_run():withmlflowhelper.managed_artifact_dir("data",stage="load_data")asartifact_dir:train_path=artifact_dir.get_path("test.csv")test_path=artifact_dir.get_path("train.csv")ifartifact_dir.loaded:# load artifactstrain=pd.read_csv(train_path)test=pd.read_csv(test_path)else:data=pd.read_csv("/shared/dir/data.csv").sample(frac=1)train=data.iloc[:100,:]test=data.iloc[100:,:]# save artifactstrain.to_csv(train_path)test.to_csv(test_path)Note:Thestageparameter must be set inmlflowhelper.managed_artifact(_dir)to enable loading.Central logging and loading behavior managementLogging and loading behavior can be managed in a central way:importmlflowhelperimportpandasaspdwithmlflowhelper.start_run():# activate loading the stage `load_data` from previous run `e1363f760b1e4ab3a9e93f856f2e9341`mlflowhelper.set_load(run_id="e1363f760b1e4ab3a9e93f856f2e9341",stages=["load_data"])# deactivate logging the stage `load_data`, in this case for example because it was loaded from a previous runmlflowhelper.set_skip_log(stages=["load_data"])withmlflowhelper.managed_artifact_dir("data",stage="load_data")asartifact_dir:train_path=artifact_dir.get_path("test.csv")test_path=artifact_dir.get_path("train.csv")ifartifact_dir.loaded:# load artifactstrain=pd.read_csv(train_path)test=pd.read_csv(test_path)else:data=pd.read_csv("/shared/dir/data.csv").sample(frac=1)train=data.iloc[:100,:]test=data.iloc[100:,:]# save artifactstrain.to_csv(train_path)test.to_csv(test_path)Note:For central managing thestageparameter must be set inmlflowhelper.managed_artifact(_dir).Easy parameter loggingmlflowhelperhelps you to never forget logging parameters again by making it easy to log all existing variables
usingmlflowhelper.log_vars.importmlflowhelperdefmain(param1,param2,param3="defaultvalue",verbose=0,*args,**kwargs):some_variable="x"withmlflowhelper.start_run():# mlflow.start_run() is also OK heremlflowhelper.log_vars(exclude=["verbose"])if__name__=='__main__':main("a","b",something_else=6)This will log:{"param1":"a","param2":"b","param3":"defaultvalue","something_else":6}Persistent dictionarymlflowhelperprovides dictionary-like implementation that persistens elements to MLFlow.frommlflowhelper.tracking.collectionsimportMlflowDictd=MlflowDict()# you can also provide tracking URI or an MlflowClientd["a"]=5deldd=MlflowDict()# you can also provide tracking URI or an MlflowClientprint(d["a"])# will give you 5OtherThere are a few more convenience functions included inmlflowhelper:TODOs / Ideascheck if loading works across experimentspurge local artifacts (check via API which runs are marked as deleted and delete their artifacts)support nested runs by creating subdirectories based on experiment and runsupport loading from central cache instead of from runsautomatically log from where and what has been loadedset tags for logged stages (to check for artifacts before loading them)consider loading extensions:does nested loading make sense (different loads for certain nested runs)?does mixed loading make sense (loading artifacts from different runs for different stages)?NoteThis project has been set up using PyScaffold 3.2.1. For details and usage
information on PyScaffold seehttps://pyscaffold.org/.