wandb 0.18.0__py3-none-macosx_10_13_x86_64.whl

wandb/sdk/lib/tracelog.py

@@ -0,0 +1,255 @@
+"""tracelog.
+
+Functions:
+    log_message_queue   - message put() to queue
+    log_message_dequeue - message get() from queue
+    log_message_send    - message sent to socket
+    log_message_recv    - message received from socket
+    log_message_process - message processed by thread
+    log_message_link    - message linked to another message
+    log_message_assert  - message encountered problem
+
+"""
+
+import datetime
+import logging
+import secrets
+import sys
+import threading
+from typing import TYPE_CHECKING, Optional, cast
+
+if TYPE_CHECKING:
+    import multiprocessing
+    import queue
+    import socket
+    from typing import Union
+
+    from wandb.proto import wandb_internal_pb2 as pb
+    from wandb.proto import wandb_server_pb2 as spb
+
+    MessageQueueType = Union[pb.Record, pb.Result]
+    MessageType = Union[pb.Record, pb.Result, spb.ServerRequest, spb.ServerResponse]
+    QueueType = Union[multiprocessing.Queue, queue.Queue]
+    TransportType = Union[socket.socket, str]
+
+
+# Supported modes:
+#   logger - tracelog output goes to python logging (default)
+#   stdout - tracelog output goes to stdout
+#   stderr - tracelog output goes to stderr
+tracelog_mode: Optional[str] = "logger"
+
+logger = logging.getLogger(__name__)
+
+
+ANNOTATE_QUEUE_NAME = "_DEBUGLOG_QUEUE_NAME"
+
+# capture stdout and stderr before anyone messes with them
+stdout_write = sys.__stdout__.write  # type: ignore
+stderr_write = sys.__stderr__.write  # type: ignore
+
+
+def _log(
+    msg_type: str,
+    log_type: str,
+    is_response: bool = False,
+    record: Optional["pb.Record"] = None,
+    result: Optional["pb.Result"] = None,
+    resource: Optional[str] = None,
+) -> None:
+    prefix = "TRACELOG(1)"
+    tname = threading.current_thread().name
+    now = datetime.datetime.now()
+    ts = now.strftime("%H%M%S.%f")
+    arrow = "<-" if is_response else "->"
+    resource = resource or "unknown"
+    uuid = ""
+    data = record or result
+    record_id = ""
+    if data:
+        uuid = data.uuid or uuid
+        record_id = data._info._tracelog_id
+    uuid = uuid or "-"
+    record_id = record_id or "-"
+    relay = ""
+    if data and data.control and data.control.relay_id:
+        relay = data.control.relay_id
+    relay = relay or "-"
+    line = f"{prefix} {arrow} {ts} {record_id:16} {log_type:7} {resource:8} {tname:16} {msg_type:32} {uuid:32} {relay:32}"
+    if tracelog_mode == "stdout":
+        stdout_write(f"{line}\n")
+    elif tracelog_mode == "stderr":
+        stderr_write(f"{line}\n")
+    elif tracelog_mode == "logger":
+        logger.info(line)
+
+
+def _record_msg_type(record: "pb.Record") -> str:
+    msg_type = str(record.WhichOneof("record_type"))
+    if msg_type == "request":
+        request = record.request
+        msg_type = str(request.WhichOneof("request_type"))
+    return msg_type
+
+
+def _result_msg_type(result: "pb.Result") -> str:
+    msg_type = str(result.WhichOneof("result_type"))
+    if msg_type == "response":
+        response = result.response
+        msg_type = str(response.WhichOneof("response_type"))
+    return msg_type
+
+
+def _log_message(
+    msg: "MessageType", log_type: str, resource: Optional[str] = None
+) -> None:
+    record: Optional[pb.Record] = None
+    result: Optional[pb.Result] = None
+    is_response = False
+    msg_type: str
+    # Note: using strings to avoid an import
+    message_type_str = type(msg).__name__
+    if message_type_str == "Record":
+        record = cast("pb.Record", msg)
+        msg_type = _record_msg_type(record)
+    elif message_type_str == "Result":
+        is_response = True
+        result = cast("pb.Result", msg)
+        msg_type = _result_msg_type(result)
+    elif message_type_str == "ServerRequest":
+        server_request = cast("spb.ServerRequest", msg)
+        msg_type = str(server_request.WhichOneof("server_request_type"))
+        if msg_type == "record_publish":
+            record = server_request.record_publish
+            sub_msg_type = _record_msg_type(record)
+            msg_type = f"pub-{sub_msg_type}"
+        elif msg_type == "record_communicate":
+            record = server_request.record_communicate
+            sub_msg_type = _record_msg_type(record)
+            msg_type = f"comm-{sub_msg_type}"
+        # print("SRV", server_request)
+    elif message_type_str == "ServerResponse":
+        is_response = True
+        server_response = cast("spb.ServerResponse", msg)
+        msg_type = str(server_response.WhichOneof("server_response_type"))
+        if msg_type == "result_communicate":
+            result = server_response.result_communicate
+            sub_msg_type = _result_msg_type(result)
+            msg_type = f"comm-{sub_msg_type}"
+    else:
+        raise AssertionError(f"Unknown message type {message_type_str}")
+    _log(
+        msg_type,
+        is_response=is_response,
+        record=record,
+        result=result,
+        log_type=log_type,
+        resource=resource,
+    )
+
+
+def _log_message_queue(msg: "MessageQueueType", q: "QueueType") -> None:
+    _annotate_message(msg)
+    resource = getattr(q, ANNOTATE_QUEUE_NAME, None)
+    _log_message(msg, "queue", resource=resource)
+
+
+def _log_message_dequeue(msg: "MessageQueueType", q: "QueueType") -> None:
+    resource = getattr(q, ANNOTATE_QUEUE_NAME, None)
+    _log_message(msg, "dequeue", resource=resource)
+
+
+def _log_message_send(msg: "MessageType", t: "TransportType") -> None:
+    _log_message(msg, "send")
+
+
+def _log_message_recv(msg: "MessageType", t: "TransportType") -> None:
+    _log_message(msg, "recv")
+
+
+def _log_message_process(msg: "MessageType") -> None:
+    _log_message(msg, "process")
+
+
+def _log_message_link(src: "MessageType", dest: "MessageType") -> None:
+    _log_message(src, "source")
+    _log_message(dest, "dest")
+
+
+def _log_message_assert(msg: "MessageType") -> None:
+    _log_message(msg, "assert")
+
+
+def _annotate_queue(q: "QueueType", name: str) -> None:
+    setattr(q, ANNOTATE_QUEUE_NAME, name)
+
+
+def _annotate_message(msg: "MessageQueueType") -> None:
+    record_id = secrets.token_hex(8)
+    msg._info._tracelog_id = record_id
+
+
+#
+# Default functions when logging is disabled
+#
+
+
+def log_message_queue(msg: "MessageQueueType", q: "QueueType") -> None:
+    return None
+
+
+def log_message_dequeue(msg: "MessageQueueType", q: "QueueType") -> None:
+    return None
+
+
+def log_message_send(msg: "MessageType", t: "TransportType") -> None:
+    return None
+
+
+def log_message_recv(msg: "MessageType", t: "TransportType") -> None:
+    return None
+
+
+def log_message_process(msg: "MessageType") -> None:
+    return None
+
+
+def log_message_link(src: "MessageType", dest: "MessageType") -> None:
+    return None
+
+
+def log_message_assert(msg: "MessageType") -> None:
+    return None
+
+
+def annotate_queue(q: "QueueType", name: str) -> None:
+    return None
+
+
+def annotate_message(msg: "MessageQueueType") -> None:
+    return None
+
+
+def enable(log_mode: Optional[str] = None) -> None:
+    global tracelog_mode
+    if log_mode:
+        tracelog_mode = log_mode
+
+    global log_message_queue
+    global log_message_dequeue
+    global log_message_send
+    global log_message_recv
+    global log_message_process
+    global log_message_link
+    global log_message_assert
+    global annotate_queue
+    global annotate_message
+    log_message_queue = _log_message_queue
+    log_message_dequeue = _log_message_dequeue
+    log_message_send = _log_message_send
+    log_message_recv = _log_message_recv
+    log_message_process = _log_message_process
+    log_message_link = _log_message_link
+    log_message_assert = _log_message_assert
+    annotate_queue = _annotate_queue
+    annotate_message = _annotate_message

wandb/sdk/lib/viz.py

@@ -0,0 +1,123 @@
+from typing import Any, Dict, Optional, Tuple
+
+from wandb.data_types import Table
+from wandb.errors import Error
+
+
+class Visualize:
+    def __init__(self, id: str, data: Table) -> None:
+        self._id = id
+        self._data = data
+
+    def get_config_value(self, key: str) -> Dict[str, Any]:
+        return {
+            "id": self._id,
+            "historyFieldSettings": {"x-axis": "_step", "key": key},
+        }
+
+    @staticmethod
+    def get_config_key(key: str) -> Tuple[str, str, str]:
+        return "_wandb", "viz", key
+
+    @property
+    def value(self) -> Table:
+        return self._data
+
+
+class CustomChart:
+    def __init__(
+        self,
+        id: str,
+        data: Table,
+        fields: Dict[str, Any],
+        string_fields: Dict[str, Any],
+        split_table: Optional[bool] = False,
+    ) -> None:
+        self._id = id
+        self._data = data
+        self._fields = fields
+        self._string_fields = string_fields
+        self._split_table = split_table
+
+    def get_config_value(
+        self,
+        panel_type: str,
+        query: Dict[str, Any],
+    ) -> Dict[str, Any]:
+        return {
+            "panel_type": panel_type,
+            "panel_config": {
+                "panelDefId": self._id,
+                "fieldSettings": self._fields,
+                "stringSettings": self._string_fields,
+                "transform": {"name": "tableWithLeafColNames"},
+                "userQuery": query,
+            },
+        }
+
+    @staticmethod
+    def get_config_key(key: str) -> Tuple[str, str, str]:
+        return "_wandb", "visualize", key
+
+    @staticmethod
+    def user_query(table_key: str) -> Dict[str, Any]:
+        return {
+            "queryFields": [
+                {
+                    "name": "runSets",
+                    "args": [{"name": "runSets", "value": "${runSets}"}],
+                    "fields": [
+                        {"name": "id", "fields": []},
+                        {"name": "name", "fields": []},
+                        {"name": "_defaultColorIndex", "fields": []},
+                        {
+                            "name": "summaryTable",
+                            "args": [{"name": "tableKey", "value": table_key}],
+                            "fields": [],
+                        },
+                    ],
+                }
+            ],
+        }
+
+    @property
+    def table(self) -> Table:
+        return self._data
+
+    @property
+    def fields(self) -> Dict[str, Any]:
+        return self._fields
+
+    @property
+    def string_fields(self) -> Dict[str, Any]:
+        return self._string_fields
+
+
+def custom_chart(
+    vega_spec_name: str,
+    data_table: Table,
+    fields: Dict[str, Any],
+    string_fields: Optional[Dict[str, Any]] = None,
+    split_table: Optional[bool] = False,
+) -> CustomChart:
+    if string_fields is None:
+        string_fields = {}
+    if not isinstance(data_table, Table):
+        raise Error(
+            f"Expected `data_table` to be `wandb.Table` type, instead got {type(data_table).__name__}"
+        )
+    return CustomChart(
+        id=vega_spec_name,
+        data=data_table,
+        fields=fields,
+        string_fields=string_fields,
+        split_table=split_table,
+    )
+
+
+def visualize(id: str, value: Table) -> Visualize:
+    if not isinstance(value, Table):
+        raise Error(
+            f"Expected `value` to be `wandb.Table` type, instead got {type(value).__name__}"
+        )
+    return Visualize(id=id, data=value)

