wandb 0.19.1rc1__py3-none-musllinux_1_2_aarch64.whl

wandb/integration/kfp/__init__.py

@@ -0,0 +1,6 @@
+__all__ = ["wandb_log", "unpatch_kfp"]
+
+from .kfp_patch import patch_kfp, unpatch_kfp
+from .wandb_logging import wandb_log
+
+patch_kfp()

wandb/integration/kfp/helpers.py

@@ -0,0 +1,28 @@
+import json
+
+
+def add_wandb_visualization(run, mlpipeline_ui_metadata_path):
+    """NOTE: To use this, you must modify your component to have an output called `mlpipeline_ui_metadata_path` AND call `wandb.init` yourself inside that component.
+
+    Example usage:
+
+    def my_component(..., mlpipeline_ui_metadata_path: OutputPath()):
+        import wandb
+        from wandb.integration.kfp.helpers import add_wandb_visualization
+
+        with wandb.init() as run:
+            add_wandb_visualization(run, mlpipeline_ui_metadata_path)
+
+            ... # the rest of your code here
+    """
+
+    def get_iframe_html(run):
+        return f'<iframe src="{run.url}?kfp=true" style="border:none;width:100%;height:100%;min-width:900px;min-height:600px;"></iframe>'
+
+    iframe_html = get_iframe_html(run)
+    metadata = {
+        "outputs": [{"type": "markdown", "storage": "inline", "source": iframe_html}]
+    }
+
+    with open(mlpipeline_ui_metadata_path, "w") as metadata_file:
+        json.dump(metadata, metadata_file)

wandb/integration/kfp/kfp_patch.py

