wandb 0.18.0rc1__py3-none-macosx_10_13_x86_64.whl

wandb/jupyter.py

@@ -0,0 +1,515 @@
+import json
+import logging
+import os
+import re
+import shutil
+import sys
+from base64 import b64encode
+from typing import Dict
+
+import requests
+from requests.compat import urljoin
+
+import wandb
+import wandb.util
+from wandb.sdk.lib import filesystem
+
+try:
+    from IPython.core.getipython import get_ipython
+    from IPython.core.magic import Magics, line_cell_magic, magics_class
+    from IPython.core.magic_arguments import argument, magic_arguments, parse_argstring
+    from IPython.display import display
+except ImportError:
+    wandb.termwarn("ipython is not supported in python 2.7, upgrade to 3.x")
+
+    class Magics:
+        pass
+
+    def magics_class(*args, **kwargs):
+        return lambda *args, **kwargs: None
+
+    def magic_arguments(*args, **kwargs):
+        return lambda *args, **kwargs: None
+
+    def argument(*args, **kwargs):
+        return lambda *args, **kwargs: None
+
+    def line_cell_magic(*args, **kwargs):
+        return lambda *args, **kwargs: None
+
+
+logger = logging.getLogger(__name__)
+
+__IFrame = None
+
+
+def maybe_display():
+    """Display a run if the user added cell magic and we have run."""
+    if __IFrame is not None:
+        return __IFrame.maybe_display()
+    return False
+
+
+def quiet():
+    if __IFrame is not None:
+        return __IFrame.opts.get("quiet")
+    return False
+
+
+class IFrame:
+    def __init__(self, path=None, opts=None):
+        self.path = path
+        self.api = wandb.Api()
+        self.opts = opts or {}
+        self.displayed = False
+        self.height = self.opts.get("height", 420)
+
+    def maybe_display(self) -> bool:
+        if not self.displayed and (self.path or wandb.run):
+            display(self)
+        return self.displayed
+
+    def _repr_html_(self):
+        try:
+            self.displayed = True
+            if self.opts.get("workspace", False):
+                if self.path is None and wandb.run:
+                    self.path = wandb.run.path
+            if isinstance(self.path, str):
+                object = self.api.from_path(self.path)
+            else:
+                object = wandb.run
+            if object is None:
+                if wandb.Api().api_key is None:
+                    return "You must be logged in to render wandb in jupyter, run `wandb.login()`"
+                else:
+                    object = self.api.project(
+                        "/".join(
+                            [
+                                wandb.Api().default_entity,
+                                wandb.util.auto_project_name(None),
+                            ]
+                        )
+                    )
+            return object.to_html(self.height, hidden=False)
+        except wandb.Error as e:
+            return f"Can't display wandb interface<br/>{e}"
+
+
+@magics_class
+class WandBMagics(Magics):
+    def __init__(self, shell, require_interaction=False):
+        super().__init__(shell)
+        self.options = {}
+
+    @magic_arguments()
+    @argument(
+        "path",
+        default=None,
+        nargs="?",
+        help="A path to a resource you want to display, defaults to wandb.run.path",
+    )
+    @argument(
+        "-w",
+        "--workspace",
+        default=False,
+        action="store_true",
+        help="Display the entire run project workspace",
+    )
+    @argument(
+        "-q",
+        "--quiet",
+        default=False,
+        action="store_true",
+        help="Display the minimal amount of output",
+    )
+    @argument(
+        "-h",
+        "--height",
+        default=420,
+        type=int,
+        help="The height of the iframe in pixels",
+    )
+    @line_cell_magic
+    def wandb(self, line, cell=None):
+        """Display wandb resources in jupyter.  This can be used as cell or line magic.
+
+        %wandb USERNAME/PROJECT/runs/RUN_ID
+        ---
+        %%wandb -h 1024
+        with wandb.init() as run:
+            run.log({"loss": 1})
+        """
+        # Record options
+        args = parse_argstring(self.wandb, line)
+        self.options["height"] = args.height
+        self.options["workspace"] = args.workspace
+        self.options["quiet"] = args.quiet
+        iframe = IFrame(args.path, opts=self.options)
+        displayed = iframe.maybe_display()
+        if cell is not None:
+            if not displayed:
+                # Store the IFrame globally and attempt to display if we have a run
+                cell = (
+                    f"wandb.jupyter.__IFrame = wandb.jupyter.IFrame(opts={self.options})\n"
+                    + cell
+                    + "\nwandb.jupyter.__IFrame = None"
+                )
+            get_ipython().run_cell(cell)
+
+
+def notebook_metadata_from_jupyter_servers_and_kernel_id():
+    servers, kernel_id = jupyter_servers_and_kernel_id()
+    for s in servers:
+        if s.get("password"):
+            raise ValueError("Can't query password protected kernel")
+        res = requests.get(
+            urljoin(s["url"], "api/sessions"), params={"token": s.get("token", "")}
+        ).json()
+        for nn in res:
+            if isinstance(nn, dict) and nn.get("kernel") and "notebook" in nn:
+                if nn["kernel"]["id"] == kernel_id:
+                    return {
+                        "root": s.get("root_dir", s.get("notebook_dir", os.getcwd())),
+                        "path": nn["notebook"]["path"],
+                        "name": nn["notebook"]["name"],
+                    }
+
+    if not kernel_id:
+        return None
+
+    # Built-in notebook server in VS Code
+    try:
+        from IPython import get_ipython
+
+        ipython = get_ipython()
+        notebook_path = ipython.kernel.shell.user_ns.get("__vsc_ipynb_file__")
+        if notebook_path:
+            return {
+                "root": os.path.dirname(notebook_path),
+                "path": notebook_path,
+                "name": os.path.basename(notebook_path),
+            }
+    except Exception:
+        return None
+
+
+def notebook_metadata(silent: bool) -> Dict[str, str]:
+    """Attempt to query jupyter for the path and name of the notebook file.
+
+    This can handle different jupyter environments, specifically:
+
+    1. Colab
+    2. Kaggle
+    3. JupyterLab
+    4. Notebooks
+    5. Other?
+    """
+    error_message = (
+        "Failed to detect the name of this notebook, you can set it manually with "
+        "the WANDB_NOTEBOOK_NAME environment variable to enable code saving."
+    )
+    try:
+        jupyter_metadata = notebook_metadata_from_jupyter_servers_and_kernel_id()
+
+        # Colab:
+        # request the most recent contents
+        ipynb = attempt_colab_load_ipynb()
+        if ipynb is not None and jupyter_metadata is not None:
+            return {
+                "root": "/content",
+                "path": jupyter_metadata["path"],
+                "name": jupyter_metadata["name"],
+            }
+
+        # Kaggle:
+        if wandb.util._is_kaggle():
+            # request the most recent contents
+            ipynb = attempt_kaggle_load_ipynb()
+            if ipynb:
+                return {
+                    "root": "/kaggle/working",
+                    "path": ipynb["metadata"]["name"],
+                    "name": ipynb["metadata"]["name"],
+                }
+
+        if jupyter_metadata:
+            return jupyter_metadata
+        if not silent:
+            logger.error(error_message)
+        return {}
+    except Exception:
+        # TODO: report this exception
+        # TODO: Fix issue this is not the logger initialized in in wandb.init()
+        # since logger is not attached, outputs to notebook
+        if not silent:
+            logger.error(error_message)
+        return {}
+
+
+def jupyter_servers_and_kernel_id():
+    """Return a list of servers and the current kernel_id.
+
+    Used to query for the name of the notebook.
+    """
+    try:
+        import ipykernel
+
+        kernel_id = re.search(
+            "kernel-(.*).json", ipykernel.connect.get_connection_file()
+        ).group(1)
+        # We're either in jupyterlab or a notebook, lets prefer the newer jupyter_server package
+        serverapp = wandb.util.get_module("jupyter_server.serverapp")
+        notebookapp = wandb.util.get_module("notebook.notebookapp")
+        servers = []
+        if serverapp is not None:
+            servers.extend(list(serverapp.list_running_servers()))
+        if notebookapp is not None:
+            servers.extend(list(notebookapp.list_running_servers()))
+        return servers, kernel_id
+    except (AttributeError, ValueError, ImportError):
+        return [], None
+
+
+def attempt_colab_load_ipynb():
+    colab = wandb.util.get_module("google.colab")
+    if colab:
+        # This isn't thread safe, never call in a thread
+        response = colab._message.blocking_request("get_ipynb", timeout_sec=5)
+        if response:
+            return response["ipynb"]
+
+
+def attempt_kaggle_load_ipynb():
+    kaggle = wandb.util.get_module("kaggle_session")
+    if kaggle:
+        try:
+            client = kaggle.UserSessionClient()
+            parsed = json.loads(client.get_exportable_ipynb()["source"])
+            # TODO: couldn't find a way to get the name of the notebook...
+            parsed["metadata"]["name"] = "kaggle.ipynb"
+            return parsed
+        except Exception:
+            logger.exception("Unable to load kaggle notebook")
+            return None
+
+
+def attempt_colab_login(app_url):
+    """This renders an iframe to wandb in the hopes it posts back an api key."""
+    from google.colab import output
+    from google.colab._message import MessageError
+    from IPython import display
+
+    display.display(
+        display.Javascript(
+            """
+        window._wandbApiKey = new Promise((resolve, reject) => {{
+            function loadScript(url) {{
+            return new Promise(function(resolve, reject) {{
+                let newScript = document.createElement("script");
+                newScript.onerror = reject;
+                newScript.onload = resolve;
+                document.body.appendChild(newScript);
+                newScript.src = url;
+            }});
+            }}
+            loadScript("https://cdn.jsdelivr.net/npm/postmate/build/postmate.min.js").then(() => {{
+            const iframe = document.createElement('iframe')
+            iframe.style.cssText = "width:0;height:0;border:none"
+            document.body.appendChild(iframe)
+            const handshake = new Postmate({{
+                container: iframe,
+                url: '{}/authorize'
+            }});
+            const timeout = setTimeout(() => reject("Couldn't auto authenticate"), 5000)
+            handshake.then(function(child) {{
+                child.on('authorize', data => {{
+                    clearTimeout(timeout)
+                    resolve(data)
+                }});
+            }});
+            }})
+        }});
+    """.format(app_url.replace("http:", "https:"))
+        )
+    )
+    try:
+        return output.eval_js("_wandbApiKey")
+    except MessageError:
+        return None
+
+
+class Notebook:
+    def __init__(self, settings):
+        self.outputs = {}
+        self.settings = settings
+        self.shell = get_ipython()
+
+    def save_display(self, exc_count, data_with_metadata):
+        self.outputs[exc_count] = self.outputs.get(exc_count, [])
+
+        # byte values such as images need to be encoded in base64
+        # otherwise nbformat.v4.new_output will throw a NotebookValidationError
+        data = data_with_metadata["data"]
+        b64_data = {}
+        for key in data:
+            val = data[key]
+            if isinstance(val, bytes):
+                b64_data[key] = b64encode(val).decode("utf-8")
+            else:
+                b64_data[key] = val
+
+        self.outputs[exc_count].append(
+            {"data": b64_data, "metadata": data_with_metadata["metadata"]}
+        )
+
+    def probe_ipynb(self):
+        """Return notebook as dict or None."""
+        relpath = self.settings._jupyter_path
+        if relpath:
+            if os.path.exists(relpath):
+                with open(relpath) as json_file:
+                    data = json.load(json_file)
+                    return data
+
+        colab_ipynb = attempt_colab_load_ipynb()
+        if colab_ipynb:
+            return colab_ipynb
+
+        kaggle_ipynb = attempt_kaggle_load_ipynb()
+        if kaggle_ipynb and len(kaggle_ipynb["cells"]) > 0:
+            return kaggle_ipynb
+
+        return
+
+    def save_ipynb(self) -> bool:
+        if not self.settings.save_code:
+            logger.info("not saving jupyter notebook")
+            return False
+        ret = False
+        try:
+            ret = self._save_ipynb()
+        except Exception as e:
+            logger.info(f"Problem saving notebook: {repr(e)}")
+        return ret
+
+    def _save_ipynb(self) -> bool:
+        relpath = self.settings._jupyter_path
+        logger.info("looking for notebook: %s", relpath)
+        if relpath:
+            if os.path.exists(relpath):
+                shutil.copy(
+                    relpath,
+                    os.path.join(
+                        self.settings._tmp_code_dir, os.path.basename(relpath)
+                    ),
+                )
+                return True
+
+        # TODO: likely only save if the code has changed
+        colab_ipynb = attempt_colab_load_ipynb()
+        if colab_ipynb:
+            try:
+                jupyter_metadata = (
+                    notebook_metadata_from_jupyter_servers_and_kernel_id()
+                )
+                nb_name = jupyter_metadata["name"]
+            except Exception:
+                nb_name = "colab.ipynb"
+            if not nb_name.endswith(".ipynb"):
+                nb_name += ".ipynb"
+            with open(
+                os.path.join(
+                    self.settings._tmp_code_dir,
+                    nb_name,
+                ),
+                "w",
+                encoding="utf-8",
+            ) as f:
+                f.write(json.dumps(colab_ipynb))
+            return True
+
+        kaggle_ipynb = attempt_kaggle_load_ipynb()
+        if kaggle_ipynb and len(kaggle_ipynb["cells"]) > 0:
+            with open(
+                os.path.join(
+                    self.settings._tmp_code_dir, kaggle_ipynb["metadata"]["name"]
+                ),
+                "w",
+                encoding="utf-8",
+            ) as f:
+                f.write(json.dumps(kaggle_ipynb))
+            return True
+
+        return False
+
+    def save_history(self):
+        """This saves all cell executions in the current session as a new notebook."""
+        try:
+            from nbformat import v4, validator, write
+        except ImportError:
+            logger.error("Run pip install nbformat to save notebook history")
+            return
+        # TODO: some tests didn't patch ipython properly?
+        if self.shell is None:
+            return
+        cells = []
+        hist = list(self.shell.history_manager.get_range(output=True))
+        if len(hist) <= 1 or not self.settings.save_code:
+            logger.info("not saving jupyter history")
+            return
+        try:
+            for _, execution_count, exc in hist:
+                if exc[1]:
+                    # TODO: capture stderr?
+                    outputs = [
+                        v4.new_output(output_type="stream", name="stdout", text=exc[1])
+                    ]
+                else:
+                    outputs = []
+                if self.outputs.get(execution_count):
+                    for out in self.outputs[execution_count]:
+                        outputs.append(
+                            v4.new_output(
+                                output_type="display_data",
+                                data=out["data"],
+                                metadata=out["metadata"] or {},
+                            )
+                        )
+                cells.append(
+                    v4.new_code_cell(
+                        execution_count=execution_count, source=exc[0], outputs=outputs
+                    )
+                )
+            if hasattr(self.shell, "kernel"):
+                language_info = self.shell.kernel.language_info
+            else:
+                language_info = {"name": "python", "version": sys.version}
+            logger.info("saving %i cells to _session_history.ipynb", len(cells))
+            nb = v4.new_notebook(
+                cells=cells,
+                metadata={
+                    "kernelspec": {
+                        "display_name": "Python %i" % sys.version_info[0],
+                        "name": "python%i" % sys.version_info[0],
+                        "language": "python",
+                    },
+                    "language_info": language_info,
+                },
+            )
+            state_path = os.path.join("code", "_session_history.ipynb")
+            wandb.run._set_config_wandb("session_history", state_path)
+            filesystem.mkdir_exists_ok(os.path.join(wandb.run.dir, "code"))
+            with open(
+                os.path.join(self.settings._tmp_code_dir, "_session_history.ipynb"),
+                "w",
+                encoding="utf-8",
+            ) as f:
+                write(nb, f, version=4)
+            with open(
+                os.path.join(wandb.run.dir, state_path), "w", encoding="utf-8"
+            ) as f:
+                write(nb, f, version=4)
+        except (OSError, validator.NotebookValidationError) as e:
+            logger.error("Unable to save ipython session history:\n%s", e)
+            pass

wandb/magic.py

@@ -0,0 +1,3 @@
+from wandb.integration import magic
+
+magic.magic_install()

wandb/mpmain/__main__.py

@@ -0,0 +1 @@
+# This module is initialized after multiprocessing spawn

wandb/old/core.py

@@ -0,0 +1,131 @@
+"""Core variables, functions, and classes that we want in the wandb
+module but are also used in modules that import the wandb module.
+
+The purpose of this module is to break circular imports.
+"""
+
+import os
+import sys
+import tempfile
+import time
+
+import click
+
+import wandb
+from wandb import env
+
+# We use the hidden version if it already exists, otherwise non-hidden.
+if os.path.exists(os.path.join(env.get_dir(os.getcwd()), ".wandb")):
+    __stage_dir__ = ".wandb" + os.sep
+elif os.path.exists(os.path.join(env.get_dir(os.getcwd()), "wandb")):
+    __stage_dir__ = "wandb" + os.sep
+else:
+    __stage_dir__ = None
+
+SCRIPT_PATH = os.path.abspath(sys.argv[0])
+wandb.START_TIME = time.time()
+LIB_ROOT = os.path.join(os.path.dirname(__file__), "..")
+IS_GIT = os.path.exists(os.path.join(LIB_ROOT, ".git"))
+
+
+def wandb_dir(root_dir=None):
+    if root_dir is None or root_dir == "":
+        try:
+            cwd = os.getcwd()
+        except OSError:
+            termwarn("os.getcwd() no longer exists, using system temp directory")
+            cwd = tempfile.gettempdir()
+        root_dir = env.get_dir(cwd)
+    path = os.path.join(root_dir, __stage_dir__ or ("wandb" + os.sep))
+    if not os.access(root_dir, os.W_OK):
+        termwarn(
+            f"Path {path} wasn't writable, using system temp directory", repeat=False
+        )
+        path = os.path.join(tempfile.gettempdir(), __stage_dir__ or ("wandb" + os.sep))
+    return path
+
+
+def _set_stage_dir(stage_dir):
+    # Used when initing a new project with "wandb init"
+    global __stage_dir__
+    __stage_dir__ = stage_dir
+
+
+class Error(Exception):
+    """Base W&B Error"""
+
+    def __init__(self, message):
+        super().__init__(message)
+        self.message = message
+
+    # For python 2 support
+    def encode(self, encoding):
+        return self.message
+
+
+class WandbWarning(Warning):
+    """Base W&B Warning"""
+
+    pass
+
+
+LOG_STRING = click.style("wandb", fg="blue", bold=True)
+ERROR_STRING = click.style("ERROR", bg="red", fg="green")
+WARN_STRING = click.style("WARNING", fg="yellow")
+PRINTED_MESSAGES = set()
+
+
+# TODO(adrian): if output has been redirected, make this write to the original STDERR
+# so it doesn't get logged to the backend
+def termlog(string="", newline=True, repeat=True):
+    """Log to standard error with formatting.
+
+    Arguments:
+        string (str, optional): The string to print
+        newline (bool, optional): Print a newline at the end of the string
+        repeat (bool, optional): If set to False only prints the string once per process
+    """
+    if string:
+        line = "\n".join([f"{LOG_STRING}: {s}" for s in string.split("\n")])
+    else:
+        line = ""
+    if not repeat and line in PRINTED_MESSAGES:
+        return
+    # Repeated line tracking limited to 1k messages
+    if len(PRINTED_MESSAGES) < 1000:
+        PRINTED_MESSAGES.add(line)
+    if os.getenv(env.SILENT):
+        from wandb import util
+        from wandb.sdk.lib import filesystem
+
+        filesystem.mkdir_exists_ok(os.path.dirname(util.get_log_file_path()))
+        with open(util.get_log_file_path(), "w") as log:
+            click.echo(line, file=log, nl=newline)
+    else:
+        click.echo(line, file=sys.stderr, nl=newline)
+
+
+def termwarn(string, **kwargs):
+    string = "\n".join([f"{WARN_STRING} {s}" for s in string.split("\n")])
+    termlog(string=string, newline=True, **kwargs)
+
+
+def termerror(string, **kwargs):
+    string = "\n".join([f"{ERROR_STRING} {s}" for s in string.split("\n")])
+    termlog(string=string, newline=True, **kwargs)
+
+
+__all__ = [
+    "__stage_dir__",
+    "SCRIPT_PATH",
+    "START_TIME",
+    "wandb_dir",
+    "_set_stage_dir",
+    "Error",
+    "WandbWarning",
+    "LOG_STRING",
+    "ERROR_STRING",
+    "termlog",
+    "termwarn",
+    "termerror",
+]