wandb 0.21.0__py3-none-win_amd64.whl → 0.21.2__py3-none-win_amd64.whl

wandb/integration/metaflow/errors.py

@@ -0,0 +1,13 @@
+import wandb
+
+
+class MissingDependencyError(Exception):
+    """An optional dependency is missing."""
+
+    def __init__(self, *args: object, warning: str) -> None:
+        super().__init__(*args)
+        self._wb_warning = warning
+
+    def warn(self) -> None:
+        """Print a warning for the problem."""
+        wandb.termwarn(self._wb_warning)

wandb/integration/metaflow/metaflow.py

@@ -1,174 +1,40 @@
-"""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
 
-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`"
-    ) from e
+from . import errors
 
 try:
-    from plum import dispatch
+    from metaflow import current
 except ImportError as e:
     raise Exception(
-        "Error: `plum-dispatch` not installed >> "
-        "This integration requires plum-dispatch! To fix, please `pip install -Uqq plum-dispatch`"
+        "Error: `metaflow` not installed >> This integration requires metaflow!"
+        " To fix, please `pip install -Uqq metaflow`"
     ) from e
 
 
 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(
-        name: str,
-        data: nn.Module,
-        models=False,
-        run=None,
-        testing=False,
-        *args,
-        **kwargs,
-    ):  # type: ignore
-        if testing:
-            return "models" if models else None
-
-        if models:
-            run.use_artifact(f"{name}:latest")
-            wandb.termlog(f"Using artifact: {name} ({type(data)})")
-
-    @dispatch
-    def wandb_track(
-        name: str,
-        data: nn.Module,
-        models=False,
-        run=None,
-        testing=False,
-        *args,
-        **kwargs,
-    ):
-        if testing:
-            return "nn.Module" if models else None
-
-        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)})")
-
-except ImportError:
-    wandb.termwarn(
-        "`pytorch` not installed >> @wandb_log(models=True) may not auto log your model!"
-    )
+    from . import data_pytorch
+except errors.MissingDependencyError as e:
+    e.warn()
+    data_pytorch = None
 
 try:
-    from sklearn.base import BaseEstimator
-
-    @dispatch
-    def _wandb_use(
-        name: str,
-        data: BaseEstimator,
-        models=False,
-        run=None,
-        testing=False,
-        *args,
-        **kwargs,
-    ):  # type: ignore
-        if testing:
-            return "models" if models else None
-
-        if models:
-            run.use_artifact(f"{name}:latest")
-            wandb.termlog(f"Using artifact: {name} ({type(data)})")
-
-    @dispatch
-    def wandb_track(
-        name: str,
-        data: BaseEstimator,
-        models=False,
-        run=None,
-        testing=False,
-        *args,
-        **kwargs,
-    ):
-        if testing:
-            return "BaseEstimator" if models else None
-
-        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)})")
-
-except ImportError:
-    wandb.termwarn(
-        "`sklearn` not installed >> @wandb_log(models=True) may not auto log your model!"
-    )
+    from . import data_sklearn
+except errors.MissingDependencyError as e:
+    e.warn()
+    data_sklearn = None
 
 
 class ArtifactProxy:
@@ -194,93 +60,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
+        return "Path"
 
-    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)})")
+    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 data_pytorch and data_pytorch.is_nn_module(data) and models:
+        return data_pytorch.track_nn_module(name, data, run, testing)
+
+    # Check for scikit-learn BaseEstimator
+    if data_sklearn and data_sklearn.is_estimator(data) and models:
+        return data_sklearn.track_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
 
 
-@dispatch
-def wandb_use(name: str, data, *args, **kwargs):
+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
+
     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 data_pytorch and data_pytorch.is_nn_module(data) and models:
+            return data_pytorch.use_nn_module(name, run, testing)
+
+        # Check for scikit-learn BaseEstimator
+        elif data_sklearn and data_sklearn.is_estimator(data) and models:
+            return data_sklearn.use_estimator(name, 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 +228,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/weave/__init__.py

@@ -0,0 +1,6 @@
+"""Weave integration for W&B."""
+
+from .interface import RunPath, active_run_path
+from .weave import setup
+
+__all__ = ("active_run_path", "RunPath", "setup")

wandb/integration/weave/interface.py

@@ -0,0 +1,49 @@
+"""Internal APIs for integrating with weave.
+
+The public functions here are intended to be called by weave and care should
+be taken to maintain backward compatibility.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+
+from wandb.sdk import wandb_setup
+
+
+@dataclasses.dataclass(frozen=True)
+class RunPath:
+    entity: str
+    """The entity to which the run is logging. Never empty."""
+
+    project: str
+    """The project to which the run is logging. Never empty."""
+
+    run_id: str
+    """The run's ID. Never empty."""
+
+
+def active_run_path() -> RunPath | None:
+    """Returns the path of an initialized, unfinished run.
+
+    Returns None if all initialized runs are finished. If there is
+    more than one active run, an arbitrary path is returned.
+    The run may be finished by the time its path is returned.
+
+    Thread-safe.
+    """
+    singleton = wandb_setup.singleton()
+
+    if (
+        (run := singleton.most_recent_active_run)
+        and run.entity
+        and run.project
+        and run.id
+    ):
+        return RunPath(
+            entity=run.entity,
+            project=run.project,
+            run_id=run.id,
+        )
+
+    return None

wandb/integration/weave/weave.py

@@ -0,0 +1,63 @@
+"""Integration module for automatic Weave initialization with W&B.
+
+This module provides automatic initialization of Weave when:
+1. Weave is installed
+2. A W&B run is active with a project
+3. Weave is imported (init-on-import)
+
+The integration can be disabled by setting the WANDB_DISABLE_WEAVE environment variable.
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+import threading
+
+import wandb
+
+_weave_init_lock = threading.Lock()
+
+_DISABLE_WEAVE = "WANDB_DISABLE_WEAVE"
+_WEAVE_PACKAGE_NAME = "weave"
+
+
+def setup(entity: str | None, project: str | None) -> None:
+    """Set up automatic Weave initialization for the current W&B run.
+
+    Args:
+        project: The W&B project name to use for Weave initialization.
+    """
+    # We can't or shouldn't init weave; return
+    if os.getenv(_DISABLE_WEAVE):
+        return
+    if not project:
+        return
+
+    # Use entity/project when available; otherwise fall back to project only
+    if entity:
+        project_path = f"{entity}/{project}"
+    else:
+        project_path = project
+
+    # If weave is not yet imported, we can't init it from here.  Instead, we'll
+    # rely on the weave library itself to detect a run and init itself.
+    if _WEAVE_PACKAGE_NAME not in sys.modules:
+        return
+
+    # If weave has already been imported, initialize immediately
+    with _weave_init_lock:
+        try:
+            # This import should have already happened, so it's effectively a no-op.
+            # We just import to keep the symbol for the init that follows
+            import weave
+        except ImportError:
+            # This should never happen; but we don't raise here to avoid
+            # breaking the wandb run init flow just in case
+            return
+
+        wandb.termlog("Initializing weave.")
+        try:
+            weave.init(project_path)
+        except Exception as e:
+            wandb.termwarn(f"Failed to automatically initialize Weave: {e}")

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