wandb 0.20.2rc20250616__py3-none-win_amd64.whl → 0.21.1__py3-none-win_amd64.whl

wandb/integration/metaflow/metaflow.py

@@ -1,124 +1,67 @@
-"""W&B Integration for Metaflow.
-
-This integration lets users apply decorators to Metaflow flows and steps to automatically log parameters and artifacts to W&B by type dispatch.
-
-- Decorating a step will enable or disable logging for certain types within that step
-- Decorating the flow is equivalent to decorating all steps with a default
-- Decorating a step after decorating the flow will overwrite the flow decoration
-
-Examples can be found at wandb/wandb/functional_tests/metaflow
-"""
-
 import inspect
 import pickle
 from functools import wraps
 from pathlib import Path
-from typing import Union
+from typing import Optional, Union
 
 import wandb
 from wandb.sdk.lib import telemetry as wb_telemetry
 
+from . import errors
+
 try:
     from metaflow import current
 except ImportError as e:
     raise Exception(
-        "Error: `metaflow` not installed >> This integration requires metaflow!  To fix, please `pip install -Uqq metaflow`"
+        "Error: `metaflow` not installed >> This integration requires metaflow!"
+        " To fix, please `pip install -Uqq metaflow`"
     ) from e
 
-try:
-    from plum import dispatch
-except ImportError as e:
-    raise Exception(
-        "Error: `plum-dispatch` not installed >> "
-        "This integration requires plum-dispatch! To fix, please `pip install -Uqq plum-dispatch`"
-    ) from e
 
+# Classes for isinstance() checks.
+_NN_MODULE = None
+_BASE_ESTIMATOR = None
 
 try:
-    import pandas as pd
-
-    @dispatch
-    def _wandb_use(
-        name: str,
-        data: pd.DataFrame,
-        datasets=False,
-        run=None,
-        testing=False,
-        *args,
-        **kwargs,
-    ):  # type: ignore
-        if testing:
-            return "datasets" if datasets else None
-
-        if datasets:
-            run.use_artifact(f"{name}:latest")
-            wandb.termlog(f"Using artifact: {name} ({type(data)})")
-
-    @dispatch
-    def wandb_track(
-        name: str,
-        data: pd.DataFrame,
-        datasets=False,
-        run=None,
-        testing=False,
-        *args,
-        **kwargs,
-    ):
-        if testing:
-            return "pd.DataFrame" if datasets else None
-
-        if datasets:
-            artifact = wandb.Artifact(name, type="dataset")
-            with artifact.new_file(f"{name}.parquet", "wb") as f:
-                data.to_parquet(f, engine="pyarrow")
-            run.log_artifact(artifact)
-            wandb.termlog(f"Logging artifact: {name} ({type(data)})")
-
-except ImportError:
-    wandb.termwarn(
-        "`pandas` not installed >> @wandb_log(datasets=True) may not auto log your dataset!"
-    )
+    from . import data_pandas
+except errors.MissingDependencyError as e:
+    e.warn()
+    data_pandas = None
 
 try:
     import torch
     import torch.nn as nn
 
-    @dispatch
-    def _wandb_use(
+    _NN_MODULE = nn.Module
+
+    def _use_torch_module(
         name: str,
         data: nn.Module,
-        models=False,
-        run=None,
-        testing=False,
-        *args,
-        **kwargs,
-    ):  # type: ignore
+        run,
+        testing: bool = False,
+    ) -> Optional[str]:
         if testing:
-            return "models" if models else None
+            return "models"
 
-        if models:
-            run.use_artifact(f"{name}:latest")
-            wandb.termlog(f"Using artifact: {name} ({type(data)})")
+        run.use_artifact(f"{name}:latest")
+        wandb.termlog(f"Using artifact: {name} ({type(data)})")
+        return None
 
-    @dispatch
-    def wandb_track(
+    def _track_torch_module(
         name: str,
         data: nn.Module,
-        models=False,
-        run=None,
-        testing=False,
-        *args,
-        **kwargs,
-    ):
+        run,
+        testing: bool = False,
+    ) -> Optional[str]:
         if testing:
-            return "nn.Module" if models else None
+            return "nn.Module"
 