@@ -0,0 +1,334 @@
+import inspect
+import itertools
+import textwrap
+from typing import Callable, List, Mapping, Optional
+
+import wandb
+
+try:
+    from kfp import __version__ as kfp_version
+    from kfp.components import structures
+    from kfp.components._components import _create_task_factory_from_component_spec
+    from kfp.components._python_op import _func_to_component_spec
+
+    from wandb.util import parse_version
+
+    MIN_KFP_VERSION = "1.6.1"
+
+    if parse_version(kfp_version) < parse_version(MIN_KFP_VERSION):
+        wandb.termwarn(
+            f"Your version of kfp {kfp_version} may not work.  This integration requires kfp>={MIN_KFP_VERSION}"
+        )
+
+except ImportError:
+    wandb.termerror("kfp not found!  Please `pip install kfp`")
+
+from .wandb_logging import wandb_log
+
+decorator_code = inspect.getsource(wandb_log)
+wandb_logging_extras = f"""
+import typing
+from typing import NamedTuple
+
+import collections
+from collections import namedtuple
+
+import kfp
+from kfp import components
+from kfp.components import InputPath, OutputPath
+
+import wandb
+
+{decorator_code}
+"""
+
+
+def full_path_exists(full_func):
+    def get_parent_child_pairs(full_func):
+        components = full_func.split(".")
+        parents, children = [], []
+        for i, _ in enumerate(components[:-1], 1):
+            parent = ".".join(components[:i])
+            child = components[i]
+            parents.append(parent)
+            children.append(child)
+        return zip(parents, children)
+
+    for parent, child in get_parent_child_pairs(full_func):
+        module = wandb.util.get_module(parent)
+        if not module or not hasattr(module, child) or getattr(module, child) is None:
+            return False
+    return True
+
+
+def patch(module_name, func):
+    module = wandb.util.get_module(module_name)
+    success = False
+
+    full_func = f"{module_name}.{func.__name__}"
+    if not full_path_exists(full_func):
+        wandb.termerror(
+            f"Failed to patch {module_name}.{func.__name__}!  Please check if this package/module is installed!"
+        )
+    else:
+        wandb.patched.setdefault(module.__name__, [])
+        # if already patched, do not patch again
+        if [module, func.__name__] not in wandb.patched[module.__name__]:
+            setattr(module, f"orig_{func.__name__}", getattr(module, func.__name__))
+            setattr(module, func.__name__, func)
+            wandb.patched[module.__name__].append([module, func.__name__])
+        success = True
+
+    return success
+
+
+def unpatch(module_name):
+    if module_name in wandb.patched:
+        for module, func in wandb.patched[module_name]:
+            setattr(module, func, getattr(module, f"orig_{func}"))
+        wandb.patched[module_name] = []
+
+
+def unpatch_kfp():
+    unpatch("kfp.components")
+    unpatch("kfp.components._python_op")
+    unpatch("wandb.integration.kfp")
+
+
+def patch_kfp():
+    to_patch = [
+        (
+            "kfp.components",
+            create_component_from_func,
+        ),
+        (
+            "kfp.components._python_op",
+            create_component_from_func,
+        ),
+        (
+            "kfp.components._python_op",
+            _get_function_source_definition,
+        ),
+        ("kfp.components._python_op", strip_type_hints),
+    ]
+
+    successes = []
+    for module_name, func in to_patch:
+        success = patch(module_name, func)
+        successes.append(success)
+    if not all(successes):
+        wandb.termerror(
+            "Failed to patch one or more kfp functions.  Patching @wandb_log decorator to no-op."
+        )
+        patch("wandb.integration.kfp", wandb_log)
+
+
+def wandb_log(
+    func=None,
+    # /,  # py38 only
+    log_component_file=True,
+):
+    """Wrap a standard python function and log to W&B.
+
+    NOTE: Because patching failed, this decorator is a no-op.
+    """
+    from functools import wraps
+
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            return func(*args, **kwargs)
+
+        return wrapper
+
+    if func is None:
+        return decorator
+    else:
+        return decorator(func)
+
+
+def _get_function_source_definition(func: Callable) -> str:
+    """Get the source code of a function.
+
+    This function is modified from KFP.  The original source is below:
+    https://github.com/kubeflow/pipelines/blob/b6406b02f45cdb195c7b99e2f6d22bf85b12268b/sdk/python/kfp/components/_python_op.py#L300-L319.
+    """
+    func_code = inspect.getsource(func)
+
+    # Function might be defined in some indented scope (e.g. in another
+    # function). We need to handle this and properly dedent the function source
+    # code
+    func_code = textwrap.dedent(func_code)
+    func_code_lines = func_code.split("\n")
+
+    # For wandb, allow decorators (so we can use the @wandb_log decorator)
+    func_code_lines = itertools.dropwhile(
+        lambda x: not (x.startswith(("def", "@wandb_log"))),
+        func_code_lines,
+    )
+
+    if not func_code_lines:
+        raise ValueError(
+            f'Failed to dedent and clean up the source of function "{func.__name__}". '
+            "It is probably not properly indented."
+        )
+
+    return "\n".join(func_code_lines)
+
+
+def create_component_from_func(
+    func: Callable,
+    output_component_file: Optional[str] = None,
+    base_image: Optional[str] = None,
+    packages_to_install: Optional[List[str]] = None,
+    annotations: Optional[Mapping[str, str]] = None,
+):
+    '''Convert a Python function to a component and returns a task factory.
+
+    The returned task factory accepts arguments and returns a task object.
+
+    This function is modified from KFP.  The original source is below:
+    https://github.com/kubeflow/pipelines/blob/b6406b02f45cdb195c7b99e2f6d22bf85b12268b/sdk/python/kfp/components/_python_op.py#L998-L1110.
+
+    Args:
+        func: The python function to convert
+        base_image: Optional. Specify a custom Docker container image to use in the component. For lightweight components, the image needs to have python 3.5+. Default is the python image corresponding to the current python environment.
+        output_component_file: Optional. Write a component definition to a local file. The produced component file can be loaded back by calling :code:`load_component_from_file` or :code:`load_component_from_uri`.
+        packages_to_install: Optional. List of [versioned] python packages to pip install before executing the user function.
+        annotations: Optional. Allows adding arbitrary key-value data to the component specification.
+
+    Returns:
+        A factory function with a strongly-typed signature taken from the python function.
+        Once called with the required arguments, the factory constructs a task instance that can run the original function in a container.
+
+    Examples:
+        The function name and docstring are used as component name and description. Argument and return annotations are used as component input/output types::
+
+            def add(a: float, b: float) -> float:
+                """Return sum of two arguments"""
+                return a + b
+
+
+            # add_op is a task factory function that creates a task object when given arguments
+            add_op = create_component_from_func(
+                func=add,
+                base_image="python:3.7",  # Optional
+                output_component_file="add.component.yaml",  # Optional
+                packages_to_install=["pandas==0.24"],  # Optional
+            )
+
+            # The component spec can be accessed through the .component_spec attribute:
+            add_op.component_spec.save("add.component.yaml")
+
+            # The component function can be called with arguments to create a task:
+            add_task = add_op(1, 3)
+
+            # The resulting task has output references, corresponding to the component outputs.
+            # When the function only has a single anonymous return value, the output name is "Output":
+            sum_output_ref = add_task.outputs["Output"]
+
+            # These task output references can be passed to other component functions, constructing a computation graph:
+            task2 = add_op(sum_output_ref, 5)
+
+
+        :code:`create_component_from_func` function can also be used as decorator::
+
+            @create_component_from_func
+            def add_op(a: float, b: float) -> float:
+                """Return sum of two arguments"""
+                return a + b
+
+        To declare a function with multiple return values, use the :code:`NamedTuple` return annotation syntax::
+
+            from typing import NamedTuple
+
+
+            def add_multiply_two_numbers(a: float, b: float) -> NamedTuple(
+                "Outputs", [("sum", float), ("product", float)]
+            ):
+                """Return sum and product of two arguments"""
+                return (a + b, a * b)
+
+
+            add_multiply_op = create_component_from_func(add_multiply_two_numbers)
+
+            # The component function can be called with arguments to create a task:
+            add_multiply_task = add_multiply_op(1, 3)
+
+            # The resulting task has output references, corresponding to the component outputs:
+            sum_output_ref = add_multiply_task.outputs["sum"]
+
+            # These task output references can be passed to other component functions, constructing a computation graph:
+            task2 = add_multiply_op(sum_output_ref, 5)
+
+        Bigger data should be read from files and written to files.
+        Use the :py:class:`kfp.components.InputPath` parameter annotation to tell the system that the function wants to consume the corresponding input data as a file. The system will download the data, write it to a local file and then pass the **path** of that file to the function.
+        Use the :py:class:`kfp.components.OutputPath` parameter annotation to tell the system that the function wants to produce the corresponding output data as a file. The system will prepare and pass the **path** of a file where the function should write the output data. After the function exits, the system will upload the data to the storage system so that it can be passed to downstream components.
+
+        You can specify the type of the consumed/produced data by specifying the type argument to :py:class:`kfp.components.InputPath` and :py:class:`kfp.components.OutputPath`. The type can be a python type or an arbitrary type name string. :code:`OutputPath('CatBoostModel')` means that the function states that the data it has written to a file has type :code:`CatBoostModel`. :code:`InputPath('CatBoostModel')` means that the function states that it expect the data it reads from a file to have type 'CatBoostModel'. When the pipeline author connects inputs to outputs the system checks whether the types match.
+        Every kind of data can be consumed as a file input. Conversely, bigger data should not be consumed by value as all value inputs pass through the command line.
+
+        Example of a component function declaring file input and output::
+
+            def catboost_train_classifier(
+                training_data_path: InputPath("CSV"),  # Path to input data file of type "CSV"
+                trained_model_path: OutputPath(
+                    "CatBoostModel"
+                ),  # Path to output data file of type "CatBoostModel"
+                number_of_trees: int = 100,  # Small output of type "Integer"
+            ) -> NamedTuple(
+                "Outputs",
+                [
+                    ("Accuracy", float),  # Small output of type "Float"
+                    ("Precision", float),  # Small output of type "Float"
+                    ("JobUri", "URI"),  # Small output of type "URI"
+                ],
+            ):
+                """Train CatBoost classification model"""
+                ...
+
+                return (accuracy, precision, recall)
+    '''
+    core_packages = ["wandb", "kfp"]
+
+    if not packages_to_install:
+        packages_to_install = core_packages
+    else:
+        packages_to_install += core_packages
+
+    component_spec = _func_to_component_spec(
+        func=func,
+        extra_code=wandb_logging_extras,
+        base_image=base_image,
+        packages_to_install=packages_to_install,
+    )
+    if annotations:
+        component_spec.metadata = structures.MetadataSpec(
+            annotations=annotations,
+        )
+
+    if output_component_file:
+        component_spec.save(output_component_file)
+
+    return _create_task_factory_from_component_spec(component_spec)
+
+
+def strip_type_hints(source_code: str) -> str:
+    """Strip type hints from source code.
+
+    This function is modified from KFP.  The original source is below:
+    https://github.com/kubeflow/pipelines/blob/b6406b02f45cdb195c7b99e2f6d22bf85b12268b/sdk/python/kfp/components/_python_op.py#L237-L248.
+    """
+    # For wandb, do not strip type hints
+
+    #     try:
+    #         return _strip_type_hints_using_lib2to3(source_code)
+    #     except Exception as ex:
+    #         print('Error when stripping type annotations: ' + str(ex))
+
+    #     try:
+    #         return _strip_type_hints_using_strip_hints(source_code)
+    #     except Exception as ex:
+    #         print('Error when stripping type annotations: ' + str(ex))
+
+    return source_code

