Skip to content

Conversation

@samhita-alla
Copy link
Contributor

@samhita-alla samhita-alla commented Oct 22, 2025

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.

Signed-off-by: Samhita Alla <aallasamhita@gmail.com>
@flyte-bot
Copy link
Contributor

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 eduardo@union.ai.

@samhita-alla samhita-alla changed the title add support for dir in wandb plugin add support for dir in wandb plugin initialization Oct 22, 2025
@samhita-alla samhita-alla marked this pull request as ready for review October 22, 2025 10:15
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants