wandb 0.21.0__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/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/custom_chart.py

@@ -85,27 +85,50 @@ def plot_table(
     This function creates a custom chart based on a Vega-Lite specification and
     a data table represented by a `wandb.Table` object. The specification needs
     to be predefined and stored in the W&B backend. The function returns a custom
-    chart object that can be logged to W&B using `wandb.log()`.
+    chart object that can be logged to W&B using `wandb.Run.log()`.
 
     Args:
-        vega_spec_name (str): The name or identifier of the Vega-Lite spec
+        vega_spec_name: The name or identifier of the Vega-Lite spec
             that defines the visualization structure.
-        data_table (wandb.Table): A `wandb.Table` object containing the data to be
+        data_table: A `wandb.Table` object containing the data to be
             visualized.
-        fields (dict[str, Any]): A mapping between the fields in the Vega-Lite spec and the
+        fields: A mapping between the fields in the Vega-Lite spec and the
             corresponding columns in the data table to be visualized.
-        string_fields (dict[str, Any] | None): A dictionary for providing values for any string constants
+        string_fields: A dictionary for providing values for any string constants
             required by the custom visualization.
-        split_table (bool): Whether the table should be split into a separate section
+        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`.
 
     Returns:
         CustomChart: A custom chart object that can be logged to W&B. To log the
-            chart, pass it to `wandb.log()`.
+            chart, pass the chart object as argument to `wandb.Run.log()`.
 
     Raises:
         wandb.Error: If `data_table` is not a `wandb.Table` object.
+
+    Example:
+    ```python
+    # Create a custom chart using a Vega-Lite spec and the data table.
+    import wandb
+
+    data = [[1, 1], [2, 2], [3, 3], [4, 4], [5, 5]]
+    table = wandb.Table(data=data, columns=["x", "y"])
+    fields = {"x": "x", "y": "y", "title": "MY TITLE"}
+
+    with wandb.init() as run:
+        # Training code goes here
+
+        # Create a custom title with `string_fields`.
+        my_custom_chart = wandb.plot_table(
+            vega_spec_name="wandb/line/v0",
+            data_table=table,
+            fields=fields,
+            string_fields={"title": "Title"},
+        )
+
+        run.log({"custom_chart": my_custom_chart})
+    ```
     """
 
     if not isinstance(data_table, wandb.Table):