-        if models:
-            artifact = wandb.Artifact(name, type="model")
-            with artifact.new_file(f"{name}.pkl", "wb") as f:
-                torch.save(data, f)
-            run.log_artifact(artifact)
-            wandb.termlog(f"Logging artifact: {name} ({type(data)})")
+        artifact = wandb.Artifact(name, type="model")
+        with artifact.new_file(f"{name}.pkl", "wb") as f:
+            torch.save(data, f)
+        run.log_artifact(artifact)
+        wandb.termlog(f"Logging artifact: {name} ({type(data)})")
+        return None
 
 except ImportError:
     wandb.termwarn(
@@ -128,42 +71,36 @@ except ImportError:
 try:
     from sklearn.base import BaseEstimator
 
-    @dispatch
-    def _wandb_use(
+    _BASE_ESTIMATOR = BaseEstimator
+
+    def _use_sklearn_estimator(
         name: str,
         data: BaseEstimator,
-        models=False,
-        run=None,
-        testing=False,
-        *args,
-        **kwargs,
-    ):  # type: ignore
+        run,
+        testing: bool = False,
+    ) -> Optional[str]:
         if testing:
-            return "models" if models else None
+            return "models"
 
-        if models:
-            run.use_artifact(f"{name}:latest")
-            wandb.termlog(f"Using artifact: {name} ({type(data)})")
+        run.use_artifact(f"{name}:latest")
+        wandb.termlog(f"Using artifact: {name} ({type(data)})")
+        return None
 
-    @dispatch
-    def wandb_track(
+    def _track_sklearn_estimator(
         name: str,
         data: BaseEstimator,
-        models=False,
-        run=None,
-        testing=False,
-        *args,
-        **kwargs,
-    ):
+        run,
+        testing: bool = False,
+    ) -> Optional[str]:
         if testing:
-            return "BaseEstimator" if models else None
+            return "BaseEstimator"
 
-        if models:
-            artifact = wandb.Artifact(name, type="model")
-            with artifact.new_file(f"{name}.pkl", "wb") as f:
-                pickle.dump(data, f)
-            run.log_artifact(artifact)
-            wandb.termlog(f"Logging artifact: {name} ({type(data)})")
+        artifact = wandb.Artifact(name, type="model")
+        with artifact.new_file(f"{name}.pkl", "wb") as f:
+            pickle.dump(data, f)
+        run.log_artifact(artifact)
+        wandb.termlog(f"Logging artifact: {name} ({type(data)})")
+        return None
 
 except ImportError:
     wandb.termwarn(
@@ -194,93 +131,166 @@ class ArtifactProxy:
         return getattr(self.flow, key)
 
 
-@dispatch
-def wandb_track(
+def _track_scalar(
     name: str,
     data: Union[dict, list, set, str, int, float, bool],
-    run=None,
-    testing=False,
-    *args,
-    **kwargs,
-):  # type: ignore
+    run,
+    testing: bool = False,
+) -> Optional[str]:
     if testing:
         return "scalar"
 
     run.log({name: data})
+    return None
 
 
-@dispatch
-def wandb_track(
-    name: str, data: Path, datasets=False, run=None, testing=False, *args, **kwargs
-):
+def _track_path(
+    name: str,
+    data: Path,
+    run,
+    testing: bool = False,
+) -> Optional[str]:
     if testing:
-        return "Path" if datasets else None
-
-    if datasets:
-        artifact = wandb.Artifact(name, type="dataset")
-        if data.is_dir():
-            artifact.add_dir(data)
-        elif data.is_file():
-            artifact.add_file(data)
-        run.log_artifact(artifact)
-        wandb.termlog(f"Logging artifact: {name} ({type(data)})")
+        return "Path"
 
+    artifact = wandb.Artifact(name, type="dataset")
+    if data.is_dir():
+        artifact.add_dir(data)
+    elif data.is_file():
+        artifact.add_file(data)
+    run.log_artifact(artifact)
+    wandb.termlog(f"Logging artifact: {name} ({type(data)})")
+    return None
 
-# this is the base case
-@dispatch
-def wandb_track(
-    name: str, data, others=False, run=None, testing=False, *args, **kwargs
-):
+
+def _track_generic(
+    name: str,
+    data,
+    run,
+    testing: bool = False,
+) -> Optional[str]:
     if testing:
-        return "generic" if others else None
+        return "generic"
+
+    artifact = wandb.Artifact(name, type="other")
+    with artifact.new_file(f"{name}.pkl", "wb") as f:
+        pickle.dump(data, f)
+    run.log_artifact(artifact)
+    wandb.termlog(f"Logging artifact: {name} ({type(data)})")
+    return None
 
+
+def wandb_track(
+    name: str,
+    data,
+    datasets: bool = False,
+    models: bool = False,
+    others: bool = False,
+    run: Optional[wandb.Run] = None,
+    testing: bool = False,
+) -> Optional[str]:
+    """Track data as wandb artifacts based on type and flags."""
+    # Check for pandas DataFrame
+    if data_pandas and data_pandas.is_dataframe(data) and datasets:
+        return data_pandas.track_dataframe(name, data, run, testing)
+
+    # Check for PyTorch Module
+    if _NN_MODULE and isinstance(data, _NN_MODULE) and models:
+        return _track_torch_module(name, data, run, testing)
+
+    # Check for scikit-learn BaseEstimator
+    if _BASE_ESTIMATOR and isinstance(data, _BASE_ESTIMATOR) and models:
+        return _track_sklearn_estimator(name, data, run, testing)
+
+    # Check for Path objects
+    if isinstance(data, Path) and datasets:
+        return _track_path(name, data, run, testing)
+
+    # Check for scalar types
+    if isinstance(data, (dict, list, set, str, int, float, bool)):
+        return _track_scalar(name, data, run, testing)
+
+    # Generic fallback
     if others:
-        artifact = wandb.Artifact(name, type="other")
-        with artifact.new_file(f"{name}.pkl", "wb") as f:
-            pickle.dump(data, f)
-        run.log_artifact(artifact)
-        wandb.termlog(f"Logging artifact: {name} ({type(data)})")
+        return _track_generic(name, data, run, testing)
 
+    # No action taken
+    return None
+
+
+def wandb_use(
+    name: str,
+    data,
+    datasets: bool = False,
+    models: bool = False,
+    others: bool = False,
+    run=None,
+    testing: bool = False,
+) -> Optional[str]:
+    """Use wandb artifacts based on data type and flags."""
+    # Skip scalar types - nothing to use
+    if isinstance(data, (dict, list, set, str, int, float, bool)):
+        return None
 
-@dispatch
-def wandb_use(name: str, data, *args, **kwargs):
     try:
-        return _wandb_use(name, data, *args, **kwargs)
+        # Check for pandas DataFrame
+        if data_pandas and data_pandas.is_dataframe(data) and datasets:
+            return data_pandas.use_dataframe(name, run, testing)
+
+        # Check for PyTorch Module
+        elif _NN_MODULE and isinstance(data, _NN_MODULE) and models:
+            return _use_torch_module(name, data, run, testing)
+
+        # Check for scikit-learn BaseEstimator
+        elif _BASE_ESTIMATOR and isinstance(data, _BASE_ESTIMATOR) and models:
+            return _use_sklearn_estimator(name, data, run, testing)
+
+        # Check for Path objects
+        elif isinstance(data, Path) and datasets:
+            return _use_path(name, data, run, testing)
+
+        # Generic fallback
+        elif others:
+            return _use_generic(name, data, run, testing)
+
+        else:
+            return None
+
     except wandb.CommError:
         wandb.termwarn(
             f"This artifact ({name}, {type(data)}) does not exist in the wandb datastore!"
-            f"If you created an instance inline (e.g. sklearn.ensemble.RandomForestClassifier), then you can safely ignore this"
-            f"Otherwise you may want to check your internet connection!"
+            " If you created an instance inline (e.g. sklearn.ensemble.RandomForestClassifier),"
+            " then you can safely ignore this. Otherwise you may want to check your internet connection!"
         )
+        return None
 
 
-@dispatch
-def wandb_use(
-    name: str, data: Union[dict, list, set, str, int, float, bool], *args, **kwargs
-):  # type: ignore
-    pass  # do nothing for these types
-
-
-@dispatch
-def _wandb_use(
-    name: str, data: Path, datasets=False, run=None, testing=False, *args, **kwargs
-):  # type: ignore
+def _use_path(
+    name: str,
+    data: Path,
+    run,
+    testing: bool = False,
+) -> Optional[str]:
     if testing:
-        return "datasets" if datasets else None
+        return "datasets"
 
-    if datasets:
-        run.use_artifact(f"{name}:latest")
-        wandb.termlog(f"Using artifact: {name} ({type(data)})")
+    run.use_artifact(f"{name}:latest")
+    wandb.termlog(f"Using artifact: {name} ({type(data)})")
+    return None
 
 
-@dispatch
-def _wandb_use(name: str, data, others=False, run=None, testing=False, *args, **kwargs):  # type: ignore
+def _use_generic(
+    name: str,
+    data,
+    run,
+    testing: bool = False,
+) -> Optional[str]:
     if testing:
-        return "others" if others else None
+        return "others"
 
-    if others:
-        run.use_artifact(f"{name}:latest")
-        wandb.termlog(f"Using artifact: {name} ({type(data)})")
+    run.use_artifact(f"{name}:latest")
+    wandb.termlog(f"Using artifact: {name} ({type(data)})")
+    return None
 
 
 def coalesce(*arg):
@@ -289,25 +299,30 @@ def coalesce(*arg):
 
 def wandb_log(
     func=None,
-    # /,  # py38 only
-    datasets=False,
-    models=False,
-    others=False,
-    settings=None,
+    /,
+    datasets: bool = False,
+    models: bool = False,
+    others: bool = False,
+    settings: Optional[wandb.Settings] = None,
 ):
-    """Automatically log parameters and artifacts to W&B by type dispatch.
+    """Automatically log parameters and artifacts to W&B.
+
+    This decorator can be applied to a flow, step, or both:
 
-    This decorator can be applied to a flow, step, or both.
-    - Decorating a step will enable or disable logging for certain types within that step
-    - Decorating the flow is equivalent to decorating all steps with a default
-    - Decorating a step after decorating the flow will overwrite the flow decoration
+    - Decorating a step enables or disables logging within that step
+    - Decorating a flow is equivalent to decorating all steps
+    - Decorating a step after decorating its flow overwrites the flow decoration
 
     Args:
-        func: (`Callable`). The method or class being decorated (if decorating a step or flow respectively).
-        datasets: (`bool`). If `True`, log datasets.  Datasets can be a `pd.DataFrame` or `pathlib.Path`.  The default value is `False`, so datasets are not logged.
-        models: (`bool`). If `True`, log models.  Models can be a `nn.Module` or `sklearn.base.BaseEstimator`.  The default value is `False`, so models are not logged.
-        others: (`bool`). If `True`, log anything pickle-able.  The default value is `False`, so files are not logged.
-        settings: (`wandb.sdk.wandb_settings.Settings`). Custom settings passed to `wandb.init`.  The default value is `None`, and is the same as passing `wandb.Settings()`.  If `settings.run_group` is `None`, it will be set to `{flow_name}/{run_id}.  If `settings.run_job_type` is `None`, it will be set to `{run_job_type}/{step_name}`
+        func: The step method or flow class to decorate.
+        datasets: Whether to log `pd.DataFrame` and `pathlib.Path`
+            types. Defaults to False.
+        models: Whether to log `nn.Module` and `sklearn.base.BaseEstimator`
+            types. Defaults to False.
+        others: If `True`, log anything pickle-able. Defaults to False.
+        settings: Custom settings to pass to `wandb.init`.
+            If `run_group` is `None`, it is set to `{flow_name}/{run_id}`.
+            If `run_job_type` is `None`, it is set to `{run_job_type}/{step_name}`.
     """
 
     @wraps(func)

wandb/integration/openai/fine_tuning.py

@@ -14,7 +14,6 @@ import wandb
 from wandb import util
 from wandb.data_types import Table
 from wandb.sdk.lib import telemetry
-from wandb.sdk.wandb_run import Run
 
 openai = util.get_module(
     name="openai",
@@ -54,7 +53,7 @@ class WandbLogger:
     _wandb_api: Optional[wandb.Api] = None
     _logged_in: bool = False
     openai_client: Optional[OpenAI] = None
-    _run: Optional[Run] = None
+    _run: Optional[wandb.Run] = None
 
     @classmethod
     def sync(

wandb/integration/sb3/sb3.py

@@ -106,9 +106,9 @@ class WandbCallback(BaseCallback):
             os.makedirs(self.model_save_path, exist_ok=True)
             self.path = os.path.join(self.model_save_path, "model.zip")
         else:
-            assert (
-                self.model_save_freq == 0
-            ), "to use the `model_save_freq` you have to set the `model_save_path` parameter"
+            assert self.model_save_freq == 0, (
+                "to use the `model_save_freq` you have to set the `model_save_path` parameter"
+            )
 
     def _init_callback(self) -> None:
         d = {}

wandb/integration/ultralytics/callback.py

@@ -102,7 +102,9 @@ class WandBUltralyticsCallback:
         model = YOLO("yolov8n.pt")
 
         # add wandb callback
-        add_wandb_callback(model, max_validation_batches=2, enable_model_checkpointing=True)
+        add_wandb_callback(
+            model, max_validation_batches=2, enable_model_checkpointing=True
+        )
 
         # train
         model.train(data="coco128.yaml", epochs=5, imgsz=640)
@@ -454,7 +456,9 @@ def add_wandb_callback(
         model = YOLO("yolov8n.pt")
 
         # add wandb callback
-        add_wandb_callback(model, max_validation_batches=2, enable_model_checkpointing=True)
+        add_wandb_callback(
+            model, max_validation_batches=2, enable_model_checkpointing=True
+        )
 
         # train
         model.train(data="coco128.yaml", epochs=5, imgsz=640)

wandb/jupyter.py

@@ -19,13 +19,13 @@ from requests.compat import urljoin
 
 import wandb
 import wandb.util
-from wandb.sdk import wandb_run, wandb_setup
+from wandb.sdk import wandb_setup
 from wandb.sdk.lib import filesystem
 
 logger = logging.getLogger(__name__)
 
 
-def display_if_magic_is_used(run: wandb_run.Run) -> bool:
+def display_if_magic_is_used(run: wandb.Run) -> bool:
     """Display a run's page if the cell has the %%wandb cell magic.
 
     Args:
@@ -53,7 +53,7 @@ class _WandbCellMagicState:
         self._height = height
         self._already_displayed = False
 
-    def display_if_allowed(self, run: wandb_run.Run) -> None:
+    def display_if_allowed(self, run: wandb.Run) -> None:
         """Display a run's iframe if one is not already displayed.
 
         Args:
@@ -93,7 +93,7 @@ def _display_by_wandb_path(path: str, *, height: int) -> None:
         )
 
 
-def _display_wandb_run(run: wandb_run.Run, *, height: int) -> None:
+def _display_wandb_run(run: wandb.Run, *, height: int) -> None:
     """Display a run (usually in an iframe).
 
     Args:
@@ -461,7 +461,7 @@ class Notebook:
 
         return False
 
-    def save_history(self, run: wandb_run.Run):
+    def save_history(self, run: wandb.Run):
         """This saves all cell executions in the current session as a new notebook."""
         try:
             from nbformat import v4, validator, write  # type: ignore

wandb/plot/__init__.py

@@ -14,6 +14,8 @@ __all__ = [
     "pr_curve",
     "confusion_matrix",
     "line_series",
+    "plot_table",
+    "visualize",  # doc:exclude
 ]
 
 from wandb.plot.bar import bar

wandb/plot/bar.py

@@ -19,11 +19,11 @@ def bar(
     """Constructs a bar chart from a wandb.Table of data.
 
     Args:
-        table (wandb.Table): A table containing the data for the bar chart.
-        label (str): The name of the column to use for the labels of each bar.
-        value (str): The name of the column to use for the values of each bar.
-        title (str): The title of the bar chart.
-        split_table (bool): Whether the table should be split into a separate section
+        table: A table containing the data for the bar chart.
+        label: The name of the column to use for the labels of each bar.
+        value: The name of the column to use for the values of each bar.
+        title: The title of the bar chart.
+        split_table: Whether the table should be split into a separate section
             in the W&B UI. If `True`, the table will be displayed in a section named
             "Custom Chart Tables". Default is `False`.
 
@@ -32,34 +32,35 @@ def bar(
             chart, pass it to `wandb.log()`.
 
     Example:
-        ```
-        import random
-        import wandb
 
-        # Generate random data for the table
-        data = [
-            ["car", random.uniform(0, 1)],
-            ["bus", random.uniform(0, 1)],
-            ["road", random.uniform(0, 1)],
-            ["person", random.uniform(0, 1)],
-        ]
+    ```python
+    import random
+    import wandb
+
+    # Generate random data for the table
+    data = [
+        ["car", random.uniform(0, 1)],
+        ["bus", random.uniform(0, 1)],
+        ["road", random.uniform(0, 1)],
+        ["person", random.uniform(0, 1)],
+    ]
 
-        # Create a table with the data
-        table = wandb.Table(data=data, columns=["class", "accuracy"])
+    # Create a table with the data
+    table = wandb.Table(data=data, columns=["class", "accuracy"])
 
-        # Initialize a W&B run and log the bar plot
-        with wandb.init(project="bar_chart") as run:
-            # Create a bar plot from the table
-            bar_plot = wandb.plot.bar(
-                table=table,
-                label="class",
-                value="accuracy",
-                title="Object Classification Accuracy",
-            )
+    # Initialize a W&B run and log the bar plot
+    with wandb.init(project="bar_chart") as run:
+        # Create a bar plot from the table
+        bar_plot = wandb.plot.bar(
+            table=table,
+            label="class",
+            value="accuracy",
+            title="Object Classification Accuracy",
+        )
 
-            # Log the bar chart to W&B
-            run.log({"bar_plot": bar_plot})
-        ```
+        # Log the bar chart to W&B
+        run.log({"bar_plot": bar_plot})
+    ```
     """
     return plot_table(
         data_table=table,