flytekit icon indicating copy to clipboard operation
flytekit copied to clipboard

add support for `dir` in wandb plugin initialization

Open samhita-alla opened this issue 2 months ago • 1 comments

Tracking issue

Why are the changes needed?

To enable users to specify dir in the wandb_init initialization, we can allow a string that refers to the FlyteDirectory parameter name. The plugin would then download that directory and pass the local path to the wandb_init call.

What changes were proposed in this pull request?

This PR adds support for materializing dir in the wandb plugin.

@task(
    container_image=image,
    secret_requests=[wandb_secret],
)
@wandb_init(
    project=WANDB_PROJECT,
    entity=WANDB_ENTITY,
    secret=wandb_secret,
    dir="run_dir",
)
def t1(run_dir: union.FlyteDirectory) -> union.FlyteDirectory:
    wandb.log({"accuracy": 0.9, "loss": 0.1})
    return union.FlyteDirectory(path=wandb.run.dir)

How was this patch tested?

Setup process

Screenshots

Check all the applicable boxes

  • [ ] I updated the documentation accordingly.
  • [ ] All new and existing tests passed.
  • [ ] All commits are signed-off.

Related PRs

Docs link

Summary by Bito

This pull request introduces a new feature allowing users to specify a 'dir' parameter during wandb plugin initialization. It enhances the plugin's ability to handle FlyteDirectory parameters, improving the flexibility and usability of wandb integration within Flyte tasks.

samhita-alla avatar Oct 22 '25 09:10 samhita-alla

Bito Automatic Review Skipped - Draft PR

Bito didn't auto-review because this pull request is in draft status.
No action is needed if you didn't intend for the agent to review it. Otherwise, to manually trigger a review, type /review in a comment and save.
You can change draft PR review settings here, or contact your Bito workspace admin at [email protected].

flyte-bot avatar Oct 22 '25 09:10 flyte-bot