helpers.py 1016 B

12345678910111213141516171819202122232425262728
  1. import json
  2. def add_wandb_visualization(run, mlpipeline_ui_metadata_path):
  3. """NOTE: To use this, you must modify your component to have an output called `mlpipeline_ui_metadata_path` AND call `wandb.init` yourself inside that component.
  4. Example usage:
  5. def my_component(..., mlpipeline_ui_metadata_path: OutputPath()):
  6. import wandb
  7. from wandb.integration.kfp.helpers import add_wandb_visualization
  8. with wandb.init() as run:
  9. add_wandb_visualization(run, mlpipeline_ui_metadata_path)
  10. ... # the rest of your code here
  11. """
  12. def get_iframe_html(run):
  13. return f'<iframe src="{run.url}?kfp=true" style="border:none;width:100%;height:100%;min-width:900px;min-height:600px;"></iframe>'
  14. iframe_html = get_iframe_html(run)
  15. metadata = {
  16. "outputs": [{"type": "markdown", "storage": "inline", "source": iframe_html}]
  17. }
  18. with open(mlpipeline_ui_metadata_path, "w") as metadata_file:
  19. json.dump(metadata, metadata_file)