wandb/sdk/lib/wburls.py

@@ -0,0 +1,46 @@
+"""Container for urls used in the wandb package.
+
+Use this anytime a URL is displayed to the user.
+
+Usage:
+    ```python
+    from wandb.sdk.lib.wburls import wburls
+
+    print(f"This is a url {wburls.get('cli_launch')}")
+    ```
+"""
+
+from typing import TYPE_CHECKING, Dict, Optional
+
+if TYPE_CHECKING:
+    from ._wburls_generated import URLS
+
+
+class WBURLs:
+    _urls_dict: Optional[Dict["URLS", str]]
+
+    def __init__(self) -> None:
+        self._urls_dict = None
+
+    def _get_urls(self) -> Dict["URLS", str]:
+        return dict(
+            cli_launch="https://wandb.me/launch",
+            doc_run="https://wandb.me/run",
+            doc_require="https://wandb.me/library-require",
+            doc_start_err="https://docs.wandb.ai/guides/track/tracking-faq#initstarterror-error-communicating-with-wandb-process-",
+            doc_artifacts_guide="https://docs.wandb.ai/guides/artifacts",
+            upgrade_server="https://wandb.me/server-upgrade",
+            multiprocess="http://wandb.me/init-multiprocess",
+            wandb_init="https://wandb.me/wandb-init",
+            wandb_server="https://wandb.me/wandb-server",
+            wandb_define_metric="https://wandb.me/define-metric",
+            wandb_core="https://wandb.me/wandb-core",
+        )
+
+    def get(self, s: "URLS") -> str:
+        if self._urls_dict is None:
+            self._urls_dict = self._get_urls()
+        return self._urls_dict[s]
+
+
+wburls = WBURLs()

