Package: tango-mlflow
tango-mlflowMLflow integration forai2-tangoIntroductiontango-mlflowis a Python library that connects theTangotoMLflow.
Tango, developed by AllenAI, is a flexible pipeline library for managing experiments and caching outputs of each step.
MLflow, on the other hand, is a platform that helps manage the Machine Learning lifecycle, including experimentation, reproducibility, and deployment.
This integration enables you to store and manage complete experimental settings and artifacts of your Tango run in MLflow, thus enhancing your experiment tracking capabilities.Here's a screenshot of the MLflow interface when executing withtango-mlflow:Tango run: The top-level run in MLflow corresponds to a single execution in Tango, and its name matches the execution name in Tango.Tango steps: Results of each step in Tango are recorded as nested runs in MLflow. The names of these nested runs correspond to the names of the steps in Tango.Parameters and Metrics: The entire settings for the execution, as well as the parameters for each step, are automatically logged in MLflow.Artifacts and Caching: The cached outputs of each step's execution are saved as artifacts under the corresponding MLflow run. These can be reused when executing Tango, enhancing efficiency and reproducibility.InstallationInstalltango-mlflowby running the following command in your terminal:pipinstalltango-mlflow[all]UsageYou can use theMLFlowWorkspacewith command line arguments as follows:tangorun--workspacemlflow://your_experiment_name--include-packagetango_mlflowAlternatively, you can define your configuration in atango.ymlfile:workspace:
  type: mlflow
  experiment_name: your_experiment_name

include_package:
  - tango_mlflowIn the abovetango.ymlconfiguration,typespecifies the workspace type as MLflow, andexperiment_namesets the name of your experiment in MLflow.
Theinclude_packagefield needs to include thetango_mlflowpackage to usetango-mlflowfunctionality.Remember to replaceyour_experiment_namewith the name of your specific experiment.To log runs remotely, set theMLFLOW_TRACKING_URIenvironment variable to a tracking server’s URI like below:exportMLFLOW_TRACKING_URI=https://mlflow.example.comFunctionalitiesLogging metrics into MLflowThetango-mlflowpackage provides theMLflowStepclass, which allows you to easily log the results of each step execution to MLflow.fromtango_mlflow.stepimportMLflowStepclassTrainModel(MLflowStep):defrun(self,**params):# pre-process...forepochinrange(max_epochs):loss=train_epoch(...)metrics=evaluate_model(...)# log metrics with mlflow_loggerforname,valueinmetrics.items():self.mlflow_logger.log_metric(name,value,step=epoch)# post-process...In the example above, theTrainModelstep inherits fromMLflowStep.
Inside the step, you can directly record metrics to the corresponding MLflow run by invokingself.mlflow_logger.log_metric(...).Please note, this functionality must be used in conjunction withMLFlowWorkspace.Summarizing Tango run metricsYou can specify a step to record its returned metrics as representative values of the Tango run by setting the class variableMLFLOW_SUMMARY = True.
This feature enables you to conveniently view metrics for each Tango run directly in the MLflow interfaceclassEvaluateModel(Step):MLFLOW_SUMMARY=True# Enables MLflow summary!defrun(self,...)->dict[str,float]:# compute metrics ...returnmetricsIn the example above, theEvaluateModelstep returns metrics that are logged as the representative values for that Tango run. These metrics are then recorded in the corresponding (top-level) MLflow run.Please note the following requirements:The return value of a step whereMLFLOW_SUMMARY = Trueis set must always bedict[str, float].You don't necessarily need to inherit fromMLflowStepto useMLFLOW_SUMMARY.Tuning hyperparameters with Optunatango-mlflowalso provides thetango-mlflow tunecommand for tuning hyperparameters withOptuna.
For more details, please refer to theexamples/breast_cancerdirectory.ExamplesBasic example:examples/eulerHyper parameter tuning withOptuna:examples/breast_cancer