Adds flytekitplugin.wandb

[thomasjpfan]

Tracking issue

Related to flyteorg/flyte#2798

Why are the changes needed?

This PR adds a wandb plugin to better integrate with Weights and Biases.

What changes were proposed in this pull request?

This PR proposes a simple wandb_init task decorator that gives the user full control on how the project, entity are named in Weights and Biases. If the id is not given, then the HOSTNAME is used for wandb's run id. Currently, HOSTNAME is set to {.executionName}-{.nodeID}-{.taskRetryAttempt}, which is unique to the node.

How was this patch tested?

This PR adds unit tests and docs to help users enable the feature.

For local testing, I've built an image with this plugin installed at ghcr.io/thomasjpfan/wandb:0.0.4 and ran:

Workflow contents

from flytekit import task, Secret, workflow

from flytekitplugins.wandb import wandb_init

WANDB_PROJECT = "flytekit-wandb-plugin"
WANDB_ENTITY = "username"
SECRET_KEY = "wandb-api-key"
SECRET_GROUP = "wandb-api-group"
wandb_secret = Secret(key=WANDB_SECRET_KEY, group=WANDB_SECRET_GROUP)


@task(
    container_image="ghcr.io/thomasjpfan/wandb:0.0.4",
    secret_requests=[wandb_secret],
)
@wandb_init(
    project=WANDB_PROJECT,
    entity=WANDB_ENTITY,
    secret=wandb_secret,
)
def train() -> float:
    from xgboost import XGBClassifier
    from wandb.integration.xgboost import WandbCallback
    from sklearn.datasets import load_iris
    from sklearn.model_selection import train_test_split

    import wandb

    X, y = load_iris(return_X_y=True)
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)
    bst = XGBClassifier(
        n_estimators=100,
        objective="binary:logistic",
        callbacks=[WandbCallback(log_model=True)],
    )
    bst.fit(X_train, y_train)

    test_score = bst.score(X_test, y_test)

    # Log custom metrics
    wandb.run.log({"test_score": test_score})
    return test_score


@workflow
def wf():
    train()

Check all the applicable boxes

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

Docs link

The associated flytesnacks docs are in: flyteorg/flytesnacks#1673

With multiple runs, this is what shows up in the Weights and Bias UI, the runs are ID with Flyte specific run_ids:

[pingsutw]

Is it necessary to specify the secret key two times? (@task and @wandb_init)

[kumare3]

I agree, you can simply use the one specified in the task. as it will be part of the context

[thomasjpfan]

I updated the PR to pass the secret object around:

wandb_secret = Secret(key=WANDB_SECRET_KEY, group=WANDB_SECRET_GROUP)

@task(
    container_image=image,
    secret_requests=[wandb_secret],
)
@wandb_init(
    project=WANDB_PROJECT,
    entity=WANDB_ENTITY,
    secret=wandb_secret,
)

We can do something more implicit by pulling the secret request out of the task, but the wandb_init decorator still needs to know the name of the Secret associated with W&B. I prefer the more explicit approach of passing the Secret object around.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 75.83%. Comparing base (edab1e3) to head (5445bc3).

❗ Current head 5445bc3 differs from pull request most recent head 7516fa3

Please upload reports for the commit 7516fa3 to get more accurate results.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #2405      +/-   ##
==========================================
- Coverage   75.85%   75.83%   -0.02%     
==========================================
  Files         181      181              
  Lines       18395    18393       -2     
  Branches     3601     3600       -1     
==========================================
- Hits        13953    13949       -4     
- Misses       3840     3841       +1     
- Partials      602      603       +1     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[kumare3]

This is incredible!

[thomasjpfan]

Giving Weights and Biases two keys here to support:

1. Our auto-generated execution id -> Default unique id for Weights and Biases
2. User provided id -> useful for continue training.

[wild-endeavor]

why lazy here out of curiosity? wandb is a dependency of this plugin, so if the env can access this file shouldn't it also be able to access wandb?

[thomasjpfan]

I saw it as a pattern in the mlflow plugin:

flytekit/plugins/flytekit-mlflow/flytekitplugins/mlflow/tracking.py

Lines 9 to 12 in 76fb7c3

	go = lazy_module("plotly.graph_objects")
	plotly_subplots = lazy_module("plotly.subplots")
	pd = lazy_module("pandas")
	mlflow = lazy_module("mlflow")

Personally, I would go with a normal import.

[eapolinario]

For direct dependencies (which is the case, since wandb is listed in setup.py) we should go with a regular import.

[wild-endeavor]

in the future, do you want to also add P, R here?

[thomasjpfan]

Yea, f we want this decorator to type correctly with @task, we'll likely end up needing ParamSpec and TypeVar here.