wandb/sdk/service/_startup_debug.py

@@ -0,0 +1,22 @@
+"""_startup_debug.
+
+Temporary helper to debug issues with wandb service startup
+"""
+
+import os
+import time
+
+
+def is_enabled() -> bool:
+    # This is very temporary to help diagnose problems seen by some
+    # customers which we are having trouble reproducing. It should be
+    # replaced by something more permanent in the future when we have
+    # proper logging for wandb-service
+    if os.environ.get("_WANDB_STARTUP_DEBUG"):
+        return True
+    return False
+
+
+def print_message(message: str) -> None:
+    time_now = time.time()
+    print("WANDB_STARTUP_DEBUG", time_now, message)

wandb/sdk/service/port_file.py

@@ -0,0 +1,53 @@
+"""port_file: write/read file containing port info."""
+
+import os
+import tempfile
+from typing import Optional
+
+
+class PortFile:
+    _sock_port: Optional[int]
+    _valid: bool
+
+    SOCK_TOKEN = "sock="
+    EOF_TOKEN = "EOF"
+
+    def __init__(self, sock_port: Optional[int] = None) -> None:
+        self._sock_port = sock_port
+        self._valid = False
+
+    def write(self, fname: str) -> None:
+        dname, bname = os.path.split(fname)
+        f = tempfile.NamedTemporaryFile(prefix=bname, dir=dname, mode="w", delete=False)
+        try:
+            tmp_filename = f.name
+            with f:
+                data = []
+                if self._sock_port:
+                    data.append(f"{self.SOCK_TOKEN}{self._sock_port}")
+                data.append(self.EOF_TOKEN)
+                port_str = "\n".join(data)
+                written = f.write(port_str)
+                assert written == len(port_str)
+            os.rename(tmp_filename, fname)
+        except Exception:
+            os.unlink(tmp_filename)
+            raise
+
+    def read(self, fname: str) -> None:
+        with open(fname) as f:
+            lines = f.readlines()
+            if lines[-1] != self.EOF_TOKEN:
+                return
+            for ln in lines:
+                if ln.startswith(self.SOCK_TOKEN):
+                    self._sock_port = int(ln[len(self.SOCK_TOKEN) :])
+            self._valid = True
+
+    @property
+    def sock_port(self) -> Optional[int]:
+        return self._sock_port
+
+    @property
+    def is_valid(self) -> bool:
+        return self._valid

wandb/sdk/service/server.py

@@ -0,0 +1,119 @@
+"""wandb server.
+
+Start up socket transport servers.
+"""
+
+import logging
+import os
+import sys
+from typing import Optional
+
+import wandb
+
+from ..lib import tracelog
+from . import _startup_debug, port_file
+from .server_sock import SocketServer
+from .streams import StreamMux
+
+
+class WandbServer:
+    _pid: Optional[int]
+    _sock_port: Optional[int]
+    _debug: bool
+    _serve_sock: bool
+    _sock_server: Optional[SocketServer]
+    _startup_debug_enabled: bool
+
+    def __init__(
+        self,
+        sock_port: Optional[int] = None,
+        port_fname: Optional[str] = None,
+        address: Optional[str] = None,
+        pid: Optional[int] = None,
+        debug: bool = True,
+        serve_sock: bool = False,
+    ) -> None:
+        self._sock_port = sock_port
+        self._port_fname = port_fname
+        self._address = address
+        self._pid = pid
+        self._debug = debug
+        self._serve_sock = serve_sock
+        self._sock_server = None
+        self._startup_debug_enabled = _startup_debug.is_enabled()
+
+        if debug:
+            logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)
+
+    def _inform_used_ports(self, sock_port: Optional[int]) -> None:
+        if not self._port_fname:
+            return
+        pf = port_file.PortFile(sock_port=sock_port)
+        pf.write(self._port_fname)
+
+    def _start_sock(self, mux: StreamMux) -> int:
+        address: str = self._address or "127.0.0.1"
+        port: int = self._sock_port or 0
+        self._sock_server = SocketServer(mux=mux, address=address, port=port)
+        try:
+            self._sock_server.start()
+            port = self._sock_server.port
+            if self._pid:
+                mux.set_pid(self._pid)
+        except KeyboardInterrupt:
+            mux.cleanup()
+            raise
+        except Exception:
+            mux.cleanup()
+            raise
+        return port
+
+    def _stop_servers(self) -> None:
+        if self._sock_server:
+            self._sock_server.stop()
+
+    def _setup_tracelog(self) -> None:
+        # TODO: remove this temporary hack, need to find a better way to pass settings
+        # to the server.  for now lets just look at the environment variable we need
+        tracelog_mode = os.environ.get("WANDB_TRACELOG")
+        if tracelog_mode:
+            tracelog.enable(tracelog_mode)
+
+    def _startup_debug_print(self, message: str) -> None:
+        if not self._startup_debug_enabled:
+            return
+        _startup_debug.print_message(message)
+
+    def _setup_proctitle(self, sock_port: Optional[int]) -> None:
+        # TODO: similar to _setup_tracelog, the internal_process should have
+        # a better way to have access to settings.
+        disable_setproctitle = os.environ.get("WANDB__DISABLE_SETPROCTITLE")
+        if disable_setproctitle:
+            return
+
+        setproctitle = wandb.util.get_optional_module("setproctitle")
+        if setproctitle:
+            service_ver = 2
+            pid = str(self._pid or 0)
+            transport = "s" if sock_port else "g"
+            port = sock_port or 0
+            # this format is similar to wandb_manager token, but it's purely informative now
+            # (consider unifying this in the future)
+            service_id = f"{service_ver}-{pid}-{transport}-{port}"
+            proc_title = f"wandb-service({service_id})"
+            self._startup_debug_print("before_setproctitle")
+            setproctitle.setproctitle(proc_title)
+            self._startup_debug_print("after_setproctitle")
+
+    def serve(self) -> None:
+        self._setup_tracelog()
+        mux = StreamMux()
+        self._startup_debug_print("before_network")
+        sock_port = self._start_sock(mux=mux) if self._serve_sock else None
+        self._startup_debug_print("after_network")
+        self._inform_used_ports(sock_port=sock_port)
+        self._startup_debug_print("after_inform")
+        self._setup_proctitle(sock_port=sock_port)
+        self._startup_debug_print("before_loop")
+        mux.loop()
+        self._stop_servers()