wandb/integration/kfp/wandb_logging.py

@@ -0,0 +1,182 @@
+def wandb_log(  # noqa: C901
+    func=None,
+    # /,  # py38 only
+    log_component_file=True,
+):
+    """Wrap a standard python function and log to W&B."""
+    import json
+    import os
+    from functools import wraps
+    from inspect import Parameter, signature
+
+    from kfp import components
+    from kfp.components import (
+        InputArtifact,
+        InputBinaryFile,
+        InputPath,
+        InputTextFile,
+        OutputArtifact,
+        OutputBinaryFile,
+        OutputPath,
+        OutputTextFile,
+    )
+
+    import wandb
+    from wandb.sdk.lib import telemetry as wb_telemetry
+
+    output_types = (OutputArtifact, OutputBinaryFile, OutputPath, OutputTextFile)
+    input_types = (InputArtifact, InputBinaryFile, InputPath, InputTextFile)
+
+    def isinstance_namedtuple(x):
+        t = type(x)
+        b = t.__bases__
+        if len(b) != 1 or b[0] is not tuple:
+            return False
+        f = getattr(t, "_fields", None)
+        if not isinstance(f, tuple):
+            return False
+        return all(isinstance(n, str) for n in f)
+
+    def get_iframe_html(run):
+        return f'<iframe src="{run.url}?kfp=true" style="border:none;width:100%;height:100%;min-width:900px;min-height:600px;"></iframe>'
+
+    def get_link_back_to_kubeflow():
+        wandb_kubeflow_url = os.getenv("WANDB_KUBEFLOW_URL")
+        return f"{wandb_kubeflow_url}/#/runs/details/{{workflow.uid}}"
+
+    def log_input_scalar(name, data, run=None):
+        run.config[name] = data
+        wandb.termlog(f"Setting config: {name} to {data}")
+
+    def log_input_artifact(name, data, type, run=None):
+        artifact = wandb.Artifact(name, type=type)
+        artifact.add_file(data)
+        run.use_artifact(artifact)
+        wandb.termlog(f"Using artifact: {name}")
+
+    def log_output_scalar(name, data, run=None):
+        if isinstance_namedtuple(data):
+            for k, v in zip(data._fields, data):
+                run.log({f"{func.__name__}.{k}": v})
+        else:
+            run.log({name: data})
+
+    def log_output_artifact(name, data, type, run=None):
+        artifact = wandb.Artifact(name, type=type)
+        artifact.add_file(data)
+        run.log_artifact(artifact)
+        wandb.termlog(f"Logging artifact: {name}")
+
+    def _log_component_file(func, run=None):
+        name = func.__name__
+        output_component_file = f"{name}.yml"
+        components._python_op.func_to_component_file(func, output_component_file)
+        artifact = wandb.Artifact(name, type="kubeflow_component_file")
+        artifact.add_file(output_component_file)
+        run.log_artifact(artifact)
+        wandb.termlog(f"Logging component file: {output_component_file}")
+
+    # Add `mlpipeline_ui_metadata_path` to signature to show W&B run in "ML Visualizations tab"
+    sig = signature(func)
+    no_default = []
+    has_default = []
+
+    for param in sig.parameters.values():
+        if param.default is param.empty:
+            no_default.append(param)
+        else:
+            has_default.append(param)
+
+    new_params = tuple(
+        (
+            *no_default,
+            Parameter(
+                "mlpipeline_ui_metadata_path",
+                annotation=OutputPath(),
+                kind=Parameter.POSITIONAL_OR_KEYWORD,
+            ),
+            *has_default,
+        )
+    )
+    new_sig = sig.replace(parameters=new_params)
+    new_anns = {param.name: param.annotation for param in new_params}
+    if "return" in func.__annotations__:
+        new_anns["return"] = func.__annotations__["return"]
+
+    def decorator(func):
+        input_scalars = {}
+        input_artifacts = {}
+        output_scalars = {}
+        output_artifacts = {}
+
+        for name, ann in func.__annotations__.items():
+            if name == "return":
+                output_scalars[name] = ann
+            elif isinstance(ann, output_types):
+                output_artifacts[name] = ann
+            elif isinstance(ann, input_types):
+                input_artifacts[name] = ann
+            else:
+                input_scalars[name] = ann
+
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            bound = new_sig.bind(*args, **kwargs)
+            bound.apply_defaults()
+
+            mlpipeline_ui_metadata_path = bound.arguments["mlpipeline_ui_metadata_path"]
+            del bound.arguments["mlpipeline_ui_metadata_path"]
+
+            with wandb.init(
+                job_type=func.__name__,
+                group="{{workflow.annotations.pipelines.kubeflow.org/run_name}}",
+            ) as run:
+                # Link back to the kfp UI
+                kubeflow_url = get_link_back_to_kubeflow()
+                run.notes = kubeflow_url
+                run.config["LINK_TO_KUBEFLOW_RUN"] = kubeflow_url
+
+                iframe_html = get_iframe_html(run)
+                metadata = {
+                    "outputs": [
+                        {
+                            "type": "markdown",
+                            "storage": "inline",
+                            "source": iframe_html,
+                        }
+                    ]
+                }
+
+                with open(mlpipeline_ui_metadata_path, "w") as metadata_file:
+                    json.dump(metadata, metadata_file)
+
+                if log_component_file:
+                    _log_component_file(func, run=run)
+
+                for name, _ in input_scalars.items():
+                    log_input_scalar(name, kwargs[name], run)
+
+                for name, ann in input_artifacts.items():
+                    log_input_artifact(name, kwargs[name], ann.type, run)
+
+                with wb_telemetry.context(run=run) as tel:
+                    tel.feature.kfp_wandb_log = True
+
+                result = func(*bound.args, **bound.kwargs)
+
+                for name, _ in output_scalars.items():
+                    log_output_scalar(name, result, run)
+
+                for name, ann in output_artifacts.items():
+                    log_output_artifact(name, kwargs[name], ann.type, run)
+
+            return result
+
+        wrapper.__signature__ = new_sig
+        wrapper.__annotations__ = new_anns
+        return wrapper
+
+    if func is None:
+        return decorator
+    else:
+        return decorator(func)

wandb/integration/langchain/__init__.py

@@ -0,0 +1,3 @@
+__all__ = ("WandbTracer",)
+
+from .wandb_tracer import WandbTracer

wandb/integration/langchain/wandb_tracer.py

@@ -0,0 +1,48 @@
+"""This module contains an integration with the LangChain library.
+
+Specifically, it exposes a `WandbTracer` class that can be used to stream
+LangChain activity to W&B. The intended usage pattern is to call
+`tracer = WandbTracer()` at the top of the script/notebook, and call
+`tracer.finish()` at the end of the script/notebook.
+ This will stream all LangChain activity to W&B.
+
+Technical Note:
+LangChain is in very rapid development - meaning their APIs and schemas are actively changing.
+As a matter of precaution, any call to LangChain apis, or use of their returned data is wrapped
+in a try/except block. This is to ensure that if a breaking change is introduced, the W&B
+integration will not break user code. The one exception to the rule is at import time. If
+LangChain is not installed, or the symbols are not in the same place, the appropriate error
+will be raised when importing this module.
+"""
+
+from packaging import version
+
+import wandb.util
+from wandb.sdk.lib import deprecate
+
+langchain = wandb.util.get_module(
+    name="langchain",
+    required="To use the LangChain WandbTracer you need to have the `langchain` python "
+    "package installed. Please install it with `pip install langchain`.",
+)
+
+if version.parse(langchain.__version__) < version.parse("0.0.188"):
+    raise ValueError(
+        "The Weights & Biases Langchain integration does not support versions 0.0.187 and lower. "
+        "To ensure proper functionality, please use version 0.0.188 or higher."
+    )
+
+# isort: off
+from langchain.callbacks.tracers import WandbTracer  # noqa: E402, I001
+
+
+class WandbTracer(WandbTracer):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        deprecate.deprecate(
+            field_name=deprecate.Deprecated.langchain_tracer,
+            warning_message="This feature is deprecated and has been moved to `langchain`. Enable tracing by setting "
+            "LANGCHAIN_WANDB_TRACING=true in your environment. See the documentation at "
+            "https://python.langchain.com/docs/ecosystem/integrations/agent_with_wandb_tracing for guidance. "
+            "Replace your current import with `from langchain.callbacks.tracers import WandbTracer`.",
+        )